Issue with custom OCR model

[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