text_rec.py 2.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566
  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 PIL
  15. from PIL import Image, ImageDraw, ImageFont
  16. import numpy as np
  17. import cv2
  18. from ...utils.fonts import PINGFANG_FONT_FILE_PATH
  19. from ...utils import logging
  20. from .base import CVResult
  21. class TextRecResult(CVResult):
  22. def _to_img(self):
  23. """Draw label on image"""
  24. image = self._img_reader.read(self["input_path"])
  25. rec_text = self["rec_text"]
  26. rec_score = self["rec_score"]
  27. image = image.convert("RGB")
  28. image_width, image_height = image.size
  29. text = f"{rec_text} ({rec_score})"
  30. font = self.adjust_font_size(image_width, text, PINGFANG_FONT_FILE_PATH)
  31. row_height = font.getbbox(text)[3]
  32. new_image_height = image_height + int(row_height * 1.2)
  33. new_image = Image.new("RGB", (image_width, new_image_height), (255, 255, 255))
  34. new_image.paste(image, (0, 0))
  35. draw = ImageDraw.Draw(new_image)
  36. draw.text(
  37. (0, image_height),
  38. text,
  39. fill=(0, 0, 0),
  40. font=font,
  41. )
  42. return new_image
  43. def adjust_font_size(self, image_width, text, font_path):
  44. font_size = int(image_width * 0.06)
  45. font = ImageFont.truetype(font_path, font_size)
  46. if int(PIL.__version__.split(".")[0]) < 10:
  47. text_width, _ = font.getsize(text)
  48. else:
  49. text_width, _ = font.getbbox(text)[2:]
  50. while text_width > image_width:
  51. font_size -= 1
  52. font = ImageFont.truetype(font_path, font_size)
  53. if int(PIL.__version__.split(".")[0]) < 10:
  54. text_width, _ = font.getsize(text)
  55. else:
  56. text_width, _ = font.getbbox(text)[2:]
  57. return font