Fix: tensor vs array type correctness

Summary:
For fg-masking depth, we assumed np.array but passed a Tensor; for defining the default depth_mask, vice versa.

Note that we change the intended behaviour for the latter, assuming that 0s are areas with empty depth. When loading depth masks, we replace NaNs with zeros, so it is sensible. It is not a BC change as that branch would crash if executed. Since there was no reports, I assume no one cared.

Reviewed By: bottler

Differential Revision: D47403588

fbshipit-source-id: 1094104176d7d767a5657b5bbc9f5a0cc9da0ede
This commit is contained in:
Roman Shapovalov 2023-07-13 04:50:03 -07:00 committed by Facebook GitHub Bot
parent 9446d91fae
commit 8164ac4081

View File

@ -583,9 +583,9 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
else None, else None,
) )
fg_mask_np: Optional[np.ndarray] = None
mask_annotation = frame_annotation.mask mask_annotation = frame_annotation.mask
if mask_annotation is not None: if mask_annotation is not None:
fg_mask_np: Optional[np.ndarray] = None
if load_blobs and self.load_masks: if load_blobs and self.load_masks:
fg_mask_np, mask_path = self._load_fg_probability(frame_annotation) fg_mask_np, mask_path = self._load_fg_probability(frame_annotation)
frame_data.mask_path = mask_path frame_data.mask_path = mask_path
@ -627,7 +627,7 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
frame_data.depth_map, frame_data.depth_map,
frame_data.depth_path, frame_data.depth_path,
frame_data.depth_mask, frame_data.depth_mask,
) = self._load_mask_depth(frame_annotation, frame_data.fg_probability) ) = self._load_mask_depth(frame_annotation, fg_mask_np)
if load_blobs and self.load_point_clouds and point_cloud is not None: if load_blobs and self.load_point_clouds and point_cloud is not None:
pcl_path = self._fix_point_cloud_path(point_cloud.path) pcl_path = self._fix_point_cloud_path(point_cloud.path)
@ -683,7 +683,7 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
def _load_mask_depth( def _load_mask_depth(
self, self,
entry: types.FrameAnnotation, entry: types.FrameAnnotation,
fg_probability: Optional[torch.Tensor], fg_mask: Optional[np.ndarray],
) -> Tuple[torch.Tensor, str, torch.Tensor]: ) -> Tuple[torch.Tensor, str, torch.Tensor]:
entry_depth = entry.depth entry_depth = entry.depth
dataset_root = self.dataset_root dataset_root = self.dataset_root
@ -693,15 +693,15 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
depth_map = load_depth(self._local_path(path), entry_depth.scale_adjustment) depth_map = load_depth(self._local_path(path), entry_depth.scale_adjustment)
if self.mask_depths: if self.mask_depths:
assert fg_probability is not None assert fg_mask is not None
depth_map *= fg_probability depth_map *= fg_mask
mask_path = entry_depth.mask_path mask_path = entry_depth.mask_path
if self.load_depth_masks and mask_path is not None: if self.load_depth_masks and mask_path is not None:
mask_path = os.path.join(dataset_root, mask_path) mask_path = os.path.join(dataset_root, mask_path)
depth_mask = load_depth_mask(self._local_path(mask_path)) depth_mask = load_depth_mask(self._local_path(mask_path))
else: else:
depth_mask = torch.ones_like(depth_map) depth_mask = (depth_map > 0.0).astype(np.float32)
return torch.tensor(depth_map), path, torch.tensor(depth_mask) return torch.tensor(depth_map), path, torch.tensor(depth_mask)