codes/table_recognition.py

import os

from datatypes.datatypes import TableRecognitionData, TableDetectionData
from codes.image_processing import ImageProcessor
from datatypes.config import Config

class TableRecognition:
    def __init__(self, feature_extractor, recognition_model, threshold):
        self.feature_extractor = feature_extractor
        self.recognition_model = recognition_model
        self.threshold = threshold

    def table_recognition_from_detection(self, recognition_image, detection_results):

        for table in detection_results.tables:
            recognised_table_results = TableRecognitionData()
            bbox = table.detection_box
            detected_tbl = recognition_image.crop(bbox)
            img_processor = ImageProcessor()
            padded_table = img_processor.image_padding(image=detected_tbl, padd=Config['table_padd'])
            width, height = padded_table.size

            recognition_encoding = self.feature_extractor(padded_table, return_tensors='pt')
            recognition_output = self.recognition_model(**recognition_encoding)
            recognition_results = self.feature_extractor.post_process_object_detection(recognition_output, threshold=0.7, target_sizes=[(height, width)])
            recognition_results = recognition_results[0]

            recognised_table_results.scores = (recognition_results['scores'].tolist())
            recognised_table_results.labels = (recognition_results['labels'].tolist())
            recognised_table_results.boxes = (recognition_results['boxes'].tolist())

            table.recognitiondata.append(recognised_table_results)
        return detection_results