How to add Custom OCR - Githubissues

I tried doing this way,

import warnings
from tempfile import NamedTemporaryFile
from typing import List, Dict

import cv2
import numpy as np
import polars as pl

from img2table.document.base import Document
from img2table.ocr.base import OCRInstance
from img2table.ocr.data import OCRDataframe

from typing import List, Tuple
from surya.ocr import run_ocr
from surya.model.detection import model as det_model_org
from surya.model.recognition import model as rec_model_org, processor

class Table_SuryaOCR(OCRInstance):
    """
    SuryaOCR instance
    """

    def __init__(self, langs: List[str] = ["en", "ar"]):
        """
        Initialization of SuryaOCR instance
        :param lang: lang parameter used in SuryaOCR
        """
        if isinstance(langs, list):
            if all([isinstance(l, str) for l in langs]):
                self.langs = langs
        else:
            raise TypeError(f"Invalid type {type(lang)} for lang argument")

        self.rec_model = rec_model_org.load_model()
        self.rec_processor = processor.load_processor()
        self.det_model, self.det_processor = det_model_org.load_model(
        ), det_model_org.load_processor()

    def content(self, document: Document) -> List[List[Tuple]]:
        # Get OCR of all images
        ocrs = [self.parse_ocr_predictions(image) for image in document.images]
        return ocrs

    def parse_ocr_predictions(self,image):
        """Post processing on SuryaOCR predictions"""        
        predictions = run_ocr([Image.fromarray(image)], [self.langs], self.det_model,
                              self.det_processor, self.rec_model, self.rec_processor)

        ocr_result = [(i.text, i.bbox) for i in predictions[0].text_lines]

        return ocr_result

    def to_ocr_dataframe(self, content: List[List]) -> OCRDataframe:
        """
        Convert hOCR HTML to OCRDataframe object
        :param content: hOCR HTML string
        :return: OCRDataframe object corresponding to content
        """
        # Create list of elements
        list_elements = list()
        for page, ocr_result in enumerate(content):

            word_id = 0
            for word in ocr_result:
                word_id += 1
                dict_word = {
                    "page": 0,
                    "class": "ocrx_word",
                    "id": f"word_{page + 1}_{word_id}",
                    "parent": f"word_{page + 1}_{word_id}",
                    "value": word[0],
                    "confidence": round(0),
                    "x1": round(word[1][0]),
                    "y1": round(word[1][1]),
                    "x2": round(word[1][2]),
                    "y2": round(word[1][3])
                }

                list_elements.append(dict_word)

        return OCRDataframe(df=pl.LazyFrame(list_elements)) if list_elements else None
xavctn / img2table

How to add Custom OCR #209
