Update app.py

app.py

@@ -4,54 +4,59 @@ from PIL import Image
 from pdf2image import convert_from_path
 from PyPDF2 import PdfReader
 import os
+import torch
 
 # Load the LayoutLMv3 model and processor
 processor = AutoProcessor.from_pretrained("microsoft/layoutlmv3-base")
 model = AutoModel.from_pretrained("microsoft/layoutlmv3-base")
 
-# Helper function to process PDF page
 def process_pdf(pdf, page_number):
-    reader = PdfReader(pdf.name)
-    num_pages = len(reader.pages)
-    if page_number > num_pages or page_number < 1:
-        return f"Invalid page number. The document has {num_pages} pages."
+    try:
+        reader = PdfReader(pdf.name)
+        num_pages = len(reader.pages)
+        if page_number > num_pages or page_number < 1:
+            return f"Invalid page number. The document has {num_pages} pages."
 
-    # Convert PDF page to image
-    images = convert_from_path(pdf.name, first_page=page_number, last_page=page_number)
-    return process_image(images[0])
+        images = convert_from_path(pdf.name, first_page=page_number, last_page=page_number)
+        return process_image(images[0])
+    except Exception as e:
+        return f"Error processing PDF: {str(e)}"
 
-# Helper function to process image (for both direct image input and PDF page conversion)
 def process_image(image):
-    image = image.convert("RGB")
-    # Tokenize and process the image with LayoutLMv3
-    encoded_inputs = processor(image, return_tensors="pt")
-    outputs = model(**encoded_inputs)
-    
-    # Placeholder for further processing (e.g., extracting tokens, detecting tables, etc.)
-    return "Image processed with LayoutLMv3."
+    try:
+        image = image.convert("RGB")
+        encoded_inputs = processor(image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**encoded_inputs)
+        
+        # Placeholder for table extraction logic
+        # You would implement your table detection and extraction here
+        # For now, we'll just return the shape of the last hidden states
+        last_hidden_states = outputs.last_hidden_state
+        return f"Image processed. Output shape: {last_hidden_states.shape}"
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
 
-# Gradio interface function
 def extract_table(file, page_number):
-    file_ext = os.path.splitext(file.name)[1].lower()
-    if file_ext == ".pdf":
-        return process_pdf(file, page_number)
-    elif file_ext in [".png", ".jpg", ".jpeg"]:
-        image = Image.open(file)
-        return process_image(image)
-    else:
-        return "Unsupported file format. Please upload a PDF or an image."
+    try:
+        file_ext = os.path.splitext(file.name)[1].lower()
+        if file_ext == ".pdf":
+            return process_pdf(file, page_number)
+        elif file_ext in [".png", ".jpg", ".jpeg"]:
+            image = Image.open(file)
+            return process_image(image)
+        else:
+            return "Unsupported file format. Please upload a PDF or an image (PNG, JPG, JPEG)."
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Upload a PDF or an image of a table")
-
     with gr.Row():
-        # Specify file types explicitly: PDFs and images (jpg, png)
-        file_input = gr.File(label="Upload PDF or Image", file_types=["pdf", "image"])
+        file_input = gr.File(label="Upload PDF or Image", file_types=["pdf", "png", "jpg", "jpeg"])
         page_input = gr.Number(label="Page number for PDF (ignore if uploading an image)", value=1, precision=0)
-
     output = gr.Textbox(label="Result")
-
     file_input.change(fn=extract_table, inputs=[file_input, page_input], outputs=output)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()