Add at new repo again
This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import numpy as np
|
||||
import unittest
|
||||
from collections import defaultdict
|
||||
import torch
|
||||
import tqdm
|
||||
from fvcore.common.benchmark import benchmark
|
||||
from fvcore.common.file_io import PathManager
|
||||
from pycocotools.coco import COCO
|
||||
from tabulate import tabulate
|
||||
from torch.nn import functional as F
|
||||
|
||||
from detectron2.data import MetadataCatalog
|
||||
from detectron2.layers.mask_ops import (
|
||||
pad_masks,
|
||||
paste_mask_in_image_old,
|
||||
paste_masks_in_image,
|
||||
scale_boxes,
|
||||
)
|
||||
from detectron2.structures import BitMasks, Boxes, BoxMode, PolygonMasks
|
||||
from detectron2.structures.masks import polygons_to_bitmask
|
||||
|
||||
|
||||
def iou_between_full_image_bit_masks(a, b):
|
||||
intersect = (a & b).sum()
|
||||
union = (a | b).sum()
|
||||
return intersect / union
|
||||
|
||||
|
||||
def rasterize_polygons_with_grid_sample(full_image_bit_mask, box, mask_size, threshold=0.5):
|
||||
x0, y0, x1, y1 = box[0], box[1], box[2], box[3]
|
||||
|
||||
img_h, img_w = full_image_bit_mask.shape
|
||||
|
||||
mask_y = np.arange(0.0, mask_size) + 0.5 # mask y sample coords in [0.5, mask_size - 0.5]
|
||||
mask_x = np.arange(0.0, mask_size) + 0.5 # mask x sample coords in [0.5, mask_size - 0.5]
|
||||
mask_y = mask_y / mask_size * (y1 - y0) + y0
|
||||
mask_x = mask_x / mask_size * (x1 - x0) + x0
|
||||
|
||||
mask_x = (mask_x - 0.5) / (img_w - 1) * 2 + -1
|
||||
mask_y = (mask_y - 0.5) / (img_h - 1) * 2 + -1
|
||||
gy, gx = torch.meshgrid(torch.from_numpy(mask_y), torch.from_numpy(mask_x))
|
||||
ind = torch.stack([gx, gy], dim=-1).to(dtype=torch.float32)
|
||||
|
||||
full_image_bit_mask = torch.from_numpy(full_image_bit_mask)
|
||||
mask = F.grid_sample(
|
||||
full_image_bit_mask[None, None, :, :].to(dtype=torch.float32),
|
||||
ind[None, :, :, :],
|
||||
align_corners=True,
|
||||
)
|
||||
|
||||
return mask[0, 0] >= threshold
|
||||
|
||||
|
||||
class TestMaskCropPaste(unittest.TestCase):
|
||||
def setUp(self):
|
||||
json_file = MetadataCatalog.get("coco_2017_val_100").json_file
|
||||
if not PathManager.isfile(json_file):
|
||||
raise unittest.SkipTest("{} not found".format(json_file))
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
json_file = PathManager.get_local_path(json_file)
|
||||
self.coco = COCO(json_file)
|
||||
|
||||
def test_crop_paste_consistency(self):
|
||||
"""
|
||||
rasterize_polygons_within_box (used in training)
|
||||
and
|
||||
paste_masks_in_image (used in inference)
|
||||
should be inverse operations to each other.
|
||||
|
||||
This function runs several implementation of the above two operations and prints
|
||||
the reconstruction error.
|
||||
"""
|
||||
|
||||
anns = self.coco.loadAnns(self.coco.getAnnIds(iscrowd=False)) # avoid crowd annotations
|
||||
|
||||
selected_anns = anns[:100]
|
||||
|
||||
ious = []
|
||||
for ann in tqdm.tqdm(selected_anns):
|
||||
results = self.process_annotation(ann)
|
||||
ious.append([k[2] for k in results])
|
||||
|
||||
ious = np.array(ious)
|
||||
mean_ious = ious.mean(axis=0)
|
||||
table = []
|
||||
res_dic = defaultdict(dict)
|
||||
for row, iou in zip(results, mean_ious):
|
||||
table.append((row[0], row[1], iou))
|
||||
res_dic[row[0]][row[1]] = iou
|
||||
print(tabulate(table, headers=["rasterize", "paste", "iou"], tablefmt="simple"))
|
||||
# assert that the reconstruction is good:
|
||||
self.assertTrue(res_dic["polygon"]["aligned"] > 0.94)
|
||||
self.assertTrue(res_dic["roialign"]["aligned"] > 0.95)
|
||||
|
||||
def process_annotation(self, ann, mask_side_len=28):
|
||||
# Parse annotation data
|
||||
img_info = self.coco.loadImgs(ids=[ann["image_id"]])[0]
|
||||
height, width = img_info["height"], img_info["width"]
|
||||
gt_polygons = [np.array(p, dtype=np.float64) for p in ann["segmentation"]]
|
||||
gt_bbox = BoxMode.convert(ann["bbox"], BoxMode.XYWH_ABS, BoxMode.XYXY_ABS)
|
||||
gt_bit_mask = polygons_to_bitmask(gt_polygons, height, width)
|
||||
|
||||
# Run rasterize ..
|
||||
torch_gt_bbox = torch.tensor(gt_bbox).to(dtype=torch.float32).reshape(-1, 4)
|
||||
box_bitmasks = {
|
||||
"polygon": PolygonMasks([gt_polygons]).crop_and_resize(torch_gt_bbox, mask_side_len)[0],
|
||||
"gridsample": rasterize_polygons_with_grid_sample(gt_bit_mask, gt_bbox, mask_side_len),
|
||||
"roialign": BitMasks(torch.from_numpy(gt_bit_mask[None, :, :])).crop_and_resize(
|
||||
torch_gt_bbox, mask_side_len
|
||||
)[0],
|
||||
}
|
||||
|
||||
# Run paste ..
|
||||
results = defaultdict(dict)
|
||||
for k, box_bitmask in box_bitmasks.items():
|
||||
padded_bitmask, scale = pad_masks(box_bitmask[None, :, :], 1)
|
||||
scaled_boxes = scale_boxes(torch_gt_bbox, scale)
|
||||
|
||||
r = results[k]
|
||||
r["old"] = paste_mask_in_image_old(
|
||||
padded_bitmask[0], scaled_boxes[0], height, width, threshold=0.5
|
||||
)
|
||||
r["aligned"] = paste_masks_in_image(
|
||||
box_bitmask[None, :, :], Boxes(torch_gt_bbox), (height, width)
|
||||
)[0]
|
||||
|
||||
table = []
|
||||
for rasterize_method, r in results.items():
|
||||
for paste_method, mask in r.items():
|
||||
mask = np.asarray(mask)
|
||||
iou = iou_between_full_image_bit_masks(gt_bit_mask.astype("uint8"), mask)
|
||||
table.append((rasterize_method, paste_method, iou))
|
||||
return table
|
||||
|
||||
def test_polygon_area(self):
|
||||
# Draw polygon boxes
|
||||
for d in [5.0, 10.0, 1000.0]:
|
||||
polygon = PolygonMasks([[[0, 0, 0, d, d, d, d, 0]]])
|
||||
area = polygon.area()[0]
|
||||
target = d ** 2
|
||||
self.assertEqual(area, target)
|
||||
|
||||
# Draw polygon triangles
|
||||
for d in [5.0, 10.0, 1000.0]:
|
||||
polygon = PolygonMasks([[[0, 0, 0, d, d, d]]])
|
||||
area = polygon.area()[0]
|
||||
target = d ** 2 / 2
|
||||
self.assertEqual(area, target)
|
||||
|
||||
|
||||
def benchmark_paste():
|
||||
S = 800
|
||||
H, W = image_shape = (S, S)
|
||||
N = 64
|
||||
torch.manual_seed(42)
|
||||
masks = torch.rand(N, 28, 28)
|
||||
|
||||
center = torch.rand(N, 2) * 600 + 100
|
||||
wh = torch.clamp(torch.randn(N, 2) * 40 + 200, min=50)
|
||||
x0y0 = torch.clamp(center - wh * 0.5, min=0.0)
|
||||
x1y1 = torch.clamp(center + wh * 0.5, max=S)
|
||||
boxes = Boxes(torch.cat([x0y0, x1y1], axis=1))
|
||||
|
||||
def func(device, n=3):
|
||||
m = masks.to(device=device)
|
||||
b = boxes.to(device=device)
|
||||
|
||||
def bench():
|
||||
for _ in range(n):
|
||||
paste_masks_in_image(m, b, image_shape)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return bench
|
||||
|
||||
specs = [{"device": torch.device("cpu"), "n": 3}]
|
||||
if torch.cuda.is_available():
|
||||
specs.append({"device": torch.device("cuda"), "n": 3})
|
||||
|
||||
benchmark(func, "paste_masks", specs, num_iters=10, warmup_iters=2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_paste()
|
||||
unittest.main()
|
@@ -0,0 +1,188 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||||
from __future__ import absolute_import, division, print_function, unicode_literals
|
||||
import numpy as np
|
||||
import unittest
|
||||
import torch
|
||||
from torchvision import ops
|
||||
|
||||
from detectron2.layers import batched_nms, batched_nms_rotated, nms_rotated
|
||||
|
||||
|
||||
def nms_edit_distance(keep1, keep2):
|
||||
"""
|
||||
Compare the "keep" result of two nms call.
|
||||
They are allowed to be different in terms of edit distance
|
||||
due to floating point precision issues, e.g.,
|
||||
if a box happen to have an IoU of 0.5 with another box,
|
||||
one implentation may choose to keep it while another may discard it.
|
||||
"""
|
||||
if torch.equal(keep1, keep2):
|
||||
# they should be equal most of the time
|
||||
return 0
|
||||
keep1, keep2 = tuple(keep1.cpu()), tuple(keep2.cpu())
|
||||
m, n = len(keep1), len(keep2)
|
||||
|
||||
# edit distance with DP
|
||||
f = [np.arange(n + 1), np.arange(n + 1)]
|
||||
for i in range(m):
|
||||
cur_row = i % 2
|
||||
other_row = (i + 1) % 2
|
||||
f[other_row][0] = i + 1
|
||||
for j in range(n):
|
||||
f[other_row][j + 1] = (
|
||||
f[cur_row][j]
|
||||
if keep1[i] == keep2[j]
|
||||
else min(min(f[cur_row][j], f[cur_row][j + 1]), f[other_row][j]) + 1
|
||||
)
|
||||
return f[m % 2][n]
|
||||
|
||||
|
||||
class TestNMSRotated(unittest.TestCase):
|
||||
def reference_horizontal_nms(self, boxes, scores, iou_threshold):
|
||||
"""
|
||||
Args:
|
||||
box_scores (N, 5): boxes in corner-form and probabilities.
|
||||
(Note here 5 == 4 + 1, i.e., 4-dim horizontal box + 1-dim prob)
|
||||
iou_threshold: intersection over union threshold.
|
||||
Returns:
|
||||
picked: a list of indexes of the kept boxes
|
||||
"""
|
||||
picked = []
|
||||
_, indexes = scores.sort(descending=True)
|
||||
while len(indexes) > 0:
|
||||
current = indexes[0]
|
||||
picked.append(current.item())
|
||||
if len(indexes) == 1:
|
||||
break
|
||||
current_box = boxes[current, :]
|
||||
indexes = indexes[1:]
|
||||
rest_boxes = boxes[indexes, :]
|
||||
iou = ops.box_iou(rest_boxes, current_box.unsqueeze(0)).squeeze(1)
|
||||
indexes = indexes[iou <= iou_threshold]
|
||||
|
||||
return torch.as_tensor(picked)
|
||||
|
||||
def _create_tensors(self, N):
|
||||
boxes = torch.rand(N, 4) * 100
|
||||
# Note: the implementation of this function in torchvision is:
|
||||
# boxes[:, 2:] += torch.rand(N, 2) * 100
|
||||
# but it does not guarantee non-negative widths/heights constraints:
|
||||
# boxes[:, 2] >= boxes[:, 0] and boxes[:, 3] >= boxes[:, 1]:
|
||||
boxes[:, 2:] += boxes[:, :2]
|
||||
scores = torch.rand(N)
|
||||
return boxes, scores
|
||||
|
||||
def test_batched_nms_rotated_0_degree_cpu(self):
|
||||
N = 2000
|
||||
num_classes = 50
|
||||
boxes, scores = self._create_tensors(N)
|
||||
idxs = torch.randint(0, num_classes, (N,))
|
||||
rotated_boxes = torch.zeros(N, 5)
|
||||
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
|
||||
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
|
||||
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
|
||||
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
|
||||
err_msg = "Rotated NMS with 0 degree is incompatible with horizontal NMS for IoU={}"
|
||||
for iou in [0.2, 0.5, 0.8]:
|
||||
backup = boxes.clone()
|
||||
keep_ref = batched_nms(boxes, scores, idxs, iou)
|
||||
assert torch.allclose(boxes, backup), "boxes modified by batched_nms"
|
||||
backup = rotated_boxes.clone()
|
||||
keep = batched_nms_rotated(rotated_boxes, scores, idxs, iou)
|
||||
assert torch.allclose(
|
||||
rotated_boxes, backup
|
||||
), "rotated_boxes modified by batched_nms_rotated"
|
||||
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
|
||||
def test_batched_nms_rotated_0_degree_cuda(self):
|
||||
N = 2000
|
||||
num_classes = 50
|
||||
boxes, scores = self._create_tensors(N)
|
||||
idxs = torch.randint(0, num_classes, (N,))
|
||||
rotated_boxes = torch.zeros(N, 5)
|
||||
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
|
||||
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
|
||||
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
|
||||
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
|
||||
err_msg = "Rotated NMS with 0 degree is incompatible with horizontal NMS for IoU={}"
|
||||
for iou in [0.2, 0.5, 0.8]:
|
||||
backup = boxes.clone()
|
||||
keep_ref = batched_nms(boxes.cuda(), scores.cuda(), idxs, iou)
|
||||
self.assertTrue(torch.allclose(boxes, backup), "boxes modified by batched_nms")
|
||||
backup = rotated_boxes.clone()
|
||||
keep = batched_nms_rotated(rotated_boxes.cuda(), scores.cuda(), idxs, iou)
|
||||
self.assertTrue(
|
||||
torch.allclose(rotated_boxes, backup),
|
||||
"rotated_boxes modified by batched_nms_rotated",
|
||||
)
|
||||
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
|
||||
|
||||
def test_nms_rotated_0_degree_cpu(self):
|
||||
N = 1000
|
||||
boxes, scores = self._create_tensors(N)
|
||||
rotated_boxes = torch.zeros(N, 5)
|
||||
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
|
||||
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
|
||||
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
|
||||
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
|
||||
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
|
||||
for iou in [0.5]:
|
||||
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
|
||||
keep = nms_rotated(rotated_boxes, scores, iou)
|
||||
self.assertLessEqual(nms_edit_distance(keep, keep_ref), 1, err_msg.format(iou))
|
||||
|
||||
def test_nms_rotated_90_degrees_cpu(self):
|
||||
N = 1000
|
||||
boxes, scores = self._create_tensors(N)
|
||||
rotated_boxes = torch.zeros(N, 5)
|
||||
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
|
||||
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
|
||||
# Note for rotated_boxes[:, 2] and rotated_boxes[:, 3]:
|
||||
# widths and heights are intentionally swapped here for 90 degrees case
|
||||
# so that the reference horizontal nms could be used
|
||||
rotated_boxes[:, 2] = boxes[:, 3] - boxes[:, 1]
|
||||
rotated_boxes[:, 3] = boxes[:, 2] - boxes[:, 0]
|
||||
|
||||
rotated_boxes[:, 4] = torch.ones(N) * 90
|
||||
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
|
||||
for iou in [0.2, 0.5, 0.8]:
|
||||
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
|
||||
keep = nms_rotated(rotated_boxes, scores, iou)
|
||||
assert torch.equal(keep, keep_ref), err_msg.format(iou)
|
||||
|
||||
def test_nms_rotated_180_degrees_cpu(self):
|
||||
N = 1000
|
||||
boxes, scores = self._create_tensors(N)
|
||||
rotated_boxes = torch.zeros(N, 5)
|
||||
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
|
||||
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
|
||||
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
|
||||
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
|
||||
rotated_boxes[:, 4] = torch.ones(N) * 180
|
||||
err_msg = "Rotated NMS incompatible between CPU and reference implementation for IoU={}"
|
||||
for iou in [0.2, 0.5, 0.8]:
|
||||
keep_ref = self.reference_horizontal_nms(boxes, scores, iou)
|
||||
keep = nms_rotated(rotated_boxes, scores, iou)
|
||||
assert torch.equal(keep, keep_ref), err_msg.format(iou)
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
|
||||
def test_nms_rotated_0_degree_cuda(self):
|
||||
N = 1000
|
||||
boxes, scores = self._create_tensors(N)
|
||||
rotated_boxes = torch.zeros(N, 5)
|
||||
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
|
||||
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
|
||||
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
|
||||
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
|
||||
err_msg = "Rotated NMS incompatible between CPU and CUDA for IoU={}"
|
||||
|
||||
for iou in [0.2, 0.5, 0.8]:
|
||||
r_cpu = nms_rotated(rotated_boxes, scores, iou)
|
||||
r_cuda = nms_rotated(rotated_boxes.cuda(), scores.cuda(), iou)
|
||||
|
||||
assert torch.equal(r_cpu, r_cuda.cpu()), err_msg.format(iou)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
@@ -0,0 +1,152 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||||
import numpy as np
|
||||
import unittest
|
||||
import cv2
|
||||
import torch
|
||||
from fvcore.common.benchmark import benchmark
|
||||
|
||||
from detectron2.layers.roi_align import ROIAlign
|
||||
|
||||
|
||||
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 _simple_roialign(self, img, box, resolution, aligned=True):
|
||||
"""
|
||||
RoiAlign with scale 1.0 and 0 sample ratio.
|
||||
"""
|
||||
if isinstance(resolution, int):
|
||||
resolution = (resolution, resolution)
|
||||
op = ROIAlign(resolution, 1.0, 0, 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 benchmark_roi_align():
|
||||
from detectron2 import _C
|
||||
|
||||
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(N, C, H, W, nboxes_per_img):
|
||||
input = torch.rand(N, C, H, W)
|
||||
boxes = []
|
||||
batch_idx = []
|
||||
for k in range(N):
|
||||
b = random_boxes([80, 80, 130, 130], 24, nboxes_per_img, H)
|
||||
# try smaller boxes:
|
||||
# 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.cuda()
|
||||
boxes = boxes.cuda()
|
||||
|
||||
def bench():
|
||||
_C.roi_align_forward(input, boxes, 1.0, 7, 7, 0, True)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return bench
|
||||
|
||||
args = [dict(N=2, C=512, H=256, W=256, nboxes_per_img=500)]
|
||||
benchmark(func, "cuda_roialign", args, num_iters=20, warmup_iters=1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if torch.cuda.is_available():
|
||||
benchmark_roi_align()
|
||||
unittest.main()
|
@@ -0,0 +1,176 @@
|
||||
# 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()
|
Reference in New Issue
Block a user