result.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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 copy
  15. import cv2
  16. import PIL
  17. from PIL import Image, ImageDraw, ImageFont
  18. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  19. from ...common.result import BaseCVResult, JsonMixin
  20. from ...utils.color_map import font_colormap, get_colormap
  21. def draw_attribute_result(img, boxes):
  22. """
  23. Args:
  24. img (PIL.Image.Image): PIL image
  25. boxes (list): a list of dictionaries representing detection box information.
  26. Returns:
  27. img (PIL.Image.Image): visualized image
  28. """
  29. font_size = int((0.024 * int(img.width) + 2) * 0.7)
  30. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  31. draw_thickness = int(max(img.size) * 0.005)
  32. draw = ImageDraw.Draw(img)
  33. label2color = {}
  34. catid2fontcolor = {}
  35. color_list = get_colormap(rgb=True)
  36. for i, dt in enumerate(boxes):
  37. text_lines, bbox, score = dt["label"], dt["coordinate"], dt["score"]
  38. if i not in label2color:
  39. color_index = i % len(color_list)
  40. label2color[i] = color_list[color_index]
  41. catid2fontcolor[i] = font_colormap(color_index)
  42. color = tuple(label2color[i]) + (255,)
  43. tuple(catid2fontcolor[i])
  44. xmin, ymin, xmax, ymax = bbox
  45. # draw box
  46. draw.line(
  47. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)],
  48. width=draw_thickness,
  49. fill=color,
  50. )
  51. # draw label
  52. current_y = ymin
  53. for line in text_lines:
  54. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  55. tw, th = draw.textsize(line, font=font)
  56. else:
  57. left, top, right, bottom = draw.textbbox((0, 0), line, font)
  58. tw, th = right - left, bottom - top + 4
  59. draw.text((5 + xmin + 1, current_y + 1), line, fill=(0, 0, 0), font=font)
  60. draw.text((5 + xmin, current_y), line, fill=color, font=font)
  61. current_y += th
  62. return img
  63. class AttributeRecResult(BaseCVResult):
  64. def _to_str(self, *args, **kwargs):
  65. data = copy.deepcopy(self)
  66. data.pop("input_img")
  67. return JsonMixin._to_str(data, *args, **kwargs)
  68. def _to_json(self, *args, **kwargs):
  69. data = copy.deepcopy(self)
  70. data.pop("input_img")
  71. return JsonMixin._to_json(data, *args, **kwargs)
  72. def _to_img(self):
  73. """apply"""
  74. image = Image.fromarray(cv2.cvtColor(self["input_img"], cv2.COLOR_BGR2RGB))
  75. boxes = [
  76. {
  77. "coordinate": box["coordinate"],
  78. "label": box["labels"],
  79. "score": box["det_score"],
  80. }
  81. for box in self["boxes"]
  82. ]
  83. image = draw_attribute_result(image, boxes)
  84. return {"res": image}