Spaces:
Sleeping
Sleeping
| from paddleocr import PaddleOCR | |
| from vietocr.tool.config import Cfg | |
| from vietocr.tool.predictor import Predictor | |
| from utils.config import Config | |
| import requests | |
| import numpy as np | |
| from PIL import Image, ImageTransform | |
| class OCRDetector: | |
| def __init__(self) -> None: | |
| self.paddle_ocr = PaddleOCR( | |
| lang="en", | |
| use_angle_cls=False, | |
| use_gpu=True if Config.device == "cpu" else False, | |
| show_log=False, | |
| ) | |
| # config['weights'] = './weights/transformerocr.pth' | |
| vietocr_config = Cfg.load_config_from_name("vgg_transformer") | |
| vietocr_config["weights"] = Config.ocr_path | |
| vietocr_config["cnn"]["pretrained"] = False | |
| vietocr_config["device"] = Config.device | |
| vietocr_config["predictor"]["beamsearch"] = False | |
| self.viet_ocr = Predictor(vietocr_config) | |
| def find_box(self, image): | |
| """Xác định box dựa vào mô hình paddle_ocr""" | |
| result = self.paddle_ocr.ocr(image, cls=False, rec=False) | |
| result = result[0] | |
| # Extracting detected components | |
| boxes = result # [res[0] for res in result] | |
| boxes = np.array(boxes).astype(int) | |
| # scores = [res[1][1] for res in result] | |
| return boxes | |
| def cut_image_polygon(self, image, box): | |
| (x1, y1), (x2, y2), (x3, y3), (x4, y4) = box | |
| w = x2 - x1 | |
| h = y4 - y1 | |
| scl = h // 7 | |
| new_box = ( | |
| [max(x1 - scl, 0), max(y1 - scl, 0)], | |
| [x2 + scl, y2 - scl], | |
| [x3 + scl, y3 + scl], | |
| [x4 - scl, y4 + scl], | |
| ) | |
| (x1, y1), (x2, y2), (x3, y3), (x4, y4) = new_box | |
| # Define 8-tuple with x,y coordinates of top-left, bottom-left, bottom-right and top-right corners and apply | |
| transform = [x1, y1, x4, y4, x3, y3, x2, y2] | |
| result = image.transform((w, h), ImageTransform.QuadTransform(transform)) | |
| return result | |
| def vietnamese_text(self, boxes, image): | |
| """Xác định text dựa vào mô hình viet_ocr""" | |
| results = [] | |
| for box in boxes: | |
| try: | |
| cut_image = self.cut_image_polygon(image, box) | |
| # cut_image = Image.fromarray(np.uint8(cut_image)) | |
| text, score = self.viet_ocr.predict(cut_image, return_prob=True) | |
| if score > Config.vietocr_threshold: | |
| results.append({"text": text, "score": score, "box": box}) | |
| except: | |
| continue | |
| return results | |
| # Merge | |
| def text_detector(self, image_path): | |
| if image_path.startswith("https://"): | |
| image = Image.open(requests.get(image_path, stream=True).raw).convert("RGB") | |
| else: | |
| image = Image.open(image_path).convert("RGB") | |
| # np_image = np.array(image) | |
| boxes = self.find_box(image_path) | |
| if not boxes.any(): | |
| return None | |
| results = self.vietnamese_text(boxes, image) | |
| if results != []: | |
| return results | |
| else: | |
| return None | |