|
|
|
|
|
import unittest |
|
import torch |
|
|
|
from detectron2.structures import Boxes, BoxMode, Instances |
|
|
|
from densepose.modeling.losses.embed_utils import CseAnnotationsAccumulator |
|
from densepose.structures import DensePoseDataRelative, DensePoseList |
|
|
|
|
|
class TestCseAnnotationsAccumulator(unittest.TestCase): |
|
def test_cse_annotations_accumulator_nodp(self): |
|
instances_lst = [ |
|
self._create_instances_nodp(), |
|
] |
|
self._test_template(instances_lst) |
|
|
|
def test_cse_annotations_accumulator_sparsedp(self): |
|
instances_lst = [ |
|
self._create_instances_sparsedp(), |
|
] |
|
self._test_template(instances_lst) |
|
|
|
def test_cse_annotations_accumulator_fulldp(self): |
|
instances_lst = [ |
|
self._create_instances_fulldp(), |
|
] |
|
self._test_template(instances_lst) |
|
|
|
def test_cse_annotations_accumulator_combined(self): |
|
instances_lst = [ |
|
self._create_instances_nodp(), |
|
self._create_instances_sparsedp(), |
|
self._create_instances_fulldp(), |
|
] |
|
self._test_template(instances_lst) |
|
|
|
def _test_template(self, instances_lst): |
|
acc = CseAnnotationsAccumulator() |
|
for instances in instances_lst: |
|
acc.accumulate(instances) |
|
packed_anns = acc.pack() |
|
self._check_correspondence(packed_anns, instances_lst) |
|
|
|
def _create_instances_nodp(self): |
|
image_shape = (480, 640) |
|
instances = Instances(image_shape) |
|
instances.gt_boxes = Boxes( |
|
torch.as_tensor( |
|
[ |
|
[40.0, 40.0, 140.0, 140.0], |
|
[160.0, 160.0, 270.0, 270.0], |
|
[40.0, 160.0, 160.0, 280.0], |
|
] |
|
) |
|
) |
|
instances.proposal_boxes = Boxes( |
|
torch.as_tensor( |
|
[ |
|
[41.0, 39.0, 142.0, 138.0], |
|
[161.0, 159.0, 272.0, 268.0], |
|
[41.0, 159.0, 162.0, 278.0], |
|
] |
|
) |
|
) |
|
|
|
return instances |
|
|
|
def _create_instances_sparsedp(self): |
|
image_shape = (540, 720) |
|
instances = Instances(image_shape) |
|
instances.gt_boxes = Boxes( |
|
torch.as_tensor( |
|
[ |
|
[50.0, 50.0, 130.0, 130.0], |
|
[150.0, 150.0, 240.0, 240.0], |
|
[50.0, 150.0, 230.0, 330.0], |
|
] |
|
) |
|
) |
|
instances.proposal_boxes = Boxes( |
|
torch.as_tensor( |
|
[ |
|
[49.0, 51.0, 131.0, 129.0], |
|
[151.0, 149.0, 241.0, 239.0], |
|
[51.0, 149.0, 232.0, 329.0], |
|
] |
|
) |
|
) |
|
instances.gt_densepose = DensePoseList( |
|
[ |
|
None, |
|
self._create_dp_data( |
|
{ |
|
"dp_x": [81.69, 153.47, 151.00], |
|
"dp_y": [162.24, 128.71, 113.81], |
|
"dp_vertex": [0, 1, 2], |
|
"ref_model": "zebra_5002", |
|
"dp_masks": [], |
|
}, |
|
{"c": (166, 133), "r": 64}, |
|
), |
|
None, |
|
], |
|
instances.gt_boxes, |
|
image_shape, |
|
) |
|
return instances |
|
|
|
def _create_instances_fulldp(self): |
|
image_shape = (680, 840) |
|
instances = Instances(image_shape) |
|
instances.gt_boxes = Boxes( |
|
torch.as_tensor( |
|
[ |
|
[65.0, 55.0, 165.0, 155.0], |
|
[170.0, 175.0, 275.0, 280.0], |
|
[55.0, 165.0, 165.0, 275.0], |
|
] |
|
) |
|
) |
|
instances.proposal_boxes = Boxes( |
|
torch.as_tensor( |
|
[ |
|
[66.0, 54.0, 166.0, 154.0], |
|
[171.0, 174.0, 276.0, 279.0], |
|
[56.0, 164.0, 166.0, 274.0], |
|
] |
|
) |
|
) |
|
instances.gt_densepose = DensePoseList( |
|
[ |
|
self._create_dp_data( |
|
{ |
|
"dp_x": [149.99, 198.62, 157.59], |
|
"dp_y": [170.74, 197.73, 123.12], |
|
"dp_vertex": [3, 4, 5], |
|
"ref_model": "cat_5001", |
|
"dp_masks": [], |
|
}, |
|
{"c": (100, 100), "r": 50}, |
|
), |
|
self._create_dp_data( |
|
{ |
|
"dp_x": [234.53, 116.72, 71.66], |
|
"dp_y": [107.53, 11.31, 142.32], |
|
"dp_vertex": [6, 7, 8], |
|
"ref_model": "dog_5002", |
|
"dp_masks": [], |
|
}, |
|
{"c": (200, 150), "r": 40}, |
|
), |
|
self._create_dp_data( |
|
{ |
|
"dp_x": [225.54, 202.61, 135.90], |
|
"dp_y": [167.46, 181.00, 211.47], |
|
"dp_vertex": [9, 10, 11], |
|
"ref_model": "elephant_5002", |
|
"dp_masks": [], |
|
}, |
|
{"c": (100, 200), "r": 45}, |
|
), |
|
], |
|
instances.gt_boxes, |
|
image_shape, |
|
) |
|
return instances |
|
|
|
def _create_dp_data(self, anns, blob_def=None): |
|
dp_data = DensePoseDataRelative(anns) |
|
if blob_def is not None: |
|
dp_data.segm[ |
|
blob_def["c"][0] - blob_def["r"] : blob_def["c"][0] + blob_def["r"], |
|
blob_def["c"][1] - blob_def["r"] : blob_def["c"][1] + blob_def["r"], |
|
] = 1 |
|
return dp_data |
|
|
|
def _check_correspondence(self, packed_anns, instances_lst): |
|
instance_idx = 0 |
|
data_idx = 0 |
|
pt_offset = 0 |
|
if packed_anns is not None: |
|
bbox_xyxy_gt = BoxMode.convert( |
|
packed_anns.bbox_xywh_gt.clone(), BoxMode.XYWH_ABS, BoxMode.XYXY_ABS |
|
) |
|
bbox_xyxy_est = BoxMode.convert( |
|
packed_anns.bbox_xywh_est.clone(), BoxMode.XYWH_ABS, BoxMode.XYXY_ABS |
|
) |
|
for instances in instances_lst: |
|
if not hasattr(instances, "gt_densepose"): |
|
instance_idx += len(instances) |
|
continue |
|
for i, dp_data in enumerate(instances.gt_densepose): |
|
if dp_data is None: |
|
instance_idx += 1 |
|
continue |
|
n_pts = len(dp_data.x) |
|
self.assertTrue( |
|
torch.allclose(dp_data.x, packed_anns.x_gt[pt_offset : pt_offset + n_pts]) |
|
) |
|
self.assertTrue( |
|
torch.allclose(dp_data.y, packed_anns.y_gt[pt_offset : pt_offset + n_pts]) |
|
) |
|
self.assertTrue(torch.allclose(dp_data.segm, packed_anns.coarse_segm_gt[data_idx])) |
|
self.assertTrue( |
|
torch.allclose( |
|
torch.ones(n_pts, dtype=torch.long) * dp_data.mesh_id, |
|
packed_anns.vertex_mesh_ids_gt[pt_offset : pt_offset + n_pts], |
|
) |
|
) |
|
self.assertTrue( |
|
torch.allclose( |
|
dp_data.vertex_ids, packed_anns.vertex_ids_gt[pt_offset : pt_offset + n_pts] |
|
) |
|
) |
|
self.assertTrue( |
|
torch.allclose(instances.gt_boxes.tensor[i], bbox_xyxy_gt[data_idx]) |
|
) |
|
self.assertTrue( |
|
torch.allclose(instances.proposal_boxes.tensor[i], bbox_xyxy_est[data_idx]) |
|
) |
|
self.assertTrue( |
|
torch.allclose( |
|
torch.ones(n_pts, dtype=torch.long) * data_idx, |
|
packed_anns.point_bbox_with_dp_indices[pt_offset : pt_offset + n_pts], |
|
) |
|
) |
|
self.assertTrue( |
|
torch.allclose( |
|
torch.ones(n_pts, dtype=torch.long) * instance_idx, |
|
packed_anns.point_bbox_indices[pt_offset : pt_offset + n_pts], |
|
) |
|
) |
|
self.assertEqual(instance_idx, packed_anns.bbox_indices[data_idx]) |
|
pt_offset += n_pts |
|
instance_idx += 1 |
|
data_idx += 1 |
|
if data_idx == 0: |
|
self.assertIsNone(packed_anns) |
|
|
