mirror of
https://github.com/OpenFOAM/ThirdParty-6.git
synced 2025-12-08 06:57:43 +00:00
137 lines
3.9 KiB
Python
137 lines
3.9 KiB
Python
import OpenEXR as oe
|
|
import Imath as im
|
|
import numpy as np
|
|
|
|
|
|
class OexrCompression:
|
|
NONE = 0
|
|
RLE = 1
|
|
ZIPS = 2
|
|
ZIP = 3
|
|
PIZ = 4
|
|
PXR24 = 5
|
|
|
|
def save_rgb(image, filePath, comp = OexrCompression.ZIP):
|
|
'''Saves the rgb (image) in OpenEXR format. Expects a 3-chan uint32 image.'''
|
|
|
|
if len(image.shape) != 3:
|
|
raise Exception("Incorrect dimensions!")
|
|
|
|
h, w, c = image.shape # expects Numpy convention (row, col) -> (height, width)
|
|
|
|
if c != 3:
|
|
raise Exception("Incorrect number of channels!")
|
|
|
|
if image.dtype != "uint32":
|
|
raise Exception("Incorrect type!, expected uint32")
|
|
|
|
try:
|
|
|
|
header = oe.Header(w, h)
|
|
header["channels"] = {"R" : im.Channel(im.PixelType(oe.UINT)),
|
|
"G" : im.Channel(im.PixelType(oe.UINT)),
|
|
"B" : im.Channel(im.PixelType(oe.UINT))}
|
|
|
|
header['compression'] = im.Compression(comp)
|
|
|
|
of = oe.OutputFile(filePath, header)
|
|
r_data = image[:, :, 0].tostring()
|
|
g_data = image[:, :, 1].tostring()
|
|
b_data = image[:, :, 2].tostring()
|
|
of.writePixels({"R" : r_data, "G" : g_data, "B" : b_data})
|
|
of.close()
|
|
|
|
except:
|
|
raise
|
|
|
|
|
|
def save_depth(image, filePath, comp = OexrCompression.ZIP):
|
|
'''Saves the zBuffer (image) in OpenEXR format. Expects a 1-chann float32 image.'''
|
|
|
|
if len(image.shape) != 2:
|
|
raise Exception("Incorrect dimensions!")
|
|
|
|
if image.dtype != "float32":
|
|
raise Exception("Incorrect type!, expected float32")
|
|
|
|
try:
|
|
h, w = image.shape # expects Numpy convention (row, col) -> (height, width)
|
|
header = oe.Header(w, h)
|
|
header["channels"] = {"Z" : im.Channel(im.PixelType(oe.FLOAT))}
|
|
|
|
header['compression'] = im.Compression(comp)
|
|
|
|
of = oe.OutputFile(filePath, header)
|
|
image_data = image.tostring()
|
|
of.writePixels({"Z" : image_data})
|
|
of.close()
|
|
|
|
except:
|
|
raise
|
|
|
|
|
|
def load_rgb(filePath):
|
|
''' Loads an rgb OpenEXR image.'''
|
|
|
|
if oe.isOpenExrFile(filePath) is not True:
|
|
raise Exception("File ", filePath, " does not exist!")
|
|
|
|
try:
|
|
ifi = oe.InputFile(filePath)
|
|
|
|
# Compute size
|
|
header = ifi.header()
|
|
dw = header["dataWindow"]
|
|
w, h = dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1
|
|
|
|
# Read the three channels
|
|
ifiType = header["channels"]["R"].type.v
|
|
if ifiType is not im.PixelType.UINT:
|
|
raise Exception("Only uint32 supported! (file is type ", ifiType)
|
|
|
|
R = ifi.channel("R", im.PixelType(ifiType))
|
|
G = ifi.channel("G", im.PixelType(ifiType))
|
|
B = ifi.channel("B", im.PixelType(ifiType))
|
|
ifi.close()
|
|
|
|
image = np.zeros((h, w, 3), dtype = np.uint32) # order = "C"
|
|
image[:, :, 0] = np.core.multiarray.fromstring(R, dtype = np.uint32).reshape(h, w)
|
|
image[:, :, 1] = np.core.multiarray.fromstring(G, dtype = np.uint32).reshape(h, w)
|
|
image[:, :, 2] = np.core.multiarray.fromstring(B, dtype = np.uint32).reshape(h, w)
|
|
|
|
except:
|
|
raise
|
|
|
|
return image
|
|
|
|
|
|
def load_depth(filePath):
|
|
''' Loads an depth OpenEXR image.'''
|
|
|
|
if oe.isOpenExrFile(filePath) is not True:
|
|
raise Exception("File ", filePath, " does not exist!")
|
|
|
|
try:
|
|
ifi = oe.InputFile(filePath)
|
|
|
|
# Compute size
|
|
header = ifi.header()
|
|
dw = header["dataWindow"]
|
|
w, h = dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1
|
|
|
|
# Read the three channels
|
|
ifiType = header["channels"]["Z"].type.v
|
|
if ifiType is not im.PixelType.FLOAT:
|
|
raise Exception("Only float32 supported! (file is type ", ifiType)
|
|
|
|
Z = ifi.channel("Z", im.PixelType(ifiType))
|
|
ifi.close()
|
|
|
|
image = np.zeros((h, w), dtype = np.float32) # order = "C"
|
|
image[:, :] = np.core.multiarray.fromstring(Z, dtype = np.float32).reshape(h, w)
|
|
|
|
except:
|
|
raise
|
|
|
|
return image
|