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,
)
fg_mask_np: Optional[np.ndarray] = None
mask_annotation = frame_annotation.mask
if mask_annotation is not None:
fg_mask_np: Optional[np.ndarray] = None
if load_blobs and self.load_masks:
fg_mask_np, mask_path = self._load_fg_probability(frame_annotation)
frame_data.mask_path = mask_path
@ -627,7 +627,7 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
frame_data.depth_map,
frame_data.depth_path,
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:
pcl_path = self._fix_point_cloud_path(point_cloud.path)
@ -683,7 +683,7 @@ class GenericFrameDataBuilder(FrameDataBuilderBase[FrameDataSubtype], ABC):
def _load_mask_depth(
self,
entry: types.FrameAnnotation,
fg_probability: Optional[torch.Tensor],
fg_mask: Optional[np.ndarray],
) -> Tuple[torch.Tensor, str, torch.Tensor]:
entry_depth = entry.depth
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)
if self.mask_depths:
assert fg_probability is not None
depth_map *= fg_probability
assert fg_mask is not None
depth_map *= fg_mask
mask_path = entry_depth.mask_path
if self.load_depth_masks and mask_path is not None:
mask_path = os.path.join(dataset_root, mask_path)
depth_mask = load_depth_mask(self._local_path(mask_path))
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)