|
5 | 5 |
|
6 | 6 | import imageio
|
7 | 7 | import numpy as np
|
8 |
| - |
9 | 8 | import tensorlayer as tl
|
10 | 9 | from tensorlayer.lazy_imports import LazyImport
|
| 10 | +import colorsys, random |
11 | 11 |
|
12 | 12 | cv2 = LazyImport("cv2")
|
13 | 13 |
|
|
16 | 16 | # matplotlib.use('Agg')
|
17 | 17 |
|
18 | 18 | __all__ = [
|
19 |
| - 'read_image', |
20 |
| - 'read_images', |
21 |
| - 'save_image', |
22 |
| - 'save_images', |
23 |
| - 'draw_boxes_and_labels_to_image', |
24 |
| - 'draw_mpii_people_to_image', |
25 |
| - 'frame', |
26 |
| - 'CNN2d', |
27 |
| - 'images2d', |
28 |
| - 'tsne_embedding', |
29 |
| - 'draw_weights', |
30 |
| - 'W', |
| 19 | + 'read_image', 'read_images', 'save_image', 'save_images', 'draw_boxes_and_labels_to_image', |
| 20 | + 'draw_mpii_people_to_image', 'frame', 'CNN2d', 'images2d', 'tsne_embedding', 'draw_weights', 'W', |
| 21 | + 'draw_boxes_and_labels_to_image_with_json' |
31 | 22 | ]
|
32 | 23 |
|
33 | 24 |
|
@@ -662,3 +653,66 @@ def draw_weights(W=None, second=10, saveable=True, shape=None, name='mnist', fig
|
662 | 653 |
|
663 | 654 |
|
664 | 655 | W = draw_weights
|
| 656 | + |
| 657 | + |
| 658 | +def draw_boxes_and_labels_to_image_with_json(image, json_result, class_list, save_name=None): |
| 659 | + """Draw bboxes and class labels on image. Return the image with bboxes. |
| 660 | +
|
| 661 | + Parameters |
| 662 | + ----------- |
| 663 | + image : numpy.array |
| 664 | + The RGB image [height, width, channel]. |
| 665 | + json_result : list of dict |
| 666 | + The object detection result with json format. |
| 667 | + classes_list : list of str |
| 668 | + For converting ID to string on image. |
| 669 | + save_name : None or str |
| 670 | + The name of image file (i.e. image.png), if None, not to save image. |
| 671 | +
|
| 672 | + Returns |
| 673 | + ------- |
| 674 | + numpy.array |
| 675 | + The saved image. |
| 676 | +
|
| 677 | + References |
| 678 | + ----------- |
| 679 | + - OpenCV rectangle and putText. |
| 680 | + - `scikit-image <http://scikit-image.org/docs/dev/api/skimage.draw.html#skimage.draw.rectangle>`__. |
| 681 | +
|
| 682 | + """ |
| 683 | + image_h, image_w, _ = image.shape |
| 684 | + num_classes = len(class_list) |
| 685 | + hsv_tuples = [(1.0 * x / num_classes, 1., 1.) for x in range(num_classes)] |
| 686 | + colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) |
| 687 | + colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors)) |
| 688 | + random.seed(0) |
| 689 | + random.shuffle(colors) |
| 690 | + random.seed(None) |
| 691 | + bbox_thick = int(0.6 * (image_h + image_w) / 600) |
| 692 | + fontScale = 0.5 |
| 693 | + |
| 694 | + for bbox_info in json_result: |
| 695 | + image_name = bbox_info['image'] |
| 696 | + category_id = bbox_info['category_id'] |
| 697 | + if category_id < 0 or category_id > num_classes: continue |
| 698 | + bbox = bbox_info['bbox'] # the order of coordinates is [x1, y2, x2, y2] |
| 699 | + score = bbox_info['score'] |
| 700 | + |
| 701 | + bbox_color = colors[category_id] |
| 702 | + c1, c2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])) |
| 703 | + cv2.rectangle(image, c1, c2, bbox_color, bbox_thick) |
| 704 | + |
| 705 | + bbox_mess = '%s: %.2f' % (class_list[category_id], score) |
| 706 | + t_size = cv2.getTextSize(bbox_mess, 0, fontScale, thickness=bbox_thick // 2)[0] |
| 707 | + c3 = (c1[0] + t_size[0], c1[1] - t_size[1] - 3) |
| 708 | + cv2.rectangle(image, c1, (np.float32(c3[0]), np.float32(c3[1])), bbox_color, -1) |
| 709 | + |
| 710 | + cv2.putText( |
| 711 | + image, bbox_mess, (c1[0], np.float32(c1[1] - 2)), cv2.FONT_HERSHEY_SIMPLEX, fontScale, (0, 0, 0), |
| 712 | + bbox_thick // 2, lineType=cv2.LINE_AA |
| 713 | + ) |
| 714 | + |
| 715 | + if save_name is not None: |
| 716 | + save_image(image, save_name) |
| 717 | + |
| 718 | + return image |
0 commit comments