det.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  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 numpy as np
  16. import math
  17. import PIL
  18. from PIL import Image, ImageDraw, ImageFont
  19. from ...utils import logging
  20. from ...utils.fonts import PINGFANG_FONT_FILE_PATH
  21. from ..utils.io import ImageWriter, ImageReader
  22. from ..utils.color_map import get_colormap, font_colormap
  23. from .base import BaseResult
  24. def draw_box(img, np_boxes, labels):
  25. """
  26. Args:
  27. img (PIL.Image.Image): PIL image
  28. np_boxes (np.ndarray): shape:[N,6], N: number of box,
  29. matix element:[class, score, x_min, y_min, x_max, y_max]
  30. labels (list): labels:['class1', ..., 'classn']
  31. Returns:
  32. img (PIL.Image.Image): visualized image
  33. """
  34. font_size = int(0.024 * int(img.width)) + 2
  35. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  36. draw_thickness = int(max(img.size) * 0.005)
  37. draw = ImageDraw.Draw(img)
  38. clsid2color = {}
  39. catid2fontcolor = {}
  40. color_list = get_colormap(rgb=True)
  41. for i, dt in enumerate(np_boxes):
  42. clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
  43. if clsid not in clsid2color:
  44. color_index = i % len(color_list)
  45. clsid2color[clsid] = color_list[color_index]
  46. catid2fontcolor[clsid] = font_colormap(color_index)
  47. color = tuple(clsid2color[clsid])
  48. font_color = tuple(catid2fontcolor[clsid])
  49. xmin, ymin, xmax, ymax = bbox
  50. # draw bbox
  51. draw.line(
  52. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)],
  53. width=draw_thickness,
  54. fill=color,
  55. )
  56. # draw label
  57. text = "{} {:.2f}".format(labels[clsid], score)
  58. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  59. tw, th = draw.textsize(text, font=font)
  60. else:
  61. left, top, right, bottom = draw.textbbox((0, 0), text, font)
  62. tw, th = right - left, bottom - top
  63. if ymin < th:
  64. draw.rectangle([(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
  65. draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
  66. else:
  67. draw.rectangle([(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
  68. draw.text((xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
  69. return img
  70. class DetResult(BaseResult):
  71. """Save Result Transform"""
  72. def __init__(self, data):
  73. super().__init__(data)
  74. self.data = data
  75. # We use pillow backend to save both numpy arrays and PIL Image objects
  76. self._img_reader.set_backend("pillow")
  77. self._img_writer.set_backend("pillow")
  78. def _get_res_img(self):
  79. """apply"""
  80. boxes = self["boxes"]
  81. img_path = self["img_path"]
  82. labels = self.data["labels"]
  83. file_name = os.path.basename(img_path)
  84. image = self._img_reader.read(img_path)
  85. image = draw_box(image, boxes, labels=labels)
  86. return image