Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 25 additions & 8 deletions exr_info/cryptomatte.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def __init__(self, exr_f: ExrInfo):

# In the manifest, some entried are added by vray, which should be ignored.
self.IGNORE_OBJS_IN_MANIFEST = ["vrayLightDome", "vrayLightMesh", "default"]
self.obj_masks = None # Save the extracted masks of all objects. Might be used by both alpha and segments.

@staticmethod
def get_coverage_for_rank(float_id: float, cr_combined: np.ndarray, rank: int) -> np.ndarray:
Expand Down Expand Up @@ -69,7 +70,7 @@ def get_mask_for_id(self, obj_hex_id: str, channels_arr: np.ndarray, level: int
will change depending on level.

Returns:
numpy.ndarray: Mask from cryptomatte for a given id. Dtype: np.uint8, Range: [0, 255]
numpy.ndarray: Mask from cryptomatte for a given id. dtype: np.float32, Range: [0, 1]
"""
float_id = self._convert_hex_id_to_float_id(obj_hex_id)

Expand All @@ -80,7 +81,11 @@ def get_mask_for_id(self, obj_hex_id: str, channels_arr: np.ndarray, level: int

coverage = sum(coverage_list)
coverage = np.clip(coverage, 0.0, 1.0)
mask = (coverage * 255).astype(np.uint8)
mask = coverage.astype(np.float32)

# Note: Max value of coverage was 1.0002, which rounded to 1.0 when casting to float16. Might imply that the
# actual data is float16.

return mask

def get_masks_for_all_objs(self, crypto_def_idx) -> OrderedDict:
Expand Down Expand Up @@ -118,7 +123,7 @@ def get_masks_for_all_objs(self, crypto_def_idx) -> OrderedDict:

return obj_masks

def get_combined_mask(self, crypto_def_idx: int = 0) -> Tuple[np.ndarray, Dict[str, int]]:
def get_segments_map(self, crypto_def_idx: int = 0) -> Tuple[np.ndarray, Dict[str, int]]:
"""
Get a single mask representing all the objects within the scene.
Each object is represented by a unique integer value, starting from 1. 0 is reserved for background.
Expand All @@ -131,24 +136,36 @@ def get_combined_mask(self, crypto_def_idx: int = 0) -> Tuple[np.ndarray, Dict[s
numpy.ndarray: Mask of all objects. Shape: [H, W], dtype: np.uint16.
dict: Mapping of the object names to mask IDs for this image.
"""
obj_masks = self.get_masks_for_all_objs(crypto_def_idx)
if self.obj_masks is None:
self.obj_masks = self.get_masks_for_all_objs(crypto_def_idx)

# Create a map of obj names to ids
name_to_mask_id_map = OrderedDict()
name_to_mask_id_map["background"] = 0 # Background is always class 0
obj_names = obj_masks.keys()
obj_names = self.obj_masks.keys()
for idx, obj_name in enumerate(obj_names):
name_to_mask_id_map[obj_name] = idx + 1

# Combine all the masks into single mask without anti-aliasing for semantic segmentation
masks = np.stack(list(obj_masks.values()), axis=0) # Shape: [N, H, W]
background_mask = 255 - masks.sum(axis=0)
# Use argmax to find which class each pixel belongs to.
masks = np.stack(list(self.obj_masks.values()), axis=0) # Shape: [N, H, W]
background_mask = 1.0 - masks.sum(axis=0)
masks = np.concatenate((np.expand_dims(background_mask, 0), masks), axis=0)
mask_combined = masks.argmax(axis=0)
mask_combined = mask_combined.astype(np.uint16)

# Convert to uint
mask_combined = (mask_combined * 255).astype(np.uint16)

return mask_combined, name_to_mask_id_map

def get_alpha_map(self, crypto_def_idx: int = 0) -> np.ndarray:
if self.obj_masks is None:
self.obj_masks = self.get_masks_for_all_objs(crypto_def_idx)

masks = np.stack(list(self.obj_masks.values()), axis=0) # Shape: [N, H, W]
alpha = np.sum(masks, axis=0)
return alpha

@staticmethod
def apply_random_colormap_to_mask(mask_combined: np.ndarray) -> np.ndarray:
"""
Expand Down
2 changes: 2 additions & 0 deletions exr_info/exr_info.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,6 +333,7 @@ def read_channels(self, channel_names: Iterator[str], cast_dtype: Optional[ExrDt
Shape: [H, W]
"""
# TODO: Can be made more efficient by reading all channels at once: https://excamera.com/articles/26/doc/openexr.html#OpenEXR.InputFile.channels
# Get the precision (dtype) of each channel in exr from the exr header.
channel_dtypes = []
for channel_name in channel_names:
channel_dtypes.append(self.get_channel_precision(channel_name))
Expand All @@ -344,6 +345,7 @@ def read_channels(self, channel_names: Iterator[str], cast_dtype: Optional[ExrDt
channel_arr = channel_arr.reshape((self.height, self.width))

if cast_dtype is not None:
# Forcefully cast array to a given datatype
if not isinstance(cast_dtype, ExrDtype):
raise ValueError(f"Expected type {ExrDtype.__name__}. Got: {type(cast_dtype)}")
channel_arr = channel_arr.astype(numpy_dtype[cast_dtype])
Expand Down