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 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 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
