Create ocr.py

ocr.py

@@ -0,0 +1,38 @@
+from paddleocr import PaddleOCR
+
+class OCR:
+    def __init__(self):
+        self.model_det  = 'models/det'
+        self.model_rec  = 'models/rec'
+        self.char_dict  = 'models/char_dict.txt'
+        self.ocr        = self.paddleocr() 
+    
+    def paddleocr(self):
+        return PaddleOCR(
+            det_model_dir       = self.model_det,
+            det_db_box_thresh   = 0.6,
+            rec_model_dir       = self.model_rec,
+            rec_image_shape     = '3, 32, 100',
+            rec_char_dict_path  = self.char_dict,
+            use_angle_cls       = True,
+            show_log            = False
+        )
+        
+    def ocr_image(self, image):
+        return self.ocr.ocr(image)
+    
+    @staticmethod
+    def extract_output(result):
+        '''
+        Extract the output of paddleocr.
+        Args:
+            result(list): The output of paddleocr.
+        Returns:
+            boxes(list): The list of boxes.
+            txts(list): The list of texts.
+            scores(list): The list of scores.
+        '''
+        boxes   = [line[0] for line in result]
+        txts    = [line[1][0] for line in result]
+        scores  = [line[1][1] for line in result]
+        return boxes, txts, scores