Update app.py

app.py

@@ -4,6 +4,7 @@ from huggingface_hub import hf_hub_download
 from transformers import AutoImageProcessor, TableTransformerForObjectDetection
 import torch
 from PIL import Image
+import fitz  # Import PyMuPDF (fitz)
 
 # Model and Image Processor Loading (ideally at the app start)
 @st.cache_resource
@@ -15,17 +16,25 @@ def load_assets():
 
 file_path, image_processor, model = load_assets()
 
+
 # App Title
-st.title(" Detection in Images")
+st.title("Table Detection in Documents")
+
+# Document Upload
+uploaded_file = st.file_uploader("Upload a document", type=["pdf", "docx", "doc"])  # Add more formats if needed 
+
+# Process Document and Display Results
+if uploaded_file:
+    doc = fitz.open(stream=uploaded_file.getvalue(), filetype="pdf")  # Open as PDF
 
-# Image Upload
-uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+    for page_index in range(len(doc)):
+        page = doc.load_page(page_index)
+        pix = page.get_pixmap()
+        image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
 
-# Process Image and Display Results
-if uploaded_file:    
-    image = Image.open(uploaded_file).convert("RGB")
-    inputs = image_processor(images=image, return_tensors="pt")    
-    outputs = model(**inputs)
+        # Table Detection (your existing logic)
+        inputs = image_processor(images=image, return_tensors="pt")    
+        outputs = model(**inputs)
 
     target_sizes = torch.tensor([image.size[::-1]])    
     results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]