visualizer.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import numpy as np
  13. import json
  14. from pathlib import Path
  15. from PIL import Image, ImageDraw, ImageFont
  16. from pycocotools.coco import COCO
  17. from ......utils.fonts import PINGFANG_FONT_FILE_PATH
  18. from ......utils import logging
  19. def colormap(rgb=False):
  20. """
  21. Get colormap
  22. The code of this function is copied from https://github.com/facebookresearch/Detectron/blob/main/detectron/\
  23. utils/colormap.py
  24. """
  25. color_list = np.array([
  26. 0xFF, 0x00, 0x00, 0xCC, 0xFF, 0x00, 0x00, 0xFF, 0x66, 0x00, 0x66, 0xFF,
  27. 0xCC, 0x00, 0xFF, 0xFF, 0x4D, 0x00, 0x80, 0xff, 0x00, 0x00, 0xFF, 0xB2,
  28. 0x00, 0x1A, 0xFF, 0xFF, 0x00, 0xE5, 0xFF, 0x99, 0x00, 0x33, 0xFF, 0x00,
  29. 0x00, 0xFF, 0xFF, 0x33, 0x00, 0xFF, 0xff, 0x00, 0x99, 0xFF, 0xE5, 0x00,
  30. 0x00, 0xFF, 0x1A, 0x00, 0xB2, 0xFF, 0x80, 0x00, 0xFF, 0xFF, 0x00, 0x4D
  31. ]).astype(np.float32)
  32. color_list = (color_list.reshape((-1, 3)))
  33. if not rgb:
  34. color_list = color_list[:, ::-1]
  35. return color_list.astype('int32')
  36. def font_colormap(color_index):
  37. """
  38. Get font color according to the index of colormap
  39. """
  40. dark = np.array([0x14, 0x0E, 0x35])
  41. light = np.array([0xFF, 0xFF, 0xFF])
  42. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  43. if color_index in light_indexs:
  44. return light.astype('int32')
  45. else:
  46. return dark.astype('int32')
  47. def draw_bbox(image, coco_info: COCO, img_id):
  48. """
  49. Draw bbox on image
  50. """
  51. try:
  52. image_info = coco_info.loadImgs(img_id)[0]
  53. font_size = int(0.024 * int(image_info['width'])) + 2
  54. except:
  55. font_size = 12
  56. font = ImageFont.truetype(
  57. PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  58. image = image.convert('RGB')
  59. draw = ImageDraw.Draw(image)
  60. image_size = image.size
  61. width = int(max(image_size) * 0.005)
  62. catid2color = {}
  63. catid2fontcolor = {}
  64. catid_num_dict = {}
  65. color_list = colormap(rgb=True)
  66. annotations = coco_info.loadAnns(coco_info.getAnnIds(imgIds=img_id))
  67. for ann in annotations:
  68. catid = ann['category_id']
  69. catid_num_dict[catid] = catid_num_dict.get(catid, 0) + 1
  70. for i, (catid, _) in enumerate(
  71. sorted(
  72. catid_num_dict.items(), key=lambda x: x[1], reverse=True)):
  73. if catid not in catid2color:
  74. color_index = i % len(color_list)
  75. catid2color[catid] = color_list[color_index]
  76. catid2fontcolor[catid] = font_colormap(color_index)
  77. for ann in annotations:
  78. catid, bbox = ann['category_id'], ann['bbox']
  79. color = tuple(catid2color[catid])
  80. font_color = tuple(catid2fontcolor[catid])
  81. if len(bbox) == 4:
  82. # draw bbox
  83. xmin, ymin, w, h = bbox
  84. xmax = xmin + w
  85. ymax = ymin + h
  86. draw.line(
  87. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
  88. (xmin, ymin)],
  89. width=width,
  90. fill=color)
  91. elif len(bbox) == 8:
  92. x1, y1, x2, y2, x3, y3, x4, y4 = bbox
  93. draw.line(
  94. [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)],
  95. width=width,
  96. fill=color)
  97. xmin = min(x1, x2, x3, x4)
  98. ymin = min(y1, y2, y3, y4)
  99. else:
  100. logging.info('Error: The shape of bbox must be [M, 4] or [M, 8]!')
  101. # draw label
  102. label = coco_info.loadCats(catid)[0]['name']
  103. text = "{}".format(label)
  104. tw, th = draw.textsize(text, font=font)
  105. if ymin < th:
  106. draw.rectangle(
  107. [(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
  108. draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
  109. else:
  110. draw.rectangle(
  111. [(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
  112. draw.text(
  113. (xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
  114. return image
  115. def draw_mask(image, coco_info: COCO, img_id):
  116. """
  117. Draw mask on image
  118. """
  119. mask_color_id = 0
  120. w_ratio = .4
  121. alpha = 0.6
  122. color_list = colormap(rgb=True)
  123. img_array = np.array(image).astype('float32')
  124. h, w = img_array.shape[:2]
  125. annotations = coco_info.loadAnns(coco_info.getAnnIds(imgIds=img_id))
  126. for ann in annotations:
  127. segm = ann['segmentation']
  128. if not segm:
  129. continue
  130. import pycocotools.mask as mask_util
  131. rles = mask_util.frPyObjects(segm, h, w)
  132. rle = mask_util.merge(rles)
  133. mask = mask_util.decode(rle) * 255
  134. color_mask = color_list[mask_color_id % len(color_list), 0:3]
  135. mask_color_id += 1
  136. for c in range(3):
  137. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  138. idx = np.nonzero(mask)
  139. img_array[idx[0], idx[1], :] *= 1.0 - alpha
  140. img_array[idx[0], idx[1], :] += alpha * color_mask
  141. return Image.fromarray(img_array.astype('uint8'))