Spaces:
Sleeping
Sleeping
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import numpy as np | |
| import unittest | |
| from copy import copy | |
| import cv2 | |
| import torch | |
| from fvcore.common.benchmark import benchmark | |
| from torch.nn import functional as F | |
| from detectron2.layers.roi_align import ROIAlign, roi_align | |
| class ROIAlignTest(unittest.TestCase): | |
| def test_forward_output(self): | |
| input = np.arange(25).reshape(5, 5).astype("float32") | |
| """ | |
| 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 | |
| """ | |
| output = self._simple_roialign(input, [1, 1, 3, 3], (4, 4), aligned=False) | |
| output_correct = self._simple_roialign(input, [1, 1, 3, 3], (4, 4), aligned=True) | |
| # without correction: | |
| old_results = [ | |
| [7.5, 8, 8.5, 9], | |
| [10, 10.5, 11, 11.5], | |
| [12.5, 13, 13.5, 14], | |
| [15, 15.5, 16, 16.5], | |
| ] | |
| # with 0.5 correction: | |
| correct_results = [ | |
| [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 an upsampled version of [[6, 7], [11, 12]] | |
| self.assertTrue(np.allclose(output.flatten(), np.asarray(old_results).flatten())) | |
| self.assertTrue( | |
| np.allclose(output_correct.flatten(), np.asarray(correct_results).flatten()) | |
| ) | |
| # Also see similar issues in tensorflow at | |
| # https://github.com/tensorflow/tensorflow/issues/26278 | |
| def test_resize(self): | |
| H, W = 30, 30 | |
| input = np.random.rand(H, W).astype("float32") * 100 | |
| box = [10, 10, 20, 20] | |
| output = self._simple_roialign(input, box, (5, 5), aligned=True) | |
| input2x = cv2.resize(input, (W // 2, H // 2), interpolation=cv2.INTER_LINEAR) | |
| box2x = [x / 2 for x in box] | |
| output2x = self._simple_roialign(input2x, box2x, (5, 5), aligned=True) | |
| diff = np.abs(output2x - output) | |
| self.assertTrue(diff.max() < 1e-4) | |
| def test_grid_sample_equivalence(self): | |
| H, W = 30, 30 | |
| input = np.random.rand(H, W).astype("float32") * 100 | |
| box = [10, 10, 20, 20] | |
| for ratio in [1, 2, 3]: | |
| output = self._simple_roialign(input, box, (5, 5), sampling_ratio=ratio) | |
| output_grid_sample = grid_sample_roi_align( | |
| torch.from_numpy(input[None, None, :, :]).float(), | |
| torch.as_tensor(box).float()[None, :], | |
| 5, | |
| 1.0, | |
| ratio, | |
| ) | |
| self.assertTrue(torch.allclose(output, output_grid_sample)) | |
| def _simple_roialign(self, img, box, resolution, sampling_ratio=0, aligned=True): | |
| """ | |
| RoiAlign with scale 1.0. | |
| """ | |
| if isinstance(resolution, int): | |
| resolution = (resolution, resolution) | |
| op = ROIAlign(resolution, 1.0, sampling_ratio, aligned=aligned) | |
| input = torch.from_numpy(img[None, None, :, :].astype("float32")) | |
| rois = [0] + list(box) | |
| rois = torch.from_numpy(np.asarray(rois)[None, :].astype("float32")) | |
| output = op.forward(input, rois) | |
| if torch.cuda.is_available(): | |
| output_cuda = op.forward(input.cuda(), rois.cuda()).cpu() | |
| self.assertTrue(torch.allclose(output, output_cuda)) | |
| return output[0, 0] | |
| def _simple_roialign_with_grad(self, img, box, resolution, device): | |
| if isinstance(resolution, int): | |
| resolution = (resolution, resolution) | |
| op = ROIAlign(resolution, 1.0, 0, aligned=True) | |
| input = torch.from_numpy(img[None, None, :, :].astype("float32")) | |
| rois = [0] + list(box) | |
| rois = torch.from_numpy(np.asarray(rois)[None, :].astype("float32")) | |
| input = input.to(device=device) | |
| rois = rois.to(device=device) | |
| input.requires_grad = True | |
| output = op.forward(input, rois) | |
| return input, output | |
| def test_empty_box(self): | |
| img = np.random.rand(5, 5) | |
| box = [3, 4, 5, 4] | |
| o = self._simple_roialign(img, box, 7) | |
| self.assertTrue(o.shape == (7, 7)) | |
| self.assertTrue((o == 0).all()) | |
| for dev in ["cpu"] + ["cuda"] if torch.cuda.is_available() else []: | |
| input, output = self._simple_roialign_with_grad(img, box, 7, torch.device(dev)) | |
| output.sum().backward() | |
| self.assertTrue(torch.allclose(input.grad, torch.zeros_like(input))) | |
| def test_empty_batch(self): | |
| input = torch.zeros(0, 3, 10, 10, dtype=torch.float32) | |
| rois = torch.zeros(0, 5, dtype=torch.float32) | |
| op = ROIAlign((7, 7), 1.0, 0, aligned=True) | |
| output = op.forward(input, rois) | |
| self.assertTrue(output.shape == (0, 3, 7, 7)) | |
| def grid_sample_roi_align(input, boxes, output_size, scale, sampling_ratio): | |
| # unlike true roi_align, this does not support different batch_idx | |
| from detectron2.projects.point_rend.point_features import ( | |
| generate_regular_grid_point_coords, | |
| get_point_coords_wrt_image, | |
| point_sample, | |
| ) | |
| N, _, H, W = input.shape | |
| R = len(boxes) | |
| assert N == 1 | |
| boxes = boxes * scale | |
| grid = generate_regular_grid_point_coords(R, output_size * sampling_ratio, device=boxes.device) | |
| coords = get_point_coords_wrt_image(boxes, grid) | |
| coords = coords / torch.as_tensor([W, H], device=coords.device) # R, s^2, 2 | |
| res = point_sample(input, coords.unsqueeze(0), align_corners=False) # 1,C, R,s^2 | |
| res = ( | |
| res.squeeze(0) | |
| .permute(1, 0, 2) | |
| .reshape(R, -1, output_size * sampling_ratio, output_size * sampling_ratio) | |
| ) | |
| res = F.avg_pool2d(res, sampling_ratio) | |
| return res | |
| def benchmark_roi_align(): | |
| def random_boxes(mean_box, stdev, N, maxsize): | |
| ret = torch.rand(N, 4) * stdev + torch.tensor(mean_box, dtype=torch.float) | |
| ret.clamp_(min=0, max=maxsize) | |
| return ret | |
| def func(shape, nboxes_per_img, sampling_ratio, device, box_size="large"): | |
| N, _, H, _ = shape | |
| input = torch.rand(*shape) | |
| boxes = [] | |
| batch_idx = [] | |
| for k in range(N): | |
| if box_size == "large": | |
| b = random_boxes([80, 80, 130, 130], 24, nboxes_per_img, H) | |
| else: | |
| b = random_boxes([100, 100, 110, 110], 4, nboxes_per_img, H) | |
| boxes.append(b) | |
| batch_idx.append(torch.zeros(nboxes_per_img, 1, dtype=torch.float32) + k) | |
| boxes = torch.cat(boxes, axis=0) | |
| batch_idx = torch.cat(batch_idx, axis=0) | |
| boxes = torch.cat([batch_idx, boxes], axis=1) | |
| input = input.to(device=device) | |
| boxes = boxes.to(device=device) | |
| def bench(): | |
| if False and sampling_ratio > 0 and N == 1: | |
| # enable to benchmark grid_sample (slower) | |
| grid_sample_roi_align(input, boxes[:, 1:], 7, 1.0, sampling_ratio) | |
| else: | |
| roi_align(input, boxes, 7, 1.0, sampling_ratio, True) | |
| if device == "cuda": | |
| torch.cuda.synchronize() | |
| return bench | |
| def gen_args(arg): | |
| args = [] | |
| for size in ["small", "large"]: | |
| for ratio in [0, 2]: | |
| args.append(copy(arg)) | |
| args[-1]["sampling_ratio"] = ratio | |
| args[-1]["box_size"] = size | |
| return args | |
| arg = dict(shape=(1, 512, 256, 256), nboxes_per_img=512, device="cuda") | |
| benchmark(func, "cuda_roialign", gen_args(arg), num_iters=20, warmup_iters=1) | |
| arg.update({"device": "cpu", "shape": (1, 256, 128, 128)}) | |
| benchmark(func, "cpu_roialign", gen_args(arg), num_iters=5, warmup_iters=1) | |
| if __name__ == "__main__": | |
| if torch.cuda.is_available(): | |
| benchmark_roi_align() | |
| unittest.main() | |