Intersection over Union for rotated rectangles

Question

Problem Statement

I am trying to find the intersection over union (IoU) metric for one to several rotated rectangles compared to many different rotated rectangles. Here are some images to help visualize this metric:

The second metric is quite close to the scenario I'm trying to calculate, the white area divided by the combined black and white area.

Example

For each small box on the left we need to determine the IoU metric for each red box on the right. The output in this case will be an array of size (5, 756) since there are 5 boxes on the left with 756 IoU metrics for each box on the right.

My solution

To solve this problem, I fill in each "anchor" box (from the above right picture) individually and store that in an array. This is a binary array.

I then take these filled in anchor boxs (now separated so each box is in their own image as above) and find the IoU metric by simple multiplication and summation. This approach will calculate the IoU metric regardless of the shape of the objects passed in. However, it is extremely inefficient both in memory and computationally. I am using the pytorch library to load the arrays onto my GPU to parallelize the computations.

Code

This code runs in python-3.x.

import numpy as np
import torch
import cv2
from timeit import Timer

def jaccard(box_a, box_b):
    denorm_bbox = torch.cat([box_a[:, :4]*768, box_a[:, 4].unsqueeze(1)*90], dim=1)
    b_boxes = list(map(lambda x:np.ceil(cv2.boxPoints(((x[1], x[0]), (x[2], x[3]), x[4]))), denorm_bbox.numpy()))
    b_imgs = torch.from_numpy(np.array([cv2.fillConvexPoly(np.zeros((768,768), dtype=np.uint8), np.int0(b), 1) for b in b_boxes]).astype(float)).float()
    intersection = torch.FloatTensor()
    summation = torch.FloatTensor()
    for b_img in b_imgs:
        intersection = torch.cat([intersection, (b_img*box_b).sum((1,2)).unsqueeze(0)])
        summation = torch.cat([summation, (b_img+box_b).sum((1,2)).unsqueeze(0)])
    return intersection / (summation - intersection + 1.0)


def main():
    anc_grids = [3,6,12]
    anc_zooms = [0.7]
    anc_ratios = [(1.,1)]
    anc_angles = np.array(range(-90, 90, 45))/90
    anchor_scales = [(anz*i,anz*j) for anz in anc_zooms for (i,j) in anc_ratios]
    k = len(anchor_scales) * len(anc_angles) # number of anchor boxes per anchor point
    anc_offsets = [1/(o*2) for o in anc_grids]

    anc_x = np.concatenate([np.repeat(np.linspace(ao, 1-ao, ag), ag) for ao,ag in zip(anc_offsets,anc_grids)])
    anc_y = np.concatenate([np.tile(np.linspace(ao, 1-ao, ag), ag) for ao,ag in zip(anc_offsets,anc_grids)])
    anc_ctrs = np.repeat(np.stack([anc_x,anc_y], axis=1), k, axis=0)
    anc_sizes = np.tile(np.concatenate([np.array([[o/ag,p/ag] for i in range(ag*ag) for o,p in anchor_scales])
                for ag in anc_grids]), (len(anc_angles),1))
    grid_sizes = torch.from_numpy(np.concatenate([np.array([ 1/ag       for i in range(ag*ag) for o,p in anchor_scales])
                   for ag in anc_grids for aa in anc_angles])).unsqueeze(1)
    anc_rots = np.tile(np.repeat(anc_angles, len(anchor_scales)), sum(i*i for i in anc_grids))[:,np.newaxis]
    anchors = torch.from_numpy(np.concatenate([anc_ctrs, anc_sizes, anc_rots], axis=1)).float()

    denorm_anchors = torch.cat([anchors[:, :4]*768, anchors[:, 4].unsqueeze(1)*90], dim=1)
    np_anchors = denorm_anchors.numpy()
    iou_anchors = list(map(lambda x:np.ceil(cv2.boxPoints(((x[1], x[0]), (x[2], x[3]), x[4]))), np_anchors))
    anchor_imgs = torch.from_numpy(np.array([cv2.fillConvexPoly(np.zeros((768,768), dtype=np.uint8), np.int0(a), 1) for a in iou_anchors]).astype(float)).float()

    test_tensor = torch.Tensor([[ 0.0807,  0.2844,  0.0174,  0.0117, -0.8440],
                                [ 0.3276,  0.0358,  0.0169,  0.0212, -0.1257],
                                [ 0.3040,  0.2904,  0.0101,  0.0157, -0.5000],
                                [ 0.0065,  0.2109,  0.0130,  0.0078, -1.0000],
                                [ 0.1895,  0.1556,  0.0143,  0.0091, -1.0000]])

    t = Timer(lambda: jaccard(test_tensor, anchor_imgs))
    print(f'Consuming {(len(anchors) * np.dtype(float).itemsize * 768 * 768)/1000000000} Gb on {"GPU" if anchors.is_cuda else "RAM"}')
    print(f'Averaging {t.timeit(number=100)/100} seconds per function call')
    print(jaccard(test_tensor, anchor_imgs))

if __name__ == '__main__':
    main()

Sample Run:

Consuming 3.567255552 Gb on RAM
Averaging 3.107201199789997 seconds per function call
tensor([[0.0020, 0.0000, 0.0020,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0026, 0.0026, 0.0026,  ..., 0.0000, 0.0000, 0.0000]])

The results of this run were on a i7-8700K for reference.

---

What I would like reviewed:

• Memory Consumption: Right now I am consuming a lot of memory for all the anchor boxes I compare each box to. What are some ways I can dramatically reduce this without giving up the accuracy of my results?

• Speed: I need this to run fast!

Answer 1

1. Review

The large bounty suggests that getting a review of this code is important. But there are some other things you could have done to make reviewing this code more productive:

1. Give us some more context. What is the purpose of this code? What do these rotated rectangles represent?

2. Tell us how the data computed in main relates to your problem. Is this the actual data you are using in your application? Or is it example/test data?

3. Write docstrings. What does jaccard do? What arguments does it take? What does it return? Both arguments seem to be arrays of rotated rectangles, but how are they specified?

4. box_a seems to be an array of 5-tuples (centre x, centre y, width, height, rotation in quarter-turns), while box_b seems to be an array of 768×768 rasterizations. Giving different kinds of data similar names is asking for confusion: it would be clearer to use a name like box_images_b. Or even better, take the two arrays of rectangles in the same format and do the rasterization for both in the function, instead of doing one in the caller and one in the function.

5. It is unusual to represent angles as anything other than degrees or radians, so the representation as quarter-turns requires an explanation.

6. Write comments. What do the values mean in the variables anc_grids, anc_zooms etc.?

7. Remove unused variables like grid_sizes.

8. Follow the advice in PEP 8 and keep the lines to 79 characters or fewer, so that we don't have to scroll the code horizontally to read it.

You've written more than a hundred answers on Code Review, so it seems to me that you ought to be aware of all this advice already! Documentation, clean code, and consistency is good for everyone, not just for beginners.

2. Use shapely for fast geometric operations

Finding the intersection-over-union for a pair of polygons is trivial if you use the shapely geometry package. For example, suppose we have two polygons:

from shapely.geometry import Polygon

a = Polygon([(1, 1), (5, 1), (5, 4), (1, 4)])
b = Polygon([(2, 2), (6, 3), (5, 7), (1, 6)])

Here a is shown as a green rectangle in the figure below, and b as a red square. The intersection over union can be computed like this:

>>> a.intersection(b).area / a.union(b).area
0.2275132275132275

This takes about 300 µs on my laptop:

>>> from timeit import timeit
>>> timeit(lambda:a.intersection(b).area / a.union(b).area, number=1000) / 1000
0.0003085929369844962

The geometric approach produces more accurate results than the rasterization approach in the post. For example, if we take the first pair of rectangles, test_tensor[0] and anchors[0], the sample run in the post shows that an intersection-over-union of 0.0020 was computed. But the correct result is 0.0016 to 2 significant digits.

3. Use rtree for finding candidate intersections

Efficiently find objects that might intersect an axis-aligned rectangle by putting the axis-aligned bounding boxes for the objects into a R-tree (for example, using the rtree package) and then query the rectangle of interest (for example, using the intersection method).

4. Example implementation

import rtree.index
from shapely.affinity import rotate, translate
from shapely.geometry import Polygon


def rect_polygon(x, y, width, height, angle):
    """Return a shapely Polygon describing the rectangle with centre at
    (x, y) and the given width and height, rotated by angle quarter-turns.

    """
    w = width / 2
    h = height / 2
    p = Polygon([(-w, -h), (w, -h), (w, h), (-w, h)])
    return translate(rotate(p, angle * 90), x, y)


def intersection_over_union(rects_a, rects_b):
    """Calculate the intersection-over-union for every pair of rectangles
    in the two arrays.

    Arguments:
    rects_a: array_like, shape=(M, 5)
    rects_b: array_like, shape=(N, 5)
        Rotated rectangles, represented as (centre x, centre y, width,
        height, rotation in quarter-turns).

    Returns:
    iou: array, shape=(M, N)
        Array whose element i, j is the intersection-over-union
        measure for rects_a[i] and rects_b[j].

    """
    m = len(rects_a)
    n = len(rects_b)
    if m > n:
        # More memory-efficient to compute it the other way round and
        # transpose.
        return intersection_over_union(rects_b, rects_a).T

    # Convert rects_a to shapely Polygon objects.
    polys_a = [rect_polygon(*r) for r in rects_a]

    # Build a spatial index for rects_a.
    index_a = rtree.index.Index()
    for i, a in enumerate(polys_a):
        index_a.insert(i, a.bounds)

    # Find candidate intersections using the spatial index.
    iou = np.zeros((m, n))
    for j, rect_b in enumerate(rects_b):
        b = rect_polygon(*rect_b)
        for i in index_a.intersection(b.bounds):
            a = polys_a[i]
            intersection_area = a.intersection(b).area
            if intersection_area:
                iou[i, j] = intersection_area / a.union(b).area

    return iou

Using the test data from the post, I find that computing the 5×756 array of results takes less than a quarter of a second:

>>> len(TEST_RECTS_A), len(TEST_RECTS_B)
(5, 756)
>>> intersection_over_union(TEST_RECTS_A, TEST_RECTS_B).shape
(5, 756)
>>> timeit(lambda:intersection_over_union(TEST_RECTS_A, TEST_RECTS_B), number=1)
0.22243631538003683

The memory used by the R-tree will be \$O(m)\$, so the overall memory usage of this approach will be dominated by the size of the result, which is \$Θ(mn)\$. You might be able to avoid building this array of results if you were willing to process the results one at a time, but that depends on how you are using this array, which you didn't tell us (see §1.1 above).

5. Further optimizations

If you're looking at further optimizations, the important thing to note is that the R-tree very efficiently disposes of the vast majority of possible intersections. With the example data from the post, there are 5×756 = 3780 pairs of rectangles, but the R-tree produces only 74 candidate pairs, of which 52 actually intersect.

So if this behaviour is representative of the real problem space, then there's little to gain by clever optimization of the intersection-over-union calculation. Profiling shows that (roughly) 66% of runtime is spent in rect_polygon (mostly inside translate and rotate), 12% in a.bounds and b.bounds, but only 3% in a.intersection(b).area and 2% in a.union(b).area.

So the first place to look for further speedups is in the affine transformation of the coordinates of the rotated rectangles. This would probably go faster if vectorized using NumPy.

Answer 2

There are two main reasons why your program is slow and using huge amounts of memory:

1. You are using a 768×768 pixel image for each overlap calculation

2. You are checking each anchor box against each red box, but most of the time there is no overlap at all.

Use exact solutions where possible

For the first issue, instead of doing image manipulation, you can instead do an exact calculation by finding the intersections of the edges of the rotated rectangles, constructing the polygon that is the overlapping area, and then calculating the area of that polygon. This is much faster and uses much less memory. See this question for a possible approach in Python.

To calculate the area of the overlapping region with such a method, you would use on the order of 100 to 200 floating point operations: checking all 4 edges of one rectangle against each 4 of the other rectangle is about 4 * 4 * 8 = 128 operations, and calculating the area is around 8 * 4 = 32 operations for the worst case polygon with 8 edges. You probably only need to store the coordinates of the intersection points, which is 8 * 2 = 16 floating point values, which is 128 bytes when using double precision. I'm assuming all temporary values from the computations can be stored in registers. Compare that to having to create a 768*768 pixel image, where each image should hold at least 3 possible value: not part of a rectangle, non-overlapping area, and overlapping area. Because packing and unpacking 2 bit values is not efficient, you probably end up using 1 byte per pixel, so that is 576 kiB of memory, 4608 times as much as the 128 bytes used for the exact method. You would have to intialize this image, then calculate draw the boxes, which can probably done with integers or fixed-point mathematics, which may or may not be any faster than floating point arithmetic. Assuming you only need to do one operation per pixel, that will still be 589824 operations, or around 3000 more than using the exact method.

Avoid testing all possible combinations

For the second issue there are various ways to deal with it. A simple approach is to calculate the bounding box for each rotated rectangle, and first check if the bounding boxes overlap before doing the more expensive IoU calculation. However, that still requires a lot of tests. Looking at the images you provided, most red rectangles are quite small relative to the whole domain. In that case, a possible approach might be to divide the domain into smaller tiles, and for each tile have a list of rectangles that are overlapping that tile. Then you check each anchor box against all of the red rectangles that are overlapping the same tiles as the anchor box itself is overlapping.

For example, if you use tiles of size 100×100, then most boxes will overlap roughly only 2 tiles. That means, with 756 boxes in a 768x768 domain, there are approximately 2 * 756 * 100²/768² ≈ 25 overlapping each tile. Each anchor box will overlap approximately 2 tiles, so you need to check only 50 red boxes for each anchor box, reducing the amount of work by a factor 15 in this case. You can play around with the tile size to find the optimum size.

Use a faster language

Last but not least, Python will only get you so far. Once you have reduced the algorithmic complexity of your solution as much as possible, you are left with the overhead of an interpreted language like Python. To get the best performance, use a language like C, C++ or Rust. Depending on how much your Python code already off-loaded most of its work to libraries written in those compiled languages, you might get a huge speed boost this way.

Total speedup and memory usage reduction

It is very hard to give an exact number of how much you can improve this code, unless you actually implement the suggested methods and measure it. I'll give an estimate, and it might easily be off by a factor of 10.

Your code reports using ~3.5 GB of RAM, and takes around 3 seconds. That's actually close to what I would expect: 5 * 756 * 768 * 768 * 1 byte/pixel ≈ 2.2 GB, and also 5 * 756 * 768 * 768 * 1 FLOP/pixel / 1 GFLOP/s ≈ 2.2 seconds.

If I put in my estimates for the time and memory usage, and scale them with the ratio of the actual measurements from your code divided by the above estimate, then with both exact solutions and tiles combined, I get a memory usage of 5 * 50 * 128 / (5 * 756 * 768 * 768) * 3.57 GB ≈ 51 kB, and 5 * 50 * 200 / (5 * 756 * 768 * 768) * 3.11 s ≈ 70 µs. Now, those are some very small numbers, and while I think they are realistic, they don't include the setup costs for the tiles, so an actual program will probably use a bit more memory and time than that.