Search code examples
pythontensorflowobject-detectionintersectionbounding-box

to find the intersection of two bounding box in tensorflow?


The coordinate of the system is

boundary coordinates (x_min, y_min, x_max, y_max).

and I want to find the intersection of two boxes set1 and set2

set1 -> (n1,4)
set2 -> (n2,4)

example 

set_1-> tensor([[0.2400, 0.2342, 0.8500, 0.8048],
        [0.1420, 0.5075, 0.2440, 0.5856],
        [0.0000, 0.5075, 0.1420, 0.5976]], device='cuda:0')
set_2-> tensor([[-0.0368, -0.0368,  0.0632,  0.0632],
        [-0.0576, -0.0576,  0.0839,  0.0839],
        [-0.0576, -0.0222,  0.0839,  0.0485],
        ...,
        [ 0.0000,  0.0000,  1.0000,  1.0000],
        [ 0.0000,  0.1818,  1.0000,  0.8182],
        [ 0.1818,  0.0000,  0.8182,  1.0000]], device='cuda:0') torch.Size([8732, 4])

if I want to get the intersection of two sets of bounding box

:return: intersection of each of the boxes in set 1 with respect to each of the boxes in set 2, a tensor of dimensions (n1, n2)

How should I write this in python and tensorflow?

def find_intersection(set_1, set_2):
    """
    Find the intersection of every box combination between two sets of boxes that are in boundary coordinates.

    :param set_1: set 1, a tensor of dimensions (n1, 4)
    :param set_2: set 2, a tensor of dimensions (n2, 4)
    :return: intersection of each of the boxes in set 1 with respect to each of the boxes in set 2, a tensor of dimensions (n1, n2)
    """
    print('set1->',set_1)
    print('set_2->',set_2)
    return tf.sets.intersection(set_1,set_2)
def find_jaccard_overlap(set_1, set_2):
    """
    Find the Jaccard Overlap (IoU) of every box combination between two sets of boxes that are in boundary coordinates.

    :param set_1: set 1, a tensor of dimensions (n1, 4)
    :param set_2: set 2, a tensor of dimensions (n2, 4)
    :return: Jaccard Overlap of each of the boxes in set 1 with respect to each of the boxes in set 2, a tensor of dimensions (n1, n2)
    """

    # Find intersections
    intersection = find_intersection(set_1, set_2)  # (n1, n2)
    print('intersection->', intersection)

    # Find areas of each box in both sets
    areas_set_1 = (set_1[:, 2] - set_1[:, 0]) * (set_1[:, 3] - set_1[:, 1])  # (n1)
    areas_set_2 = (set_2[:, 2] - set_2[:, 0]) * (set_2[:, 3] - set_2[:, 1])  # (n2)

    # Find the union
    # PyTorch auto-broadcasts singleton dimensions
    union = areas_set_1.unsqueeze(1) + areas_set_2.unsqueeze(0) - intersection  # (n1, n2)

    return intersection / union  # (n1, n2)

Solution

  • You can do that like this:

    import tensorflow as tf
    
    def box_intersections(set1, set2):
        set1 = tf.expand_dims(set1, axis=-1)
        x_min = tf.math.maximum(set1[:, 0], set2[:, 0])
        y_min = tf.math.maximum(set1[:, 1], set2[:, 1])
        x_max = tf.math.minimum(set1[:, 2], set2[:, 2])
        y_max = tf.math.minimum(set1[:, 3], set2[:, 3])
        dx = tf.math.maximum(x_max - x_min, 0)
        dy = tf.math.maximum(y_max - y_min, 0)
        return dx * dy
    
    # Example
    tf.random.set_seed(0)
    # Make random boxes
    n1 = 10
    n2 = 20
    set1 = tf.random.uniform((n1, 2))
    set1 = tf.concat([set1, set1 + tf.random.uniform((n1, 2))], axis=1)
    set2 = tf.random.uniform((n2, 2))
    set2 = tf.concat([set2, set2 + tf.random.uniform((n2, 2))], axis=1)
    # Compute intersections
    intersect = box_intersections(set1, set2)
    print(intersect.shape)
    # (10, 20)