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