Skip to content

Commit

Permalink
Draw labels name (apache#9496)
Browse files Browse the repository at this point in the history
* Draw labels name

Draw labels name next to corresponding bounding boxes when the mapping of id to names is specified.

* Update detection.py

* Update detection.py
Fixed missing variable declaration
Fixing lint issues
  • Loading branch information
lvaleriu authored and anirudh2290 committed Sep 19, 2018
1 parent 9f6eca7 commit 01fd82b
Showing 1 changed file with 14 additions and 1 deletion.
15 changes: 14 additions & 1 deletion python/mxnet/image/detection.py
Original file line number Diff line number Diff line change
Expand Up @@ -803,7 +803,7 @@ def check_label_shape(self, label_shape):
raise ValueError(msg)

def draw_next(self, color=None, thickness=2, mean=None, std=None, clip=True,
waitKey=None, window_name='draw_next'):
waitKey=None, window_name='draw_next', id2labels=None):
"""Display next image with bounding boxes drawn.
Parameters
Expand All @@ -822,6 +822,8 @@ def draw_next(self, color=None, thickness=2, mean=None, std=None, clip=True,
Hold the window for waitKey milliseconds if set, skip ploting if None
window_name : str
Plot window name if waitKey is set.
id2labels : dict
Mapping of labels id to labels name.
Returns
-------
Expand Down Expand Up @@ -889,6 +891,17 @@ def draw_next(self, color=None, thickness=2, mean=None, std=None, clip=True,
y2 = int(label[i, 4] * height)
bc = np.random.rand(3) * 255 if not color else color
cv2.rectangle(image, (x1, y1), (x2, y2), bc, thickness)
if id2labels is not None:
cls_id = int(label[i, 0])
if cls_id in id2labels:
cls_name = id2labels[cls_id]
text = "{:s}".format(cls_name)
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.5
text_height = cv2.getTextSize(text, font, font_scale, 2)[0][1]
tc = (255, 255, 255)
tpos = (x1 + 5, y1 + text_height + 5)
cv2.putText(image, text, tpos, font, font_scale, tc, 2)
if waitKey is not None:
cv2.imshow(window_name, image)
cv2.waitKey(waitKey)
Expand Down

0 comments on commit 01fd82b

Please sign in to comment.