EXR DWAA compression

ear_dwaa_compression.py

import os
import sys
import subprocess


def get_size(start_path):
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(start_path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            # skip if it is symbolic link
            if not os.path.islink(fp):
                total_size += os.path.getsize(fp)

    if os.path.isfile(start_path) and not os.path.islink(start_path):
        total_size += os.path.getsize(start_path)

    return total_size


def sizeof_fmt(num, suffix='B'):
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


def compress(path):
    size_saved = 0
    for root, dirs, files in os.walk(path):
        for file in files:
            file = os.path.join(root, file)

            if not file.endswith(".exr"):
                continue

            # Ignore z depth renders.
            if "_Z" in file:
                continue

            # Ignore existing dwaa compressed exrs.
            p = subprocess.Popen(
                "magick identify -format \"%[compression]\""
                " \"{}\"".format(file),
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                stdin=subprocess.PIPE,
                cwd=os.path.dirname(file)
            )

            output = p.communicate()[0]

            if p.returncode != 0:
                raise ValueError(output)

            if "DWAA" in str(output):
                continue

            # Process exr.
            print("Processing {}".format(file))
            temp_file = file.replace(".exr", "_compressed.exr")
            p = subprocess.Popen(
                (
                    "magick convert -background \"rgba(0, 0, 0, 0)\" \"{}\""
                    " -compress DWAA \"{}\"".format(file, temp_file)
                ),
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                stdin=subprocess.PIPE,
                cwd=os.path.dirname(file)
            )

            output = p.communicate()[0]

            if p.returncode != 0:
                raise ValueError(output)

            original_size = get_size(file)
            new_size = get_size(temp_file)
            size_saved += original_size - new_size
            print("Saved: {}".format(sizeof_fmt(size_saved)))

            os.remove(file)
            os.rename(temp_file, file)


if __name__ == "__main__":
    compress(sys.argv[1])