mirror of https://github.com/RE-OWOD/RE-OWOD
177 lines
6.6 KiB
Python
177 lines
6.6 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
|
import logging
|
|
import unittest
|
|
import cv2
|
|
import torch
|
|
from torch.autograd import Variable, gradcheck
|
|
|
|
from detectron2.layers.roi_align import ROIAlign
|
|
from detectron2.layers.roi_align_rotated import ROIAlignRotated
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ROIAlignRotatedTest(unittest.TestCase):
|
|
def _box_to_rotated_box(self, box, angle):
|
|
return [
|
|
(box[0] + box[2]) / 2.0,
|
|
(box[1] + box[3]) / 2.0,
|
|
box[2] - box[0],
|
|
box[3] - box[1],
|
|
angle,
|
|
]
|
|
|
|
def _rot90(self, img, num):
|
|
num = num % 4 # note: -1 % 4 == 3
|
|
for _ in range(num):
|
|
img = img.transpose(0, 1).flip(0)
|
|
return img
|
|
|
|
def test_forward_output_0_90_180_270(self):
|
|
for i in range(4):
|
|
# i = 0, 1, 2, 3 corresponding to 0, 90, 180, 270 degrees
|
|
img = torch.arange(25, dtype=torch.float32).reshape(5, 5)
|
|
"""
|
|
0 1 2 3 4
|
|
5 6 7 8 9
|
|
10 11 12 13 14
|
|
15 16 17 18 19
|
|
20 21 22 23 24
|
|
"""
|
|
box = [1, 1, 3, 3]
|
|
rotated_box = self._box_to_rotated_box(box=box, angle=90 * i)
|
|
|
|
result = self._simple_roi_align_rotated(img=img, box=rotated_box, resolution=(4, 4))
|
|
|
|
# Here's an explanation for 0 degree case:
|
|
# point 0 in the original input lies at [0.5, 0.5]
|
|
# (the center of bin [0, 1] x [0, 1])
|
|
# point 1 in the original input lies at [1.5, 0.5], etc.
|
|
# since the resolution is (4, 4) that divides [1, 3] x [1, 3]
|
|
# into 4 x 4 equal bins,
|
|
# the top-left bin is [1, 1.5] x [1, 1.5], and its center
|
|
# (1.25, 1.25) lies at the 3/4 position
|
|
# between point 0 and point 1, point 5 and point 6,
|
|
# point 0 and point 5, point 1 and point 6, so it can be calculated as
|
|
# 0.25*(0*0.25+1*0.75)+(5*0.25+6*0.75)*0.75 = 4.5
|
|
result_expected = torch.tensor(
|
|
[
|
|
[4.5, 5.0, 5.5, 6.0],
|
|
[7.0, 7.5, 8.0, 8.5],
|
|
[9.5, 10.0, 10.5, 11.0],
|
|
[12.0, 12.5, 13.0, 13.5],
|
|
]
|
|
)
|
|
# This is also an upsampled version of [[6, 7], [11, 12]]
|
|
|
|
# When the box is rotated by 90 degrees CCW,
|
|
# the result would be rotated by 90 degrees CW, thus it's -i here
|
|
result_expected = self._rot90(result_expected, -i)
|
|
|
|
assert torch.allclose(result, result_expected)
|
|
|
|
def test_resize(self):
|
|
H, W = 30, 30
|
|
input = torch.rand(H, W) * 100
|
|
box = [10, 10, 20, 20]
|
|
rotated_box = self._box_to_rotated_box(box, angle=0)
|
|
output = self._simple_roi_align_rotated(img=input, box=rotated_box, resolution=(5, 5))
|
|
|
|
input2x = cv2.resize(input.numpy(), (W // 2, H // 2), interpolation=cv2.INTER_LINEAR)
|
|
input2x = torch.from_numpy(input2x)
|
|
box2x = [x / 2 for x in box]
|
|
rotated_box2x = self._box_to_rotated_box(box2x, angle=0)
|
|
output2x = self._simple_roi_align_rotated(img=input2x, box=rotated_box2x, resolution=(5, 5))
|
|
assert torch.allclose(output2x, output)
|
|
|
|
def _simple_roi_align_rotated(self, img, box, resolution):
|
|
"""
|
|
RoiAlignRotated with scale 1.0 and 0 sample ratio.
|
|
"""
|
|
op = ROIAlignRotated(output_size=resolution, spatial_scale=1.0, sampling_ratio=0)
|
|
input = img[None, None, :, :]
|
|
|
|
rois = [0] + list(box)
|
|
rois = torch.tensor(rois, dtype=torch.float32)[None, :]
|
|
result_cpu = op.forward(input, rois)
|
|
if torch.cuda.is_available():
|
|
result_cuda = op.forward(input.cuda(), rois.cuda())
|
|
assert torch.allclose(result_cpu, result_cuda.cpu())
|
|
return result_cpu[0, 0]
|
|
|
|
def test_empty_box(self):
|
|
img = torch.rand(5, 5)
|
|
out = self._simple_roi_align_rotated(img, [2, 3, 0, 0, 0], (7, 7))
|
|
self.assertTrue((out == 0).all())
|
|
|
|
def test_roi_align_rotated_gradcheck_cpu(self):
|
|
dtype = torch.float64
|
|
device = torch.device("cpu")
|
|
roi_align_rotated_op = ROIAlignRotated(
|
|
output_size=(5, 5), spatial_scale=0.5, sampling_ratio=1
|
|
).to(dtype=dtype, device=device)
|
|
x = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=True)
|
|
# roi format is (batch index, x_center, y_center, width, height, angle)
|
|
rois = torch.tensor(
|
|
[[0, 4.5, 4.5, 9, 9, 0], [0, 2, 7, 4, 4, 0], [0, 7, 7, 4, 4, 0]],
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
|
|
def func(input):
|
|
return roi_align_rotated_op(input, rois)
|
|
|
|
assert gradcheck(func, (x,)), "gradcheck failed for RoIAlignRotated CPU"
|
|
assert gradcheck(func, (x.transpose(2, 3),)), "gradcheck failed for RoIAlignRotated CPU"
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
|
|
def test_roi_align_rotated_gradient_cuda(self):
|
|
"""
|
|
Compute gradients for ROIAlignRotated with multiple bounding boxes on the GPU,
|
|
and compare the result with ROIAlign
|
|
"""
|
|
# torch.manual_seed(123)
|
|
dtype = torch.float64
|
|
device = torch.device("cuda")
|
|
pool_h, pool_w = (5, 5)
|
|
|
|
roi_align = ROIAlign(output_size=(pool_h, pool_w), spatial_scale=1, sampling_ratio=2).to(
|
|
device=device
|
|
)
|
|
|
|
roi_align_rotated = ROIAlignRotated(
|
|
output_size=(pool_h, pool_w), spatial_scale=1, sampling_ratio=2
|
|
).to(device=device)
|
|
|
|
x = torch.rand(1, 1, 10, 10, dtype=dtype, device=device, requires_grad=True)
|
|
# x_rotated = x.clone() won't work (will lead to grad_fun=CloneBackward)!
|
|
x_rotated = Variable(x.data.clone(), requires_grad=True)
|
|
|
|
# roi_rotated format is (batch index, x_center, y_center, width, height, angle)
|
|
rois_rotated = torch.tensor(
|
|
[[0, 4.5, 4.5, 9, 9, 0], [0, 2, 7, 4, 4, 0], [0, 7, 7, 4, 4, 0]],
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
|
|
y_rotated = roi_align_rotated(x_rotated, rois_rotated)
|
|
s_rotated = y_rotated.sum()
|
|
s_rotated.backward()
|
|
|
|
# roi format is (batch index, x1, y1, x2, y2)
|
|
rois = torch.tensor(
|
|
[[0, 0, 0, 9, 9], [0, 0, 5, 4, 9], [0, 5, 5, 9, 9]], dtype=dtype, device=device
|
|
)
|
|
|
|
y = roi_align(x, rois)
|
|
s = y.sum()
|
|
s.backward()
|
|
|
|
assert torch.allclose(
|
|
x.grad, x_rotated.grad
|
|
), "gradients for ROIAlign and ROIAlignRotated mismatch on CUDA"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|