Silent-J
Silent-J

Reputation: 322

Multiple bounding-boxes with cv2.rectangle()

I currently have coordinates data for bounding boxes, contained in a nested data structure like so:

 defaultdict(list,
            {'giraffe': [{'conf': 0.9869,
               'coords': {'center_x': 0.360333,
                'center_y': 0.532274,
                'height': 0.596343,
                'width': 0.144651}},
              {'conf': 0.253321,
               'coords': {'center_x': 0.016296,
                'center_y': 0.565007,
                'height': 0.580526,
                'width': 0.03498}}],
             'zebra': [{'conf': 0.998863,
               'coords': {'center_x': 0.545974,
                'center_y': 0.693267,
                'height': 0.301859,
                'width': 0.257102}}]})

I want to iterate through the data structure (img_obj_data) and draw the rectangles for each object per object_class.

I want to then save the image (with boxes drawn), so I can open it later.

My first attempt is below:

import cv2

img = cv2.imread(img_path)
img_h, img_w = img.shape[:2]

for obj_class in img_obj_data.keys():
    for sub_dict in img_obj_data[obj_class]:
        x, y, w, h = sub_dict['coords'].values()
        
        # coords cannot be floats
        x = int(x*img_w)
        y = int(y*img_h)
        x_max = int(w*img_w)
        y_max = int(y*img_h)

        cv2.rectangle(img, (x, y), (x_max, y_max), color=(0, 255, 0), thickness=2)
cv2.imwrite('/content/foobar.jpg', img)

I now have run into 2 problems:

Problem 1) The bounding boxes are not properly aligned with the objects and are cropped out of the image. (The coordinates are originally floats, I multiplied them by the image width and height but I have a hunch that this was the incorrect thing to do?)

Problem 2)

My code at the moment colours all boxes the same. How can I go about this so that the colours differ per object class?

Here's the image in question, and the resulting image from my code: enter image description here

enter image description here

Upvotes: 1

Views: 6548

Answers (1)

pond27
pond27

Reputation: 278

One of the examples is as follows. *If necessary, please switch h and w for your environment.

import cv2
from random import randint

img = cv2.imread(img_path)
img_h, img_w = img.shape[:2]
colors = {}

for obj_class in img_obj_data.keys():
    if obj_class not in colors:
        colors[obj_class] = [randint(0, 255), randint(0, 255), randint(0, 255)]

    for sub_dict in img_obj_data[obj_class]:
        x, y, h, w = sub_dict['coords'].values()
        # coords cannot be floats
        x_min = int((x-w/2)*img_w)
        y_min = int((y-h/2)*img_h)
        x_max = int((x+w/2)*img_w)
        y_max = int((y+h/2)*img_h)
        cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=colors[obj_class], thickness=2)
cv2.imwrite('/content/foobar.jpg', img)

This example changes colors randomly by class name but if you know class names, you can define colors in advance.

enter image description here

Upvotes: 6

Related Questions