exr_split_channels.py

import OpenEXR, Imath
import numpy as np
import matplotlib.pyplot as plt

def exr_split_channels(exr_fname):
  """
  Load a Mitsuba render generated multi-channel OpenEXR file.
  Split the channels according to the information stored. 
  """
    ## Load all images
    fin = OpenEXR.InputFile(exr_fname)
    pt = Imath.PixelType(Imath.PixelType.FLOAT)

    dw = fin.header()['dataWindow']
    img_size = (dw.max.y - dw.min.y + 1, dw.max.x - dw.min.x + 1)

    # Get the rendered image
    def clamp_pixels(img):
        img[img > 1.0] = 1.0
        return img

    img_rgb = [ clamp_pixels(
        np.fromstring(
            fin.channel('color.' + c, pt),
            dtype = np.float32)).reshape(img_size) for c in "RGB" ]
    img_merged = reduce( lambda a,b: np.dstack((a,b)), img_rgb )

    # Get the normals
    normals_rgb = [ np.fromstring(
        fin.channel('normal.' + c, pt), dtype = np.float32).reshape(img_size) for c in "RGB" ]
    normals = reduce( lambda a,b: np.dstack((a,b)), normals_rgb )
    
        # Get the distances
    distances = np.fromstring(fin.channel('distance.Y', pt), dtype = np.float32).reshape(img_size)

    # Gives the pixel to shape mapping
    shape_masks = np.rint(np.fromstring(fin.channel('index.Y', pt), dtype = np.float32))
    shape_masks.shape = img_size
    shape_masks = shape_masks.astype(np.uint32)

    ## Plot the images
    plot_list = np.asarray([
        [img_merged, 'Rendered Image'],
        [normals, 'Normal'],
        [distances, 'Distance'],
        [shape_masks, 'Shape Index']])
    plot_list.shape = (2,2,2)

    fig, ax_list = plt.subplots(2, 2, sharex = 'all', sharey = 'all')
    for i in [0,1]:
        for j in [0,1]:
            img, title = plot_list[i,j]
            ax_list[i,j].imshow(img)
            ax_list[i,j].axis('off')
            ax_list[i,j].set_title(title)


    fig.subplots_adjust(wspace = 0.01, hspace = 0.2,
                        top = 0.9, bottom = 0.1, left=0, right=1)
    fig.savefig(exr_fname.split('.')[0] + '.pdf')
    #plt.show()
    return [img_merged, normals, distances, shape_masks]
    
if __name__ == "__main__":
    import sys, os
    exr_fname = 'multi_channel.exr'
    if len(sys.argv) > 1:
        exr_fname = sys.argv[1]
    if os.path.exists(exr_fname):
        exr_split_channels(exr_fname)
    else:
        print "Error: file " + exr_fname + " does not exist"