preprocess/detectron2/tests/layers/test_roi_align.py

# 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()