det.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import cv2
  16. import numpy as np
  17. import PIL
  18. from PIL import Image, ImageDraw, ImageFont
  19. from ...utils.fonts import PINGFANG_FONT_FILE_PATH
  20. from ..utils.color_map import get_colormap, font_colormap
  21. from .base import CVResult
  22. def draw_box(img, boxes):
  23. """
  24. Args:
  25. img (PIL.Image.Image): PIL image
  26. boxes (list): a list of dictionaries representing detection box information.
  27. Returns:
  28. img (PIL.Image.Image): visualized image
  29. """
  30. font_size = int(0.018 * int(img.width)) + 2
  31. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  32. draw_thickness = int(max(img.size) * 0.002)
  33. draw = ImageDraw.Draw(img)
  34. label2color = {}
  35. catid2fontcolor = {}
  36. color_list = get_colormap(rgb=True)
  37. for i, dt in enumerate(boxes):
  38. # clsid = dt["cls_id"]
  39. label, bbox, score = dt["label"], dt["coordinate"], dt["score"]
  40. if label not in label2color:
  41. color_index = i % len(color_list)
  42. label2color[label] = color_list[color_index]
  43. catid2fontcolor[label] = font_colormap(color_index)
  44. color = tuple(label2color[label])
  45. font_color = tuple(catid2fontcolor[label])
  46. if len(bbox) == 4:
  47. # draw bbox of normal object detection
  48. xmin, ymin, xmax, ymax = bbox
  49. rectangle = [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)]
  50. elif len(bbox) == 8:
  51. # draw bbox of rotated object detection
  52. x1, y1, x2, y2, x3, y3, x4, y4 = bbox
  53. rectangle = [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)]
  54. xmin = min(x1, x2, x3, x4)
  55. ymin = min(y1, y2, y3, y4)
  56. else:
  57. raise ValueError(
  58. f"Only support bbox format of [xmin,ymin,xmax,ymax] or [x1,y1,x2,y2,x3,y3,x4,y4], got bbox of shape {len(bbox)}."
  59. )
  60. # draw bbox
  61. draw.line(
  62. rectangle,
  63. width=draw_thickness,
  64. fill=color,
  65. )
  66. # draw label
  67. text = "{} {:.2f}".format(dt["label"], score)
  68. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  69. tw, th = draw.textsize(text, font=font)
  70. else:
  71. left, top, right, bottom = draw.textbbox((0, 0), text, font)
  72. tw, th = right - left, bottom - top + 4
  73. if ymin < th:
  74. draw.rectangle([(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
  75. draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
  76. else:
  77. draw.rectangle([(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
  78. draw.text((xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
  79. return img
  80. class DetResult(CVResult):
  81. """Save Result Transform"""
  82. def _to_img(self):
  83. """apply"""
  84. boxes = self["boxes"]
  85. image = self._img_reader.read(self["input_path"])
  86. image = draw_box(image, boxes)
  87. return image