Avoid unintentional "merging" in cluttered object detection

Question

I have a problem that has bothered me quite some time.

With modern methods object detectors can often be accurately trained, even with small to medium sized datasets. However, there is one thing where I always see errors and that is when objects are very close together. The problem seems to be that the commonly used clustering methods like NMS will often combine predictions in a wrong way, like in this example:

So far, I found that one rather hacky way to solve this is to only annotate some part of the object as the parts won't be close together, but I wonder if there could be a nicer way to solve this. Many solutions like bottom-up keypoint detection look promising, but they seem to only show their strength with large training datasets - with less data I saw a lot of confusion between the objects.

My question - has anybody experienced something similar and maybe has a promising paper or implementation for a solution that could work without huge datasets?

Answer 1

I might have a solution, but based on a hard assumption, which is the following:

all expected predicted boxes share a characteristic, in this case they are all rectangles with big height/length ratio

This is important cause the fact that you have a big box in your final prediction means that that box had a confidence score high enough to eat out potential smaller boxes, which we know are correct but that the model failed to predict with high confidence. So what need to be done is helping the model by leveraging some a priori knowledge to filter out completely "bad boxes", before proceeding applying MNS.

Below you can see the results of a toy script, in which I first generate some random boxe (top left plot), sometimes "good", which mean the long side is more than 4 times the short side, and sometimes "bad", which means the box is more squarish. When applying NMS to the boxes (top right plot), you can see that some of the bad squarish boxes are not discarded in the predictions. We know though that good predictions should not be squarish, hence we can apply a simple filter to the original boxes before applying NMS. In this case the filter could be:

$long side / short side > 4$

If we filter the boxes (bottom left plot), we remain with a set of potential predictions made only of good boxes, and if we apply MNS to this new filtered set (bottom right plot), areas that before were labeled with bad boxes are now getting labelled with good boxes, see for example the green rectangle on the left at y=400, which was eaten out by a large box in the original set.

of course this is a and crafted solution specific for this case, it might not work perfectly, but I feel it's easier to solve your issue following this post processing path rather than redesigning NMS or by different model training regimes (since it seems that on overall your model is not performing that bad).

import numpy as np
from numpy.random import randint as rind
import matplotlib.pyplot as plt

def nms(dets, thresh=0.1):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]
    return keep

def nms_filer(dets):
    good_boxes = (dets[:, 2] - dets[:, 0]) / (dets[:, 3] - dets[:, 1]) > 4
    print(f"Dropped {dets.shape[0] - len(good_boxes)} boxes")
    dets = dets[good_boxes]
    return dets

def random_boxes(tot=100, min_=0, max_=1000):
    boxes = np.zeros((tot, 5))
    for t in range(tot):
        if np.random.rand() < 0.6:
            min_lenght = 250
            min_height = 30
        else:
            min_lenght = 250
            min_height = 250
        x1 = rind(min_, max_ - min_lenght)
        x2 = rind(x1 + min_lenght, min(max_, x1 + min_lenght + rind(10,min_lenght/2)))
        y1 = rind(min_, max_ - min_height)
        y2 = rind(y1 + min_height, min(max_, y1 + min_height + rind(10,min_height/2)))
        score = 0.5 + np.random.rand() / 2
        boxes[t, :] = np.asarray([x1, y1, x2, y2, score])
    return boxes

def plot_boxes(boxes):
    f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2)
    
    for b in range(boxes.shape[0]):
        x1, y1, x2, y2, score = boxes[b]
        x = [x1,x2,x2,x1,x1]
        y = [y1,y1,y2,y2,y1]
        ax1.plot(x, y, linewidth=2)
        ax1.set_title("Original boxes")

    indices_keep = nms(boxes)
    boxes_nms_original = boxes[indices_keep]
    for b in range(boxes_nms_original.shape[0]):
        x1, y1, x2, y2, score = boxes_nms_original[b]
        x = [x1,x2,x2,x1,x1]
        y = [y1,y1,y2,y2,y1]
        ax2.plot(x, y, linewidth=2)
        ax2.set_title("NMS original boxes")

    boxes_filtered = nms_filer(boxes)
    for b in range(boxes_filtered.shape[0]):
        x1, y1, x2, y2, score = boxes_filtered[b]
        x = [x1,x2,x2,x1,x1]
        y = [y1,y1,y2,y2,y1]
        ax3.plot(x, y, linewidth=2)
        ax3.set_title("Filtered boxes")

    indices_keep = nms(boxes_filtered)
    boxes_nms_constrain = boxes_filtered[indices_keep]
    for b in range(boxes_nms_constrain.shape[0]):
        x1, y1, x2, y2, score = boxes_nms_constrain[b]
        x = [x1,x2,x2,x1,x1]
        y = [y1,y1,y2,y2,y1]
        ax4.plot(x, y, linewidth=2)
        ax4.set_title("NMS filtered boxes")

    plt.show()


if __name__=="__main__":

    boxes = random_boxes()
    plot_boxes(boxes)