Update app.py

app.py

@@ -7,19 +7,46 @@ import numpy as np
 import csv
 import pandas as pd
 
-from transformers import TableTransformerImageProcessor, AutoModelForObjectDetection
+from torchvision import transforms
+
+from transformers import AutoModelForObjectDetection
 import torch
 
 import easyocr
 
 import gradio as gr
 
+
+class MaxResize(object):
+    def __init__(self, max_size=800):
+        self.max_size = max_size
+
+    def __call__(self, image):
+        width, height = image.size
+        current_max_size = max(width, height)
+        scale = self.max_size / current_max_size
+        resized_image = image.resize((int(round(scale*width)), int(round(scale*height))))
+        
+        return resized_image
+
+detection_transform = transforms.Compose([
+    MaxResize(800),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+structure_transform = transforms.Compose([
+    MaxResize(1000),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
 # load table detection model
-processor = TableTransformerImageProcessor(max_size=800)
+# processor = TableTransformerImageProcessor(max_size=800)
 model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-detection", revision="no_timm")
 
 # load table structure recognition model
-structure_processor = TableTransformerImageProcessor(max_size=1000)
+# structure_processor = TableTransformerImageProcessor(max_size=1000)
 structure_model = AutoModelForObjectDetection.from_pretrained("microsoft/table-transformer-structure-recognition-v1.1-all")
 
 # load EasyOCR reader
@@ -117,7 +144,8 @@ def visualize_detected_tables(img, det_tables):
 
 def detect_and_crop_table(image):
     # prepare image for the model
-    pixel_values = processor(image, return_tensors="pt").pixel_values
+    # pixel_values = processor(image, return_tensors="pt").pixel_values
+    pixel_values = detection_transform(image).unsqueeze(0)
 
     # forward pass
     with torch.no_grad():
@@ -140,7 +168,8 @@ def detect_and_crop_table(image):
 
 def recognize_table(image):
     # prepare image for the model
-    pixel_values = structure_processor(images=image, return_tensors="pt").pixel_values
+    # pixel_values = structure_processor(images=image, return_tensors="pt").pixel_values
+    pixel_values = structure_transform(image).unsqueeze(0)
 
     # forward pass
     with torch.no_grad():