Abijith's picture
Upload 5 files
38744b1
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