# 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)
@st.cache_resource
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}"        
        )