Convert EXR files to numpy float32 matrices

exr_to_np.py

# -*- coding:utf-8 -*-

"""
Place this script on a directory and convert the .exr files it contains
into np.float32 matrices
"""


import os
import OpenEXR
import numpy as np


def depth_exr_to_numpy(exr_path,
                       typemap={"HALF": np.float16, "FLOAT": np.float32}):
    """
    """
    print("[processing {}]".format(exr_path))
    # load EXR and extract shape
    exr = OpenEXR.InputFile(exr_path)
    print(exr.header())
    dw = exr.header()["dataWindow"]
    shape = (dw.max.y - dw.min.y + 1, dw.max.x - dw.min.x + 1)
    #
    arr_maps = {}
    for ch_name, ch in exr.header()["channels"].items():
        print("reading channel", ch_name)
        # This, and __str__ seem to be the only ways to get typename
        exr_typename = ch.type.names[ch.type.v]
        np_type = typemap[exr_typename]
        # convert channel to np array
        bytestring = exr.channel(ch_name, ch.type)
        arr = np.frombuffer(bytestring, dtype=np_type).reshape(shape)
        arr_maps[ch_name] = arr
    #
    rgb_equal = ((arr_maps["R"] == arr_maps["G"]).all() and
                 (arr_maps["R"] == arr_maps["B"]).all())
    assert rgb_equal, "this function assumes that R, G, B must be identical!"
    #
    return arr_maps["R"]


def plot_arr(arr):
    """
    arr has to be np.float32
    """
    print("plotting", arr.shape, arr.dtype, arr.min(), arr.max())
    import matplotlib.pyplot as plt
    plt.imshow(arr)
    plt.show()


def convert_depth_exr_files_to_np_float32(*exr_paths):
    """
    """
    return [depth_exr_to_numpy(p).astype(np.float32) for p in exr_paths]


if __name__ == "__main__":
    dir_path = "."
    exr_paths = [f for f in os.listdir(dir_path) if f.endswith(".exr")]
    np_arrs = convert_depth_exr_files_to_np_float32(*exr_paths)
    # for arr in np_arrs:
    #     plot_arr(arr)