00001 import six
00002 import cv2
00003
00004
00005 def voc_colormap(nlabels):
00006 colors = []
00007 for i in six.moves.range(nlabels):
00008 r, g, b = 0, 0, 0
00009 for j in range(8):
00010 if i & (1 << 0):
00011 r |= 1 << (7 - j)
00012 if i & (1 << 1):
00013 g |= 1 << (7 - j)
00014 if i & (1 << 2):
00015 b |= 1 << (7 - j)
00016 i >>= 3
00017 colors.append([r, g, b])
00018 return colors
00019
00020
00021 def vis_bboxes(img, bboxes, labels,
00022 font_scale=0.8,
00023 thickness=1,
00024 font_face=cv2.FONT_HERSHEY_SIMPLEX,
00025 text_color=(255, 255, 255),
00026 max_label_num=1024):
00027 """Visualize bounding boxes inside image.
00028
00029 """
00030 if len(bboxes) != len(labels):
00031 raise ValueError("len(bboxes) and len(labels) should be same "
00032 "we get len(bboxes):{}, len(lables):{}"
00033 .format(len(bboxes), len(labels)))
00034 colormap = voc_colormap(max_label_num)
00035
00036 CV_AA = 16
00037 for bbox, label in zip(bboxes, labels):
00038 color = colormap[label % max_label_num]
00039 x1, y1, w, h = bbox
00040 x2 = x1 + w
00041 y2 = y1 + h
00042 x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
00043 cv2.rectangle(img, (x1, y1), (x2, y2), color, 2, CV_AA)
00044
00045 label_name = str(label)
00046 img_bbox = img[y1:y2, x1:]
00047
00048 text = label_name
00049 size, baseline = cv2.getTextSize(
00050 text, font_face, font_scale, thickness)
00051 cv2.rectangle(
00052 img_bbox, (0, 0), (size[0], size[1] + baseline),
00053 color=color, thickness=-1)
00054 cv2.putText(img_bbox, text, (0, size[1]),
00055 font_face, font_scale, text_color, thickness)
00056 return img