Search code examples
pythonnumpyopencvobject-detection

Generating bounding boxes from heatmap data

I have the heatmap data for a vehicle detection project I'm working on but I'm at a loss for where to go next. I want to draw a generalized bounding box around the 'hottest' portions of the image. My first thought was to just draw a box over all portions that overlap but something's telling me there is a more accurate way to do this. Any help would be appreciated! Unfortunately my reputation prevents me from posting images. Here's how I'm creating the heatmap:

# Positive prediction window coordinate structure: ((x1, y1), (x2, y2))
def create_heatmap(bounding_boxes_list):
     # Create a black image the same size as the input data
     heatmap = np.zeros(shape=(375, 1242))
     # Traverse the list of bounding box locations in test image
     for bounding_box in bounding_boxes_list:
          heatmap[bounding_box[0][1]:bounding_box[1][1], bounding_box[0][ 0]:bounding_box[1][0]] += 1

return heatmap

Here's the link to the heatmap I have

Here's a general idea of what I had in mind

Solution

  • Otsu's threshold and contour detection on the binary image should do it. Using this screenshotted image without the axis lines:

    enter image description here

    enter image description here

    import cv2

# Grayscale then Otsu's threshold
image = cv2.imread('1.png')
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]

# Find contours
cnts = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = cnts[0] if len(cnts) == 2 else cnts[1]
for c in cnts:
    x,y,w,h = cv2.boundingRect(c)
    cv2.rectangle(image, (x, y), (x + w, y + h), (36,255,12), 2)

cv2.imshow('thresh', thresh)
cv2.imshow('image', image)
cv2.waitKey()