How to add Custom OCR

[akshayrakate085]

Hi, I wanted to extract the table data using the Surya-ocr. How can I achieve that?

[akshayrakate085]

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

[DipanshuJuneja]

This is my code that I'm using. The issue is that I'm not seeing any tables returned and no errors, however all the print statements in the code below seem to be executing fine and I'm seeing the OCR results of the document. I'm passing my custom model to the object at the time of instantiating the class.

from tempfile import NamedTemporaryFile
from typing import List
import cv2
import numpy as np
import polars as pl
import os

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

from typing import List

class RapidImage2Table(OCRInstance):
    """
    """
    def __init__(self, ocr_model):
        self.ocr = ocr_model

    def hocr(self, image: np.ndarray) -> List:
        with NamedTemporaryFile(suffix='.jpg', delete=False) as tmp_f:
            tmp_file = tmp_f.name
            # Write image to temporary file
            cv2.imwrite(tmp_file, image)

            # Get OCR
            ocr_result = self.ocr(tmp_file)

        # Remove temporary file
        while os.path.exists(tmp_file):
            try:
                os.remove(tmp_file)
            except PermissionError:
                pass

        # Get result
        print("OCR RESULT", ocr_result)
        # ocr_result = ocr_result.pop()
        return [[bbox, word, round(conf, 2)] for bbox, word, conf in ocr_result[0]] if ocr_result else []

    def content(self, document: Document) -> List[List]:
        # Get OCR of all images
        # print("OCR, IMG", self.ocr(document.images[0])[0], document.images[0])

        ocrs = [self.hocr(image) for image in document.images]
        print("OCRS", ocrs)

        return ocrs

    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 bbox, word, conf in ocr_result:
                word_id += 1
                dict_word = {
                    "page": page,
                    "class": "ocrx_word",
                    "id": f"word_{page + 1}_{word_id}",
                    "parent": f"word_{page + 1}_{word_id}",
                    "value": word,
                    "confidence": conf,
                    "x1": round(min([edge[0] for edge in bbox])),
                    "y1": round(min([edge[1] for edge in bbox])),
                    "x2": round(max([edge[0] for edge in bbox])),
                    "y2": round(max([edge[1] for edge in bbox]))
                }

                list_elements.append(dict_word)
        print("LIST EL", list_elements)

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