Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Labelling issue and Bounding box drawing issue. #8

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
272 changes: 182 additions & 90 deletions bbox_visualizer/bbox_visualizer.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,44 @@
import cv2
import logging


def draw_rectangle(img,
bbox,
bbox_color=(255, 255, 255),
thickness=3,
is_opaque=False,
alpha=0.5):
def check_and_modify_bbox(bbox, img_size, margin=0):
"""
Checks if the bounding box is inside the given image.
If not, the coordinates are trimmed inorder to fit inside the image.

Trimming criteria:
- for xmin and ymin:
- xmin and ymin are trimmed to 0, if their value is negative i.e. out of the
scope of image.
- for xmax and ymax:
- xmax and ymax are trimmed to image_width and image_height respectivel


Parameters
----------
bbox : list
a list containing x_min, y_min, x_max and y_max of the rectangle positions
img_size: tuple
a tuple containing image width, image height and color channel.
margin: integer (default = 0)
space between edge of the image and bounding box incase the box is trimmed such that
it's size is exactly as the image.

Returns
-------
list
bounding box coordinates (xmin, ymin, xmax, ymax)
"""
bbox = [value if value > 0 else margin for value in bbox]
bbox[2] = bbox[2] if bbox[2] < img_size[1] else img_size[1] - margin
bbox[3] = bbox[3] if bbox[3] < img_size[0] else img_size[0] - margin
return bbox


def draw_rectangle(
img, bbox, bbox_color=(255, 255, 255), thickness=3, is_opaque=False, alpha=0.5
):
"""Draws the rectangle around the object

Parameters
Expand All @@ -29,29 +61,33 @@ def draw_rectangle(img,
ndarray
the image with the bounding box drawn
"""

bbox = check_and_modify_bbox(bbox, img.shape)
output = img.copy()
if not is_opaque:
cv2.rectangle(output, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
bbox_color, thickness)
cv2.rectangle(
output, (bbox[0], bbox[1]), (bbox[2], bbox[3]), bbox_color, thickness
)
else:
overlay = img.copy()

cv2.rectangle(overlay, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
bbox_color, -1)
cv2.rectangle(overlay, (bbox[0], bbox[1]), (bbox[2], bbox[3]), bbox_color, -1)
cv2.addWeighted(overlay, alpha, output, 1 - alpha, 0, output)

return output


def add_label(img,
label,
bbox,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
top=True):
"""adds label, inside or outside the rectangle
def add_label(
img,
label,
bbox,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
top=True,
):
"""
adds label, inside or outside the rectangle.
if label cannot be drawn on outside of the box, it puts label inside the box.

Parameters
----------
Expand All @@ -75,34 +111,55 @@ def add_label(img,
ndarray
the image with the label written
"""

text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][0]
text_height = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][1]

if top:
if top and bbox[1] - text_height > text_height:
label_bg = [bbox[0], bbox[1], bbox[0] + text_width, bbox[1] - 30]
if draw_bg:
cv2.rectangle(img, (label_bg[0], label_bg[1]),
(label_bg[2] + 5, label_bg[3]), text_bg_color, -1)
cv2.putText(img, label, (bbox[0] + 5, bbox[1] - 5),
cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
cv2.rectangle(
img,
(label_bg[0], label_bg[1]),
(label_bg[2] + 5, label_bg[3]),
text_bg_color,
-1,
)
cv2.putText(
img,
label,
(bbox[0] + 5, bbox[1] - 5),
cv2.FONT_HERSHEY_SIMPLEX,
1,
text_color,
2,
)

else:
label_bg = [bbox[0], bbox[1], bbox[0] + text_width, bbox[1] + 30]
if draw_bg:
cv2.rectangle(img, (label_bg[0], label_bg[1]),
(label_bg[2] + 5, label_bg[3]), text_bg_color, -1)
cv2.putText(img, label, (bbox[0] + 5, bbox[1] - 5 + 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
cv2.rectangle(
img,
(label_bg[0], label_bg[1]),
(label_bg[2] + 5, label_bg[3]),
text_bg_color,
-1,
)
cv2.putText(
img,
label,
(bbox[0] + 5, bbox[1] - 5 + 30),
cv2.FONT_HERSHEY_SIMPLEX,
1,
text_color,
2,
)

return img


def add_T_label(img,
label,
bbox,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0)):
def add_T_label(
img, label, bbox, draw_bg=True, text_bg_color=(255, 255, 255), text_color=(0, 0, 0)
):
"""adds a T label to the rectangle, originating from the top of the rectangle

Parameters
Expand Down Expand Up @@ -131,30 +188,45 @@ def add_T_label(img,

# draw vertical line
x_center = (bbox[0] + bbox[2]) // 2
y_top = bbox[1] - 50
cv2.line(img, (x_center, bbox[1]), (x_center, y_top), text_bg_color, 3)
line_top = y_top = bbox[1] - 50

# draw rectangle with label
y_bottom = y_top
y_top = y_bottom - text_height - 5

if y_top < 0:
logging.warning(
"Labelling style 'T' going out of frame. Falling back to normal labeling."
)
return add_label(img, label, bbox)

cv2.line(img, (x_center, bbox[1]), (x_center, line_top), text_bg_color, 3)
x_left = x_center - (text_width // 2) - 5
x_right = x_center + (text_width // 2) + 5
if draw_bg:
cv2.rectangle(img, (x_left, y_top - 3), (x_right, y_bottom),
text_bg_color, -1)
cv2.putText(img, label, (x_left + 5, y_bottom - 7),
cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
cv2.rectangle(img, (x_left, y_top - 3), (x_right, y_bottom), text_bg_color, -1)
cv2.putText(
img,
label,
(x_left + 5, y_bottom - 7),
cv2.FONT_HERSHEY_SIMPLEX,
1,
text_color,
2,
)

return img


def draw_flag_with_label(img,
label,
bbox,
write_label=True,
line_color=(255, 255, 255),
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0)):
def draw_flag_with_label(
img,
label,
bbox,
write_label=True,
line_color=(255, 255, 255),
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
):
"""draws a pole from the middle of the object that is to be labeled and adds the label to the flag

Parameters
Expand Down Expand Up @@ -183,8 +255,14 @@ def draw_flag_with_label(img,
# draw vertical line

x_center = (bbox[0] + bbox[2]) // 2
y_bottom = int((bbox[1] * .75 + bbox[3] * .25))
y_bottom = int((bbox[1] * 0.75 + bbox[3] * 0.25))
y_top = bbox[1] - (y_bottom - bbox[1])
if y_top < 0:
logging.warning(
"Labelling style 'Flag' going out of frame. Falling back to normal labeling."
)
img = draw_rectangle(img, bbox, bbox_color=line_color)
return add_label(img, label, bbox)

start_point = (x_center, y_top)
end_point = (x_center, y_bottom)
Expand All @@ -194,16 +272,29 @@ def draw_flag_with_label(img,
# write label

if write_label:
text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1,
2)[0][0]
text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][0]
label_bg = [
start_point[0], start_point[1], start_point[0] + text_width,
start_point[1] + 30
start_point[0],
start_point[1],
start_point[0] + text_width,
start_point[1] + 30,
]
cv2.rectangle(img, (label_bg[0], label_bg[1]),
(label_bg[2] + 5, label_bg[3]), text_bg_color, -1)
cv2.putText(img, label, (start_point[0] + 5, start_point[1] - 5 + 30),
cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
cv2.rectangle(
img,
(label_bg[0], label_bg[1]),
(label_bg[2] + 5, label_bg[3]),
text_bg_color,
-1,
)
cv2.putText(
img,
label,
(start_point[0] + 5, start_point[1] - 5 + 30),
cv2.FONT_HERSHEY_SIMPLEX,
1,
text_color,
2,
)

return img

Expand All @@ -213,12 +304,9 @@ def draw_flag_with_label(img,
# INSTEAD OF THE FUNCTIONS BELOW


def draw_multiple_rectangles(img,
bboxes,
bbox_color=(255, 255, 255),
thickness=3,
is_opaque=False,
alpha=0.5):
def draw_multiple_rectangles(
img, bboxes, bbox_color=(255, 255, 255), thickness=3, is_opaque=False, alpha=0.5
):
"""draws multiple rectangles

img : ndarray
Expand All @@ -239,20 +327,20 @@ def draw_multiple_rectangles(img,
ndarray
the image with the bounding boxes drawn
"""

for bbox in bboxes:
img = draw_rectangle(img, bbox, bbox_color, thickness, is_opaque,
alpha)
img = draw_rectangle(img, bbox, bbox_color, thickness, is_opaque, alpha)
return img


def add_multiple_labels(img,
labels,
bboxes,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
top=True):
def add_multiple_labels(
img,
labels,
bboxes,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
top=True,
):
"""add labels, inside or outside the rectangles

Parameters
Expand All @@ -279,18 +367,19 @@ def add_multiple_labels(img,
"""

for label, bbox in zip(labels, bboxes):
img = add_label(img, label, bbox, draw_bg, text_bg_color, text_color,
top)
img = add_label(img, label, bbox, draw_bg, text_bg_color, text_color, top)

return img


def add_multiple_T_labels(img,
labels,
bboxes,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0)):
def add_multiple_T_labels(
img,
labels,
bboxes,
draw_bg=True,
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
):
"""adds T labels to the rectangles, each originating from the top of the rectangle

Parameters
Expand Down Expand Up @@ -320,13 +409,15 @@ def add_multiple_T_labels(img,
return img


def draw_multiple_flags_with_labels(img,
labels,
bboxes,
write_label=True,
line_color=(255, 255, 255),
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0)):
def draw_multiple_flags_with_labels(
img,
labels,
bboxes,
write_label=True,
line_color=(255, 255, 255),
text_bg_color=(255, 255, 255),
text_color=(0, 0, 0),
):
"""draws poles from the middle of the objects that are to be labeled and adds the labels to the flags

Parameters
Expand All @@ -353,6 +444,7 @@ def draw_multiple_flags_with_labels(img,
"""

for label, bbox in zip(labels, bboxes):
img = draw_flag_with_label(img, label, bbox, write_label, line_color,
text_bg_color, text_color)
img = draw_flag_with_label(
img, label, bbox, write_label, line_color, text_bg_color, text_color
)
return img
Binary file added images/test_images/source_bird.jpg
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading