result.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  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. from pathlib import Path
  16. import copy
  17. import math
  18. import random
  19. import numpy as np
  20. import cv2
  21. import PIL
  22. from PIL import Image, ImageDraw, ImageFont
  23. from ....utils.fonts import PINGFANG_FONT_FILE_PATH, create_font
  24. from ...common.result import BaseCVResult
  25. class OCRResult(BaseCVResult):
  26. """OCR result"""
  27. def get_minarea_rect(self, points: np.ndarray) -> np.ndarray:
  28. """
  29. Get the minimum area rectangle for the given points using OpenCV.
  30. Args:
  31. points (np.ndarray): An array of 2D points.
  32. Returns:
  33. np.ndarray: An array of 2D points representing the corners of the minimum area rectangle
  34. in a specific order (clockwise or counterclockwise starting from the top-left corner).
  35. """
  36. bounding_box = cv2.minAreaRect(points)
  37. points = sorted(list(cv2.boxPoints(bounding_box)), key=lambda x: x[0])
  38. index_a, index_b, index_c, index_d = 0, 1, 2, 3
  39. if points[1][1] > points[0][1]:
  40. index_a = 0
  41. index_d = 1
  42. else:
  43. index_a = 1
  44. index_d = 0
  45. if points[3][1] > points[2][1]:
  46. index_b = 2
  47. index_c = 3
  48. else:
  49. index_b = 3
  50. index_c = 2
  51. box = np.array(
  52. [points[index_a], points[index_b], points[index_c], points[index_d]]
  53. ).astype(np.int32)
  54. return box
  55. def _to_img(self) -> PIL.Image:
  56. """
  57. Converts the internal data to a PIL Image with detection and recognition results.
  58. Returns:
  59. PIL.Image: An image with detection boxes, texts, and scores blended on it.
  60. """
  61. boxes = self["rec_boxes"]
  62. txts = self["rec_texts"]
  63. image = self["doc_preprocessor_res"]["output_img"]
  64. h, w = image.shape[0:2]
  65. image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  66. img_left = Image.fromarray(image_rgb)
  67. img_right = np.ones((h, w, 3), dtype=np.uint8) * 255
  68. random.seed(0)
  69. draw_left = ImageDraw.Draw(img_left)
  70. for idx, (box, txt) in enumerate(zip(boxes, txts)):
  71. try:
  72. color = (
  73. random.randint(0, 255),
  74. random.randint(0, 255),
  75. random.randint(0, 255),
  76. )
  77. box = np.array(box)
  78. if len(box) > 4:
  79. pts = [(x, y) for x, y in box.tolist()]
  80. draw_left.polygon(pts, outline=color, width=8)
  81. box = self.get_minarea_rect(box)
  82. height = int(0.5 * (max(box[:, 1]) - min(box[:, 1])))
  83. box[:2, 1] = np.mean(box[:, 1])
  84. box[2:, 1] = np.mean(box[:, 1]) + min(20, height)
  85. draw_left.polygon(box, fill=color)
  86. img_right_text = draw_box_txt_fine(
  87. (w, h), box, txt, PINGFANG_FONT_FILE_PATH
  88. )
  89. pts = np.array(box, np.int32).reshape((-1, 1, 2))
  90. cv2.polylines(img_right_text, [pts], True, color, 1)
  91. img_right = cv2.bitwise_and(img_right, img_right_text)
  92. except:
  93. continue
  94. img_left = Image.blend(Image.fromarray(image_rgb), img_left, 0.5)
  95. img_show = Image.new("RGB", (w * 2, h), (255, 255, 255))
  96. img_show.paste(img_left, (0, 0, w, h))
  97. img_show.paste(Image.fromarray(img_right), (w, 0, w * 2, h))
  98. model_settings = self["model_settings"]
  99. if model_settings["use_doc_preprocessor"]:
  100. return {
  101. **self["doc_preprocessor_res"].img,
  102. f"ocr_res_img": img_show,
  103. }
  104. else:
  105. return {f"ocr_res_img": img_show}
  106. # Adds a function comment according to Google Style Guide
  107. def draw_box_txt_fine(
  108. img_size: tuple, box: np.ndarray, txt: str, font_path: str
  109. ) -> np.ndarray:
  110. """
  111. Draws text in a box on an image with fine control over size and orientation.
  112. Args:
  113. img_size (tuple): The size of the output image (width, height).
  114. box (np.ndarray): A 4x2 numpy array defining the corners of the box in (x, y) order.
  115. txt (str): The text to draw inside the box.
  116. font_path (str): The path to the font file to use for drawing the text.
  117. Returns:
  118. np.ndarray: An image with the text drawn in the specified box.
  119. """
  120. box_height = int(
  121. math.sqrt((box[0][0] - box[3][0]) ** 2 + (box[0][1] - box[3][1]) ** 2)
  122. )
  123. box_width = int(
  124. math.sqrt((box[0][0] - box[1][0]) ** 2 + (box[0][1] - box[1][1]) ** 2)
  125. )
  126. if box_height > 2 * box_width and box_height > 30:
  127. img_text = Image.new("RGB", (box_height, box_width), (255, 255, 255))
  128. draw_text = ImageDraw.Draw(img_text)
  129. if txt:
  130. font = create_font(txt, (box_height, box_width), font_path)
  131. draw_text.text([0, 0], txt, fill=(0, 0, 0), font=font)
  132. img_text = img_text.transpose(Image.ROTATE_270)
  133. else:
  134. img_text = Image.new("RGB", (box_width, box_height), (255, 255, 255))
  135. draw_text = ImageDraw.Draw(img_text)
  136. if txt:
  137. font = create_font(txt, (box_width, box_height), font_path)
  138. draw_text.text([0, 0], txt, fill=(0, 0, 0), font=font)
  139. pts1 = np.float32(
  140. [[0, 0], [box_width, 0], [box_width, box_height], [0, box_height]]
  141. )
  142. pts2 = np.array(box, dtype=np.float32)
  143. M = cv2.getPerspectiveTransform(pts1, pts2)
  144. img_text = np.array(img_text, dtype=np.uint8)
  145. img_right_text = cv2.warpPerspective(
  146. img_text,
  147. M,
  148. img_size,
  149. flags=cv2.INTER_NEAREST,
  150. borderMode=cv2.BORDER_CONSTANT,
  151. borderValue=(255, 255, 255),
  152. )
  153. return img_right_text