When bounding boxes are cached in metadata, don’t crash on load_masks=False

Summary:
We currently support caching bounding boxes in MaskAnnotation. If present, they are not re-computed from the mask. However, the masks need to be loaded for the bbox to be set.

This diff fixes that. Even if load_masks / load_blobs are unset, the bounding box can be picked up from the metadata.

Reviewed By: bottler

Differential Revision: D45144918

fbshipit-source-id: 8a2e2c115e96070b6fcdc29cbe57e1cee606ddcd
This commit is contained in:
Roman Shapovalov
2023-04-20 07:28:45 -07:00
committed by Facebook GitHub Bot
parent 0e3138eca8
commit 7aeedd17a4
2 changed files with 25 additions and 24 deletions

View File

@@ -17,6 +17,7 @@ from pytorch3d.implicitron.dataset import types
from pytorch3d.implicitron.dataset.dataset_base import FrameData
from pytorch3d.implicitron.dataset.frame_data import FrameDataBuilder
from pytorch3d.implicitron.dataset.utils import (
get_bbox_from_mask,
load_16big_png_depth,
load_1bit_png_mask,
load_depth,
@@ -107,11 +108,14 @@ class TestFrameDataBuilder(TestCaseMixin, unittest.TestCase):
)
self.frame_data.effective_image_size_hw = self.frame_data.image_size_hw
(
self.frame_data.fg_probability,
self.frame_data.mask_path,
self.frame_data.bbox_xywh,
) = self.frame_data_builder._load_fg_probability(self.frame_annotation)
fg_mask_np, mask_path = self.frame_data_builder._load_fg_probability(
self.frame_annotation
)
self.frame_data.mask_path = mask_path
self.frame_data.fg_probability = safe_as_tensor(fg_mask_np, torch.float)
mask_thr = self.frame_data_builder.box_crop_mask_thr
bbox_xywh = get_bbox_from_mask(fg_mask_np, mask_thr)
self.frame_data.bbox_xywh = safe_as_tensor(bbox_xywh, torch.long)
self.assertIsNotNone(self.frame_data.mask_path)
self.assertTrue(torch.is_tensor(self.frame_data.fg_probability))