def exr_channel_to_numpy(exr_file: OpenEXR.InputFile,
channel_name: str,
reshape: Optional[Tuple[int, int]],
dtype: ExrDtype = ExrDtype.FLOAT32) -> np.ndarray:
"""Extracts a channel in an EXR file into a numpy array
Args:
exr_file (OpenEXR.InputFile): The opened EXR file object
channel_name (str): The name of the channel to be converted to numpy
reshape (None or Tuple(height, width): If given, will reshape the 2D numpy array into (height, width)
dtype (ExrDtype): Whether the data in channel is of float32 or float16 type
Returns:
numpy.ndarray: The extracted channel in form of numpy array
"""
if dtype == ExrDtype.FLOAT32:
point_type = Imath.PixelType(Imath.PixelType.FLOAT)
np_type = np.float32
else:
point_type = Imath.PixelType(Imath.PixelType.HALF)
np_type = np.float16
channel_arr = np.frombuffer(exr_file.channel(channel_name, point_type),
dtype=np_type)
if reshape:
channel_arr = channel_arr.reshape(reshape)
return channel_arr
def get_pixels_data(self, exr: OpenEXR.InputFile) -> dict:
pixels_data = {}
for i, target_channel in enumerate(self.get_target_channels()):
label_index = i
if self.is_depth():
label_index = 0
pixels_data[target_channel] = exr.channel(self.get_labels()[label_index], self.get_pixel_type())
return pixels_data
def read_channels_from_exr(exr: OpenEXR.InputFile,
channel_names: Sequence[str]) -> np.ndarray:
"""Reads a single channel from an EXR file and returns it as a numpy array."""
channels_header = exr.header()['channels']
window = exr.header()['dataWindow']
width = window.max.x - window.min.x + 1
height = window.max.y - window.min.y + 1
outputs = []
for channel_name in channel_names:
channel_type = channels_header[channel_name].type.v
numpy_type = {
Imath.PixelType.HALF: np.float16,
Imath.PixelType.FLOAT: np.float32,
Imath.PixelType.UINT: np.uint32,
}[channel_type]
array = np.frombuffer(exr.channel(channel_name), numpy_type)
array = array.reshape([height, width])
outputs.append(array)
# TODO: verify that the types are all the same?
return np.stack(outputs, axis=-1)
def __getitem__(self, index):
if isinstance(index, str):
return super(ImageSequence, self).__getitem__(index)
if index == self.cache_index:
return self.cache * numpy.exp(self.exposure_offset)
if not isinstance(index, int) and index in self.frame_range:
raise IndexError
exrfile = InputFile(self.file_pattern % index)
channels = ('R', 'G', 'B')
output = numpy.ndarray(dtype=numpy.float, shape=(self.height, self.width, channels.__len__()))
for channel_index, channel in enumerate(channels):
output[:, :, channel_index] = numpy.fromstring(exrfile.channel(channel, Imath.PixelType(Imath.PixelType.HALF)), dtype=numpy.float16).astype(numpy.float).reshape([self.height, self.width])
self.cache = output
self.cache_index = index
return output * numpy.exp(self.exposure_offset)
def get_imsize(exr_file: OpenEXR.InputFile) -> Tuple[int, int]:
"""Get the height and width of image within and EXR file
Args:
exr_file (OpenEXR.InputFile): The opened EXR file object.
Returns:
int, int: Height, Width of image
"""
header = exr_file.header()
dw = header['dataWindow']
height = int(dw.max.y - dw.min.y + 1)
width = int(dw.max.x - dw.min.x + 1)
return height, width
def exr_to_np(fname, rem=False):
return channels_to_ndarray(InputFile(fname), "RGBA")
def setup(self): sample = InputFile(self.file_pattern % self.frame_range.lower_limit) self.header = sample.header() dw = self.header['dataWindow'] (self.width, self.height) = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1) self.shape = (self.height, self.width, len(sample.header()['channels']))
def extract_mask(
exr_file: OpenEXR.InputFile,
extract_id_mapping_from_manifest: bool = True
) -> Tuple[np.ndarray, np.ndarray, Dict]:
"""Get a mask of all the objects in an EXR image from the cryptomatte
Args:
exr_file (OpenEXR.InputFile): The opened EXR file object
extract_id_mapping_from_manifest (bool): In latest renders, the ID that each object in cryptomatte maps to
in the output mask is present in the object's name in the manifest.
Eg: cords_2, means that the object is of type "cords" and it maps to
a value of 2 in the output mask
Returns:
numpy.ndarray: Mask of all objects in scene. Each object has a unique value. Dtype: np.float16, Shape: (H, W)
numpy.ndarray: RGB visualization of the mask. Dtype: np.uint8, Shape: (H, W, 3)
dict[]
"""
# Get the manifest (mapping of object names to hash ids)
header = exr_file.header()
manifest = None
for key in header:
if MANIFEST_IDENTIFIER in key:
manifest = json.loads(header[key], object_pairs_hook=OrderedDict)
break
if manifest is None:
raise RuntimeError(
'The EXR file\'s header does not contain the manifest for cryptomattes'
)
# Convert hash ids to float ids
float_ids = []
obj_ids = []
for ii, obj_name in enumerate(sorted(manifest)):
hex_id = manifest[obj_name]
bytes_val = bytes.fromhex(hex_id)
float_val = struct.unpack('>f', bytes_val)[0]
float_ids.append(float_val)
obj_ids.append(ii)
# Extract the crypto layers from EXR file
cr_combined = get_crypto_layers_from_exr(exr_file)
# Extract mask of each object in manifest
mask_list = []
id_list = []
id_mapping = OrderedDict() # Mapping the name of each obj to obj id
for float_id, obj_id, obj_name in zip(float_ids, obj_ids,
sorted(manifest)):
# Ignore the vrayLightDome.
if obj_name == 'vrayLightDome':
continue
mask = get_mask_for_id(float_id, cr_combined)
mask_list.append(mask)
if extract_id_mapping_from_manifest:
# The object type and ID is encoded in the object's name in manifest
obj_type, obj_id_manifest = obj_name.split('_')
id_mapping[obj_name] = int(obj_id_manifest)
id_list.append(int(obj_id_manifest))
else:
# The ID will be same as index of object in manifest. IDs extracted like should should be saved in a
# separate file so that the mapping from IDs to objects is available
id_mapping[obj_name] = obj_id
id_list.append(obj_id)
# Combine all the masks into single mask
height, width = get_imsize(exr_file)
mask_combined = np.zeros((height, width), dtype=np.uint16)
mask_combined_rgb = np.zeros((height, width, 3), dtype=np.uint8)
for mask, obj_id in zip(mask_list, id_list):
mask_combined[mask > MASK_THRESHOLD] = obj_id
hue = random.random()
sat, val = 0.7, 0.7
r, g, b = colorsys.hsv_to_rgb(hue, sat, val)
rgb = []
for col in [r, g, b]:
col_np = np.array(col, dtype=np.float32)
col_np = (np.clip(col_np * 255, 0, 255)).astype(np.uint8)
rgb.append(col_np)
mask_combined_rgb[mask > MASK_THRESHOLD, :] = rgb
return mask_combined, mask_combined_rgb, id_mapping
