upgrade AbstractTextRecognizer (#5)

Co-authored-by: Nikita Shevtsov <[email protected]>

CHANGELOG.md

@@ -1,5 +1,8 @@
 Changelog
 =========
+v0.2.1 (2023-08-16)
+-------------------
+* `recognize_bbox` and `recognize_bboxes` methods are added in `AbstractTextRecognizer`
 
 v0.2 (2023-08-01)
 -------------------

VERSION

@@ -1 +1 @@
-0.2
+0.2.1

dedocutils/text_recognition/abstract_text_recognizer.py

@@ -1,8 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 
+from dedocutils.data_structures import BBox
+
 
 class AbstractTextRecognizer(ABC):
     """
@@ -18,3 +20,16 @@ def recognize(self, image: np.ndarray, parameters: Optional[dict] = None) -> str
         :return: recognized text
         """
         pass
+
+    def recognize_bbox(self, image: np.ndarray, bbox: BBox, need_rotate: bool = False) -> str:
+        line_image = image[bbox.y_top_left:bbox.y_bottom_right, bbox.x_top_left:bbox.x_bottom_right]
+
+        if need_rotate:
+            line_image = np.rot90(line_image)
+
+        line_image = np.pad(line_image, [(15, 15), (15, 15), (0, 0)], constant_values=255)
+        text = self.recognize(line_image)
+        return text
+
+    def recognize_bboxes(self, image: np.ndarray, bboxes: List[BBox]) -> List[str]:
+        return [self.recognize_bbox(image, bbox) for bbox in bboxes]

tests/unit_tests/test_classes.py

@@ -7,7 +7,6 @@
 
 
 class TestClasses(unittest.TestCase):
-
     def test_text_detection(self) -> None:
         file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "data", "document_example.png"))
         text_detector = DoctrTextDetector()

tests/unit_tests/test_text_recognizer.py

@@ -0,0 +1,28 @@
+import os.path
+import unittest
+
+import cv2
+
+from dedocutils.text_detection import DoctrTextDetector
+from dedocutils.text_recognition import TesseractTextRecognizer
+
+
+class TestTextRecognizer(unittest.TestCase):
+    @unittest.skip
+    def test_tesseract_recognize(self) -> None:
+        file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "data", "document_example.png"))
+        text_recognizer = TesseractTextRecognizer()
+        text = text_recognizer.recognize(cv2.imread(file_path))
+        self.assertIn("Document example", text)
+
+    @unittest.skip
+    def test_tesseract_recognize_bboxes(self) -> None:
+        file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "data", "document_example.png"))
+
+        text_detector = DoctrTextDetector()
+        text_recognizer = TesseractTextRecognizer()
+
+        img = cv2.imread(file_path)
+        bboxes = text_detector.detect(img)
+        texts = text_recognizer.recognize_bboxes(img, bboxes)
+        self.assertEquals("Document", texts[-1])