def _load_crop_fg_probability()

in dataset/co3d_dataset.py [0:0]

• 15 lines of code
• 6 McCabe index (conditional complexity)

    def _load_crop_fg_probability(self, entry):
        clamp_bbox_xyxy = None
        if (self.load_masks or self.box_crop) and entry.mask is not None:
            full_path = os.path.join(self.dataset_root, entry.mask.path)
            mask = _load_mask(full_path)

            if mask.shape[-2:] != entry.image.size:
                raise ValueError(
                    f"bad mask size: {mask.shape[-2:]} vs {entry.image.size}!"
                )

            bbox_xywh = torch.tensor(_get_bbox_from_mask(mask, self.box_crop_mask_thr))

            if self.box_crop:
                clamp_bbox_xyxy = _get_clamp_bbox(bbox_xywh, self.box_crop_context)
                mask = _crop_around_box(mask, clamp_bbox_xyxy, full_path)

            fg_probability, _, _ = self._resize_image(mask, mode="nearest")
        return fg_probability, full_path, bbox_xywh, clamp_bbox_xyxy