visualizer.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. # -*- coding: UTF-8 -*-
  2. ################################################################################
  3. #
  4. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  5. #
  6. ################################################################################
  7. """
  8. Author: PaddlePaddle Authors
  9. """
  10. import os
  11. import numpy as np
  12. import json
  13. from pathlib import Path
  14. import PIL
  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. [
  27. 0xFF,
  28. 0x00,
  29. 0x00,
  30. 0xCC,
  31. 0xFF,
  32. 0x00,
  33. 0x00,
  34. 0xFF,
  35. 0x66,
  36. 0x00,
  37. 0x66,
  38. 0xFF,
  39. 0xCC,
  40. 0x00,
  41. 0xFF,
  42. 0xFF,
  43. 0x4D,
  44. 0x00,
  45. 0x80,
  46. 0xFF,
  47. 0x00,
  48. 0x00,
  49. 0xFF,
  50. 0xB2,
  51. 0x00,
  52. 0x1A,
  53. 0xFF,
  54. 0xFF,
  55. 0x00,
  56. 0xE5,
  57. 0xFF,
  58. 0x99,
  59. 0x00,
  60. 0x33,
  61. 0xFF,
  62. 0x00,
  63. 0x00,
  64. 0xFF,
  65. 0xFF,
  66. 0x33,
  67. 0x00,
  68. 0xFF,
  69. 0xFF,
  70. 0x00,
  71. 0x99,
  72. 0xFF,
  73. 0xE5,
  74. 0x00,
  75. 0x00,
  76. 0xFF,
  77. 0x1A,
  78. 0x00,
  79. 0xB2,
  80. 0xFF,
  81. 0x80,
  82. 0x00,
  83. 0xFF,
  84. 0xFF,
  85. 0x00,
  86. 0x4D,
  87. ]
  88. ).astype(np.float32)
  89. color_list = color_list.reshape((-1, 3))
  90. if not rgb:
  91. color_list = color_list[:, ::-1]
  92. return color_list.astype("int32")
  93. def font_colormap(color_index):
  94. """
  95. Get font color according to the index of colormap
  96. """
  97. dark = np.array([0x14, 0x0E, 0x35])
  98. light = np.array([0xFF, 0xFF, 0xFF])
  99. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  100. if color_index in light_indexs:
  101. return light.astype("int32")
  102. else:
  103. return dark.astype("int32")
  104. def draw_bbox(image, coco_info: COCO, img_id):
  105. """
  106. Draw bbox on image
  107. """
  108. try:
  109. image_info = coco_info.loadImgs(img_id)[0]
  110. font_size = int(0.024 * int(image_info["width"])) + 2
  111. except:
  112. font_size = 12
  113. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  114. image = image.convert("RGB")
  115. draw = ImageDraw.Draw(image)
  116. image_size = image.size
  117. width = int(max(image_size) * 0.005)
  118. catid2color = {}
  119. catid2fontcolor = {}
  120. catid_num_dict = {}
  121. color_list = colormap(rgb=True)
  122. annotations = coco_info.loadAnns(coco_info.getAnnIds(imgIds=img_id))
  123. for ann in annotations:
  124. catid = ann["category_id"]
  125. catid_num_dict[catid] = catid_num_dict.get(catid, 0) + 1
  126. for i, (catid, _) in enumerate(
  127. sorted(catid_num_dict.items(), key=lambda x: x[1], reverse=True)
  128. ):
  129. if catid not in catid2color:
  130. color_index = i % len(color_list)
  131. catid2color[catid] = color_list[color_index]
  132. catid2fontcolor[catid] = font_colormap(color_index)
  133. for ann in annotations:
  134. catid, bbox = ann["category_id"], ann["bbox"]
  135. color = tuple(catid2color[catid])
  136. font_color = tuple(catid2fontcolor[catid])
  137. if len(bbox) == 4:
  138. # draw bbox
  139. xmin, ymin, w, h = bbox
  140. xmax = xmin + w
  141. ymax = ymin + h
  142. draw.line(
  143. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)],
  144. width=width,
  145. fill=color,
  146. )
  147. elif len(bbox) == 8:
  148. x1, y1, x2, y2, x3, y3, x4, y4 = bbox
  149. draw.line(
  150. [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)],
  151. width=width,
  152. fill=color,
  153. )
  154. xmin = min(x1, x2, x3, x4)
  155. ymin = min(y1, y2, y3, y4)
  156. else:
  157. logging.info("Error: The shape of bbox must be [M, 4] or [M, 8]!")
  158. # draw label
  159. label = coco_info.loadCats(catid)[0]["name"]
  160. text = "{}".format(label)
  161. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  162. tw, th = draw.textsize(text, font=font)
  163. else:
  164. left, top, right, bottom = draw.textbbox((0, 0), text, font)
  165. tw, th = right - left, bottom - top
  166. if ymin < th:
  167. draw.rectangle([(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
  168. draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
  169. else:
  170. draw.rectangle([(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
  171. draw.text((xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
  172. return image