mirror of https://github.com/FoundationVision/GLEE
263 lines
11 KiB
Python
263 lines
11 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
import logging
|
|
import unittest
|
|
import torch
|
|
|
|
from detectron2.config import get_cfg
|
|
from detectron2.export import scripting_with_instances
|
|
from detectron2.layers import ShapeSpec
|
|
from detectron2.modeling.backbone import build_backbone
|
|
from detectron2.modeling.proposal_generator import RPN, build_proposal_generator
|
|
from detectron2.modeling.proposal_generator.proposal_utils import (
|
|
add_ground_truth_to_proposals,
|
|
find_top_rpn_proposals,
|
|
)
|
|
from detectron2.structures import Boxes, ImageList, Instances, RotatedBoxes
|
|
from detectron2.utils.events import EventStorage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RPNTest(unittest.TestCase):
|
|
def get_gt_and_features(self):
|
|
num_images = 2
|
|
images_tensor = torch.rand(num_images, 20, 30)
|
|
image_sizes = [(10, 10), (20, 30)]
|
|
images = ImageList(images_tensor, image_sizes)
|
|
image_shape = (15, 15)
|
|
num_channels = 1024
|
|
features = {"res4": torch.rand(num_images, num_channels, 1, 2)}
|
|
gt_boxes = torch.tensor([[1, 1, 3, 3], [2, 2, 6, 6]], dtype=torch.float32)
|
|
gt_instances = Instances(image_shape)
|
|
gt_instances.gt_boxes = Boxes(gt_boxes)
|
|
return (gt_instances, features, images, image_sizes)
|
|
|
|
def test_rpn(self):
|
|
torch.manual_seed(121)
|
|
cfg = get_cfg()
|
|
backbone = build_backbone(cfg)
|
|
proposal_generator = RPN(cfg, backbone.output_shape())
|
|
(gt_instances, features, images, image_sizes) = self.get_gt_and_features()
|
|
with EventStorage(): # capture events in a new storage to discard them
|
|
proposals, proposal_losses = proposal_generator(
|
|
images, features, [gt_instances[0], gt_instances[1]]
|
|
)
|
|
|
|
expected_losses = {
|
|
"loss_rpn_cls": torch.tensor(0.08011703193),
|
|
"loss_rpn_loc": torch.tensor(0.101470276),
|
|
}
|
|
for name in expected_losses.keys():
|
|
err_msg = "proposal_losses[{}] = {}, expected losses = {}".format(
|
|
name, proposal_losses[name], expected_losses[name]
|
|
)
|
|
self.assertTrue(torch.allclose(proposal_losses[name], expected_losses[name]), err_msg)
|
|
|
|
self.assertEqual(len(proposals), len(image_sizes))
|
|
for proposal, im_size in zip(proposals, image_sizes):
|
|
self.assertEqual(proposal.image_size, im_size)
|
|
|
|
expected_proposal_box = torch.tensor([[0, 0, 10, 10], [7.2702, 0, 10, 10]])
|
|
expected_objectness_logit = torch.tensor([0.1596, -0.0007])
|
|
self.assertTrue(
|
|
torch.allclose(proposals[0].proposal_boxes.tensor, expected_proposal_box, atol=1e-4)
|
|
)
|
|
self.assertTrue(
|
|
torch.allclose(proposals[0].objectness_logits, expected_objectness_logit, atol=1e-4)
|
|
)
|
|
|
|
def verify_rpn(self, conv_dims, expected_conv_dims):
|
|
torch.manual_seed(121)
|
|
cfg = get_cfg()
|
|
cfg.MODEL.RPN.CONV_DIMS = conv_dims
|
|
backbone = build_backbone(cfg)
|
|
proposal_generator = RPN(cfg, backbone.output_shape())
|
|
for k, conv in enumerate(proposal_generator.rpn_head.conv):
|
|
self.assertEqual(expected_conv_dims[k], conv.out_channels)
|
|
return proposal_generator
|
|
|
|
def test_rpn_larger_num_convs(self):
|
|
conv_dims = [64, 64, 64, 64, 64]
|
|
proposal_generator = self.verify_rpn(conv_dims, conv_dims)
|
|
(gt_instances, features, images, image_sizes) = self.get_gt_and_features()
|
|
with EventStorage(): # capture events in a new storage to discard them
|
|
proposals, proposal_losses = proposal_generator(
|
|
images, features, [gt_instances[0], gt_instances[1]]
|
|
)
|
|
expected_losses = {
|
|
"loss_rpn_cls": torch.tensor(0.08122821152),
|
|
"loss_rpn_loc": torch.tensor(0.10064548254),
|
|
}
|
|
for name in expected_losses.keys():
|
|
err_msg = "proposal_losses[{}] = {}, expected losses = {}".format(
|
|
name, proposal_losses[name], expected_losses[name]
|
|
)
|
|
self.assertTrue(torch.allclose(proposal_losses[name], expected_losses[name]), err_msg)
|
|
|
|
def test_rpn_conv_dims_not_set(self):
|
|
conv_dims = [-1, -1, -1]
|
|
expected_conv_dims = [1024, 1024, 1024]
|
|
self.verify_rpn(conv_dims, expected_conv_dims)
|
|
|
|
def test_rpn_scriptability(self):
|
|
cfg = get_cfg()
|
|
proposal_generator = RPN(cfg, {"res4": ShapeSpec(channels=1024, stride=16)}).eval()
|
|
num_images = 2
|
|
images_tensor = torch.rand(num_images, 30, 40)
|
|
image_sizes = [(32, 32), (30, 40)]
|
|
images = ImageList(images_tensor, image_sizes)
|
|
features = {"res4": torch.rand(num_images, 1024, 1, 2)}
|
|
|
|
fields = {"proposal_boxes": Boxes, "objectness_logits": torch.Tensor}
|
|
proposal_generator_ts = scripting_with_instances(proposal_generator, fields)
|
|
|
|
proposals, _ = proposal_generator(images, features)
|
|
proposals_ts, _ = proposal_generator_ts(images, features)
|
|
|
|
for proposal, proposal_ts in zip(proposals, proposals_ts):
|
|
self.assertEqual(proposal.image_size, proposal_ts.image_size)
|
|
self.assertTrue(
|
|
torch.equal(proposal.proposal_boxes.tensor, proposal_ts.proposal_boxes.tensor)
|
|
)
|
|
self.assertTrue(torch.equal(proposal.objectness_logits, proposal_ts.objectness_logits))
|
|
|
|
def test_rrpn(self):
|
|
torch.manual_seed(121)
|
|
cfg = get_cfg()
|
|
cfg.MODEL.PROPOSAL_GENERATOR.NAME = "RRPN"
|
|
cfg.MODEL.ANCHOR_GENERATOR.NAME = "RotatedAnchorGenerator"
|
|
cfg.MODEL.ANCHOR_GENERATOR.SIZES = [[32, 64]]
|
|
cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS = [[0.25, 1]]
|
|
cfg.MODEL.ANCHOR_GENERATOR.ANGLES = [[0, 60]]
|
|
cfg.MODEL.RPN.BBOX_REG_WEIGHTS = (1, 1, 1, 1, 1)
|
|
cfg.MODEL.RPN.HEAD_NAME = "StandardRPNHead"
|
|
backbone = build_backbone(cfg)
|
|
proposal_generator = build_proposal_generator(cfg, backbone.output_shape())
|
|
num_images = 2
|
|
images_tensor = torch.rand(num_images, 20, 30)
|
|
image_sizes = [(10, 10), (20, 30)]
|
|
images = ImageList(images_tensor, image_sizes)
|
|
image_shape = (15, 15)
|
|
num_channels = 1024
|
|
features = {"res4": torch.rand(num_images, num_channels, 1, 2)}
|
|
gt_boxes = torch.tensor([[2, 2, 2, 2, 0], [4, 4, 4, 4, 0]], dtype=torch.float32)
|
|
gt_instances = Instances(image_shape)
|
|
gt_instances.gt_boxes = RotatedBoxes(gt_boxes)
|
|
with EventStorage(): # capture events in a new storage to discard them
|
|
proposals, proposal_losses = proposal_generator(
|
|
images, features, [gt_instances[0], gt_instances[1]]
|
|
)
|
|
|
|
expected_losses = {
|
|
"loss_rpn_cls": torch.tensor(0.04291602224),
|
|
"loss_rpn_loc": torch.tensor(0.145077362),
|
|
}
|
|
for name in expected_losses.keys():
|
|
err_msg = "proposal_losses[{}] = {}, expected losses = {}".format(
|
|
name, proposal_losses[name], expected_losses[name]
|
|
)
|
|
self.assertTrue(torch.allclose(proposal_losses[name], expected_losses[name]), err_msg)
|
|
|
|
expected_proposal_box = torch.tensor(
|
|
[
|
|
[-1.77999556, 0.78155339, 68.04367828, 14.78156471, 60.59333801],
|
|
[13.82740974, -1.50282836, 34.67269897, 29.19676590, -3.81942749],
|
|
[8.10392570, -0.99071521, 145.39100647, 32.13126373, 3.67242432],
|
|
[5.00000000, 4.57370186, 10.00000000, 9.14740372, 0.89196777],
|
|
]
|
|
)
|
|
|
|
expected_objectness_logit = torch.tensor([0.10924313, 0.09881870, 0.07649877, 0.05858029])
|
|
|
|
torch.set_printoptions(precision=8, sci_mode=False)
|
|
|
|
self.assertEqual(len(proposals), len(image_sizes))
|
|
|
|
proposal = proposals[0]
|
|
# It seems that there's some randomness in the result across different machines:
|
|
# This test can be run on a local machine for 100 times with exactly the same result,
|
|
# However, a different machine might produce slightly different results,
|
|
# thus the atol here.
|
|
err_msg = "computed proposal boxes = {}, expected {}".format(
|
|
proposal.proposal_boxes.tensor, expected_proposal_box
|
|
)
|
|
self.assertTrue(
|
|
torch.allclose(proposal.proposal_boxes.tensor[:4], expected_proposal_box, atol=1e-5),
|
|
err_msg,
|
|
)
|
|
|
|
err_msg = "computed objectness logits = {}, expected {}".format(
|
|
proposal.objectness_logits, expected_objectness_logit
|
|
)
|
|
self.assertTrue(
|
|
torch.allclose(proposal.objectness_logits[:4], expected_objectness_logit, atol=1e-5),
|
|
err_msg,
|
|
)
|
|
|
|
def test_find_rpn_proposals_inf(self):
|
|
N, Hi, Wi, A = 3, 3, 3, 3
|
|
proposals = [torch.rand(N, Hi * Wi * A, 4)]
|
|
pred_logits = [torch.rand(N, Hi * Wi * A)]
|
|
pred_logits[0][1][3:5].fill_(float("inf"))
|
|
find_top_rpn_proposals(proposals, pred_logits, [(10, 10)], 0.5, 1000, 1000, 0, False)
|
|
|
|
def test_find_rpn_proposals_tracing(self):
|
|
N, Hi, Wi, A = 3, 50, 50, 9
|
|
proposal = torch.rand(N, Hi * Wi * A, 4)
|
|
pred_logit = torch.rand(N, Hi * Wi * A)
|
|
|
|
def func(proposal, logit, image_size):
|
|
r = find_top_rpn_proposals(
|
|
[proposal], [logit], [image_size], 0.7, 1000, 1000, 0, False
|
|
)[0]
|
|
size = r.image_size
|
|
if not isinstance(size, torch.Tensor):
|
|
size = torch.tensor(size)
|
|
return (size, r.proposal_boxes.tensor, r.objectness_logits)
|
|
|
|
other_inputs = []
|
|
# test that it generalizes to other shapes
|
|
for Hi, Wi, shp in [(30, 30, 60), (10, 10, 800)]:
|
|
other_inputs.append(
|
|
(
|
|
torch.rand(N, Hi * Wi * A, 4),
|
|
torch.rand(N, Hi * Wi * A),
|
|
torch.tensor([shp, shp]),
|
|
)
|
|
)
|
|
torch.jit.trace(
|
|
func, (proposal, pred_logit, torch.tensor([100, 100])), check_inputs=other_inputs
|
|
)
|
|
|
|
def test_append_gt_to_proposal(self):
|
|
proposals = Instances(
|
|
(10, 10),
|
|
**{
|
|
"proposal_boxes": Boxes(torch.empty((0, 4))),
|
|
"objectness_logits": torch.tensor([]),
|
|
"custom_attribute": torch.tensor([]),
|
|
}
|
|
)
|
|
gt_boxes = Boxes(torch.tensor([[0, 0, 1, 1]]))
|
|
|
|
self.assertRaises(AssertionError, add_ground_truth_to_proposals, [gt_boxes], [proposals])
|
|
|
|
gt_instances = Instances((10, 10))
|
|
gt_instances.gt_boxes = gt_boxes
|
|
|
|
self.assertRaises(
|
|
AssertionError, add_ground_truth_to_proposals, [gt_instances], [proposals]
|
|
)
|
|
|
|
gt_instances.custom_attribute = torch.tensor([1])
|
|
gt_instances.custom_attribute2 = torch.tensor([1])
|
|
new_proposals = add_ground_truth_to_proposals([gt_instances], [proposals])[0]
|
|
|
|
self.assertEqual(new_proposals.custom_attribute[0], 1)
|
|
# new proposals should only include the attributes in proposals
|
|
self.assertRaises(AttributeError, lambda: new_proposals.custom_attribute2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|