from glob import glob
import sys
import array
import OpenEXR
import Imath
import scipy.io
import scipy.misc
import numpy as np

# This could eventualy be useful
sA = scipy.io.loadmat('solidAngles')
d = np.dstack((sA['solidAngles'], sA['solidAngles'], sA['solidAngles']))

for f in glob('*-img.exr'):
    print(f)
    f_mask = f.rsplit("-", 1)[0] + "-mask.exr"
    f_mask_hdl = OpenEXR.InputFile(f_mask)
    FLOAT = Imath.PixelType(Imath.PixelType.FLOAT)
    Y = np.array(array.array('f', f_mask_hdl.channel("Y", FLOAT)).tolist())

    f_hdl = OpenEXR.InputFile(f)

    # Read the three color channels as 32-bit floats
    FLOAT = Imath.PixelType(Imath.PixelType.FLOAT)
    (R,G,B) = [array.array('f', f_hdl.channel(Chan, FLOAT)).tolist() for Chan in ("R", "G", "B")]
    R = np.array(R)
    G = np.array(G)
    B = np.array(B)
    #import pdb; pdb.set_trace()

    # Apply mask (Is that what we want?)
    #R, G, B = [X[Y>0.5] for X in (R, G, B)]

    # Set masked zones to black
    R[Y<0.5] = 0

    # Display example
    R = R.reshape(256,256)
    scipy.misc.imsave('r.png', R)