det.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  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, boxes):
  25. """
  26. Args:
  27. img (PIL.Image.Image): PIL image
  28. boxes (list): a list of dictionaries representing detection box information.
  29. Returns:
  30. img (PIL.Image.Image): visualized image
  31. """
  32. font_size = int(0.024 * int(img.width)) + 2
  33. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  34. draw_thickness = int(max(img.size) * 0.005)
  35. draw = ImageDraw.Draw(img)
  36. clsid2color = {}
  37. catid2fontcolor = {}
  38. color_list = get_colormap(rgb=True)
  39. for i, dt in enumerate(boxes):
  40. clsid, bbox, score = dt["cls_id"], dt["coordinate"], dt["score"]
  41. if clsid not in clsid2color:
  42. color_index = i % len(color_list)
  43. clsid2color[clsid] = color_list[color_index]
  44. catid2fontcolor[clsid] = font_colormap(color_index)
  45. color = tuple(clsid2color[clsid])
  46. font_color = tuple(catid2fontcolor[clsid])
  47. xmin, ymin, xmax, ymax = bbox
  48. # draw bbox
  49. draw.line(
  50. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)],
  51. width=draw_thickness,
  52. fill=color,
  53. )
  54. # draw label
  55. text = "{} {:.2f}".format(dt["label"], score)
  56. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  57. tw, th = draw.textsize(text, font=font)
  58. else:
  59. left, top, right, bottom = draw.textbbox((0, 0), text, font)
  60. tw, th = right - left, bottom - top
  61. if ymin < th:
  62. draw.rectangle([(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
  63. draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
  64. else:
  65. draw.rectangle([(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
  66. draw.text((xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
  67. return img
  68. class DetResult(BaseResult):
  69. """Save Result Transform"""
  70. def __init__(self, data):
  71. super().__init__(data)
  72. # We use pillow backend to save both numpy arrays and PIL Image objects
  73. self._img_reader.set_backend("pillow")
  74. self._img_writer.set_backend("pillow")
  75. def _get_res_img(self):
  76. """apply"""
  77. boxes = self["boxes"]
  78. img_path = self["img_path"]
  79. file_name = os.path.basename(img_path)
  80. image = self._img_reader.read(img_path)
  81. image = draw_box(image, boxes)
  82. return image