Reputation: 322
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:
Upvotes: 1
Views: 6548
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.
Upvotes: 6