Abijith/Table-Detection-and-Recognition-using-DETR

1import os2 3from datatypes.datatypes import ImageData4from datatypes.datatypes import TableDetectionData5 6class TableDetection():7    def __init__(self, feature_extractor, detection_model, threshold):8        self.feature_extractor = feature_extractor9        self.detection_model = detection_model10        self.threshold = threshold11 12    def table_detection_from_image(self, detection_image):13 14        table_data_extraction = ImageData([])15        image_width, image_height = detection_image.size16        detection_encoding = self.feature_extractor(detection_image, return_tensors='pt')17        detection_output = self.detection_model(**detection_encoding)18        detection_results = self.feature_extractor.post_process_object_detection(detection_output, threshold=0.3, target_sizes=[(image_height, image_width)])19        detection_results = detection_results[0]20        # copying the detections21        for score, label, bbox in zip((detection_results['scores']).tolist(), (detection_results['labels']).tolist(), (detection_results['boxes']).tolist()):22            detection_table_results = TableDetectionData()23            detection_table_results.detection_score = score24            detection_table_results.detection_label = label25            detection_table_results.detection_box = bbox26            table_data_extraction.tables.append(detection_table_results)27        return table_data_extraction28 29