GenStereo / genstereo /dataset /EXRloader.py
FQiao's picture
Upload 144 files
b4754db verified
import numpy
import OpenEXR
import Imath
import imageio
import glob
import os
def exr2hdr(exrpath):
File = OpenEXR.InputFile(exrpath)
PixType = Imath.PixelType(Imath.PixelType.FLOAT)
DW = File.header()['dataWindow']
CNum = len(File.header()['channels'].keys())
if (CNum > 1):
Channels = ['R', 'G', 'B']
CNum = 3
else:
Channels = ['G']
Size = (DW.max.x - DW.min.x + 1, DW.max.y - DW.min.y + 1)
Pixels = [numpy.fromstring(File.channel(c, PixType), dtype=numpy.float32) for c in Channels]
hdr = numpy.zeros((Size[1],Size[0],CNum),dtype=numpy.float32)
if (CNum == 1):
hdr[:,:,0] = numpy.reshape(Pixels[0],(Size[1],Size[0]))
else:
hdr[:,:,0] = numpy.reshape(Pixels[0],(Size[1],Size[0]))
hdr[:,:,1] = numpy.reshape(Pixels[1],(Size[1],Size[0]))
hdr[:,:,2] = numpy.reshape(Pixels[2],(Size[1],Size[0]))
return hdr
def writehdr(hdrpath,hdr):
h, w, c = hdr.shape
if c == 1:
hdr = numpy.pad(hdr, ((0, 0), (0, 0), (0, 2)), 'constant')
hdr[:,:,1] = hdr[:,:,0]
hdr[:,:,2] = hdr[:,:,0]
imageio.imwrite(hdrpath,hdr,format='hdr')
def load_exr(filename):
hdr = exr2hdr(filename)
h, w, c = hdr.shape
if c == 1:
hdr = numpy.squeeze(hdr)
return hdr
def test_exr():
files = glob.glob('D:/MLProjects/data/home/*.exr')
savepath = 'D:/MLProjects/data/home'
total = len(files)
count = 0
print ('Files Num:', total)
for file in files:
hdr = exr2hdr(file)
filename,file_ext = os.path.splitext(file)
filename = os.path.basename(filename)
filename = filename + '.hdr'
curpath = os.path.join(savepath,filename)
writehdr(curpath,hdr)
count = count + 1
print ('process:', count, '/', total)
if __name__ == '__main__':
test_exr()