def _load_mask_depth()

in dataset/co3d_dataset.py [0:0]

• 27 lines of code
• 5 McCabe index (conditional complexity)

    def _load_mask_depth(self, entry, clamp_bbox_xyxy, fg_probability):
        path = os.path.join(self.dataset_root, entry.depth.path)
        depth_map = _load_depth(path, entry.depth.scale_adjustment)

        if self.box_crop:
            depth_bbox_xyxy = _rescale_bbox(
                clamp_bbox_xyxy, entry.image.size, depth_map.shape[-2:]
            )
            depth_map = _crop_around_box(depth_map, depth_bbox_xyxy, path)

        depth_map, _, _ = self._resize_image(depth_map, mode="nearest")

        if self.mask_depths:
            assert fg_probability is not None
            depth_map *= fg_probability

        if self.load_depth_masks:
            assert entry.depth.mask_path is not None
            mask_path = os.path.join(self.dataset_root, entry.depth.mask_path)
            depth_mask = _load_depth_mask(mask_path)

            if self.box_crop:
                depth_mask_bbox_xyxy = _rescale_bbox(
                    clamp_bbox_xyxy, entry.image.size, depth_mask.shape[-2:]
                )
                depth_mask = _crop_around_box(
                    depth_mask, depth_mask_bbox_xyxy, mask_path
                )

            depth_mask, _, _ = self._resize_image(depth_mask, mode="nearest")
        else:
            depth_mask = torch.ones_like(depth_map)

        return depth_map, path, depth_mask