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
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.