File size: 1,402 Bytes
38744b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
import os
from datatypes.datatypes import ImageData
from datatypes.datatypes import TableDetectionData
class TableDetection():
def __init__(self, feature_extractor, detection_model, threshold):
self.feature_extractor = feature_extractor
self.detection_model = detection_model
self.threshold = threshold
def table_detection_from_image(self, detection_image):
table_data_extraction = ImageData([])
image_width, image_height = detection_image.size
detection_encoding = self.feature_extractor(detection_image, return_tensors='pt')
detection_output = self.detection_model(**detection_encoding)
detection_results = self.feature_extractor.post_process_object_detection(detection_output, threshold=0.3, target_sizes=[(image_height, image_width)])
detection_results = detection_results[0]
# copying the detections
for score, label, bbox in zip((detection_results['scores']).tolist(), (detection_results['labels']).tolist(), (detection_results['boxes']).tolist()):
detection_table_results = TableDetectionData()
detection_table_results.detection_score = score
detection_table_results.detection_label = label
detection_table_results.detection_box = bbox
table_data_extraction.tables.append(detection_table_results)
return table_data_extraction
|