Spaces:
Running
Running
| # Object Detection | |
| import streamlit as st | |
| 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) | |
| def load_assets(): | |
| file_path = hf_hub_download(repo_id="nielsr/example-pdf", repo_type="dataset", filename="example_pdf.png") | |
| image_processor = AutoImageProcessor.from_pretrained("microsoft/table-transformer-detection") | |
| model = TableTransformerForObjectDetection.from_pretrained("microsoft/table-transformer-detection") | |
| return file_path, image_processor, model | |
| file_path, image_processor, model = load_assets() | |
| # App Title | |
| 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 | |
| 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) | |
| # 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] | |
| st.image(image) # Display the uploaded image | |
| for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): | |
| box = [round(i, 2) for i in box.tolist()] | |
| st.write( | |
| f"Detected {model.config.id2label[label.item()]} with confidence " | |
| f"{round(score.item(), 3)} at location {box}" | |
| ) |