result.py 3.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  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 cv2
  15. import numpy as np
  16. import PIL
  17. from PIL import Image, ImageDraw, ImageFont
  18. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  19. from ...utils.color_map import get_colormap
  20. from ...common.result import BaseVideoResult
  21. class TopkVideoResult(BaseVideoResult):
  22. def _to_video(self):
  23. """Draw label on image"""
  24. labels = self.get("label_names", self["class_ids"])
  25. label_str = f"{labels[0]} {self['scores'][0]:.2f}"
  26. video_reader = self._video_reader
  27. video = video_reader.read(self["input_path"])
  28. video = list(video)
  29. write_fps = video_reader.get_fps()
  30. video_list = []
  31. for i in range(len(video)):
  32. image = Image.fromarray(video[i].asnumpy())
  33. image_size = image.size
  34. draw = ImageDraw.Draw(image)
  35. min_font_size = int(image_size[0] * 0.02)
  36. max_font_size = int(image_size[0] * 0.05)
  37. for font_size in range(max_font_size, min_font_size - 1, -1):
  38. font = ImageFont.truetype(
  39. PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8"
  40. )
  41. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  42. text_width_tmp, text_height_tmp = draw.textsize(label_str, font)
  43. else:
  44. left, top, right, bottom = draw.textbbox((0, 0), label_str, font)
  45. text_width_tmp, text_height_tmp = right - left, bottom - top
  46. if text_width_tmp <= image_size[0]:
  47. break
  48. else:
  49. font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, min_font_size)
  50. color_list = get_colormap(rgb=True)
  51. color = tuple(color_list[0])
  52. font_color = tuple(self._get_font_colormap(3))
  53. if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
  54. text_width, text_height = draw.textsize(label_str, font)
  55. else:
  56. left, top, right, bottom = draw.textbbox((0, 0), label_str, font)
  57. text_width, text_height = right - left, bottom - top
  58. rect_left = 3
  59. rect_top = 3
  60. rect_right = rect_left + text_width + 3
  61. rect_bottom = rect_top + text_height + 6
  62. draw.rectangle(
  63. [(rect_left, rect_top), (rect_right, rect_bottom)], fill=color
  64. )
  65. text_x = rect_left + 3
  66. text_y = rect_top
  67. draw.text((text_x, text_y), label_str, fill=font_color, font=font)
  68. image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
  69. video_list.append(image)
  70. return np.array(video_list), write_fps
  71. def _get_font_colormap(self, color_index):
  72. """
  73. Get font colormap
  74. """
  75. dark = np.array([0x14, 0x0E, 0x35])
  76. light = np.array([0xFF, 0xFF, 0xFF])
  77. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  78. if color_index in light_indexs:
  79. return light.astype("int32")
  80. else:
  81. return dark.astype("int32")