168 lines
6.3 KiB
Python
168 lines
6.3 KiB
Python
import pytest
|
|
import torch
|
|
import torchvision.models
|
|
from common_utils import assert_equal
|
|
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor, TwoMLPHead
|
|
from torchvision.models.detection.roi_heads import RoIHeads
|
|
from torchvision.models.detection.rpn import AnchorGenerator, RegionProposalNetwork, RPNHead
|
|
from torchvision.ops import MultiScaleRoIAlign
|
|
|
|
|
|
class TestModelsDetectionNegativeSamples:
|
|
def _make_empty_sample(self, add_masks=False, add_keypoints=False):
|
|
images = [torch.rand((3, 100, 100), dtype=torch.float32)]
|
|
boxes = torch.zeros((0, 4), dtype=torch.float32)
|
|
negative_target = {
|
|
"boxes": boxes,
|
|
"labels": torch.zeros(0, dtype=torch.int64),
|
|
"image_id": 4,
|
|
"area": (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]),
|
|
"iscrowd": torch.zeros((0,), dtype=torch.int64),
|
|
}
|
|
|
|
if add_masks:
|
|
negative_target["masks"] = torch.zeros(0, 100, 100, dtype=torch.uint8)
|
|
|
|
if add_keypoints:
|
|
negative_target["keypoints"] = torch.zeros(17, 0, 3, dtype=torch.float32)
|
|
|
|
targets = [negative_target]
|
|
return images, targets
|
|
|
|
def test_targets_to_anchors(self):
|
|
_, targets = self._make_empty_sample()
|
|
anchors = [torch.randint(-50, 50, (3, 4), dtype=torch.float32)]
|
|
|
|
anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
|
|
aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
|
|
rpn_anchor_generator = AnchorGenerator(anchor_sizes, aspect_ratios)
|
|
rpn_head = RPNHead(4, rpn_anchor_generator.num_anchors_per_location()[0])
|
|
|
|
head = RegionProposalNetwork(rpn_anchor_generator, rpn_head, 0.5, 0.3, 256, 0.5, 2000, 2000, 0.7, 0.05)
|
|
|
|
labels, matched_gt_boxes = head.assign_targets_to_anchors(anchors, targets)
|
|
|
|
assert labels[0].sum() == 0
|
|
assert labels[0].shape == torch.Size([anchors[0].shape[0]])
|
|
assert labels[0].dtype == torch.float32
|
|
|
|
assert matched_gt_boxes[0].sum() == 0
|
|
assert matched_gt_boxes[0].shape == anchors[0].shape
|
|
assert matched_gt_boxes[0].dtype == torch.float32
|
|
|
|
def test_assign_targets_to_proposals(self):
|
|
|
|
proposals = [torch.randint(-50, 50, (20, 4), dtype=torch.float32)]
|
|
gt_boxes = [torch.zeros((0, 4), dtype=torch.float32)]
|
|
gt_labels = [torch.tensor([[0]], dtype=torch.int64)]
|
|
|
|
box_roi_pool = MultiScaleRoIAlign(featmap_names=["0", "1", "2", "3"], output_size=7, sampling_ratio=2)
|
|
|
|
resolution = box_roi_pool.output_size[0]
|
|
representation_size = 1024
|
|
box_head = TwoMLPHead(4 * resolution**2, representation_size)
|
|
|
|
representation_size = 1024
|
|
box_predictor = FastRCNNPredictor(representation_size, 2)
|
|
|
|
roi_heads = RoIHeads(
|
|
# Box
|
|
box_roi_pool,
|
|
box_head,
|
|
box_predictor,
|
|
0.5,
|
|
0.5,
|
|
512,
|
|
0.25,
|
|
None,
|
|
0.05,
|
|
0.5,
|
|
100,
|
|
)
|
|
|
|
matched_idxs, labels = roi_heads.assign_targets_to_proposals(proposals, gt_boxes, gt_labels)
|
|
|
|
assert matched_idxs[0].sum() == 0
|
|
assert matched_idxs[0].shape == torch.Size([proposals[0].shape[0]])
|
|
assert matched_idxs[0].dtype == torch.int64
|
|
|
|
assert labels[0].sum() == 0
|
|
assert labels[0].shape == torch.Size([proposals[0].shape[0]])
|
|
assert labels[0].dtype == torch.int64
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"fasterrcnn_resnet50_fpn",
|
|
"fasterrcnn_mobilenet_v3_large_fpn",
|
|
"fasterrcnn_mobilenet_v3_large_320_fpn",
|
|
],
|
|
)
|
|
def test_forward_negative_sample_frcnn(self, name):
|
|
model = torchvision.models.get_model(
|
|
name, weights=None, weights_backbone=None, num_classes=2, min_size=100, max_size=100
|
|
)
|
|
|
|
images, targets = self._make_empty_sample()
|
|
loss_dict = model(images, targets)
|
|
|
|
assert_equal(loss_dict["loss_box_reg"], torch.tensor(0.0))
|
|
assert_equal(loss_dict["loss_rpn_box_reg"], torch.tensor(0.0))
|
|
|
|
def test_forward_negative_sample_mrcnn(self):
|
|
model = torchvision.models.detection.maskrcnn_resnet50_fpn(
|
|
weights=None, weights_backbone=None, num_classes=2, min_size=100, max_size=100
|
|
)
|
|
|
|
images, targets = self._make_empty_sample(add_masks=True)
|
|
loss_dict = model(images, targets)
|
|
|
|
assert_equal(loss_dict["loss_box_reg"], torch.tensor(0.0))
|
|
assert_equal(loss_dict["loss_rpn_box_reg"], torch.tensor(0.0))
|
|
assert_equal(loss_dict["loss_mask"], torch.tensor(0.0))
|
|
|
|
def test_forward_negative_sample_krcnn(self):
|
|
model = torchvision.models.detection.keypointrcnn_resnet50_fpn(
|
|
weights=None, weights_backbone=None, num_classes=2, min_size=100, max_size=100
|
|
)
|
|
|
|
images, targets = self._make_empty_sample(add_keypoints=True)
|
|
loss_dict = model(images, targets)
|
|
|
|
assert_equal(loss_dict["loss_box_reg"], torch.tensor(0.0))
|
|
assert_equal(loss_dict["loss_rpn_box_reg"], torch.tensor(0.0))
|
|
assert_equal(loss_dict["loss_keypoint"], torch.tensor(0.0))
|
|
|
|
def test_forward_negative_sample_retinanet(self):
|
|
model = torchvision.models.detection.retinanet_resnet50_fpn(
|
|
weights=None, weights_backbone=None, num_classes=2, min_size=100, max_size=100
|
|
)
|
|
|
|
images, targets = self._make_empty_sample()
|
|
loss_dict = model(images, targets)
|
|
|
|
assert_equal(loss_dict["bbox_regression"], torch.tensor(0.0))
|
|
|
|
def test_forward_negative_sample_fcos(self):
|
|
model = torchvision.models.detection.fcos_resnet50_fpn(
|
|
weights=None, weights_backbone=None, num_classes=2, min_size=100, max_size=100
|
|
)
|
|
|
|
images, targets = self._make_empty_sample()
|
|
loss_dict = model(images, targets)
|
|
|
|
assert_equal(loss_dict["bbox_regression"], torch.tensor(0.0))
|
|
assert_equal(loss_dict["bbox_ctrness"], torch.tensor(0.0))
|
|
|
|
def test_forward_negative_sample_ssd(self):
|
|
model = torchvision.models.detection.ssd300_vgg16(weights=None, weights_backbone=None, num_classes=2)
|
|
|
|
images, targets = self._make_empty_sample()
|
|
loss_dict = model(images, targets)
|
|
|
|
assert_equal(loss_dict["bbox_regression"], torch.tensor(0.0))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|