from pathlib import Path
import gradio as gr
import pymupdf
from ultralytics import YOLO
from PIL import Image
from huggingface_hub import hf_hub_download

SAMPLES = Path(__file__).parent / "samples"

IMAGE_SAMPLES = [
    SAMPLES / "image1.png",
    SAMPLES / "image2.png",
    SAMPLES / "image3.png",
    SAMPLES / "image4.png",
]

AVAILABLE_MODELS = {
    "yolo11n": ("Armaggheddon/yolo11-document-layout", "yolo11n_doc_layout.pt"),
    "yolo11s": ("Armaggheddon/yolo11-document-layout", "yolo11s_doc_layout.pt"),
    "yolo11m": ("Armaggheddon/yolo11-document-layout", "yolo11m_doc_layout.pt"),
}
current_model = "yolo11n"
model = None

def load_model(selected_model):
    global model
    if model is None or current_model != selected_model:
        repo_id, filename = AVAILABLE_MODELS[selected_model]
        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
        model = YOLO(model_path)

def model_runner(image, conf=0.25, iou=0.45):
    result = model.predict(source=image, save=False, verbose=False, conf=conf, iou=iou, imgsz=1280)
    result_img = result[0].plot()
    return result_img

def process_input(selected_model, pdf_input, image_input, conf=0.25, iou=0.45):
    if pdf_input is None and image_input is None:
        return gr.Error("Please upload a PDF or an image file.")
    
    load_model(selected_model)
    pages = []
    if pdf_input is not None and pdf_input.endswith(".pdf"):
        doc = pymupdf.open(pdf_input)
        for page in doc:
            pix = page.get_pixmap(dpi=200) # if A4 should result in above 1400px width
            pil_img = pix.pil_image()
            result_img = model_runner(pil_img)
            pages.append(result_img)
    elif image_input is not None and image_input.endswith((".png", ".jpg", ".jpeg")):
        image = image_input
        result_img = model_runner(image)
        pages.append(result_img)

    else:
        return gr.Error("Unsupported file type. Please upload a PDF or an image file with .pdf, .jpg or .jpeg extension.")

    return ((page, f"Page {i+1}") for i, page in enumerate(pages))

with gr.Blocks() as demo:
    gr.Markdown("# YOLO11 Document Layout 🔎📄")
    gr.Markdown(
"""
Detects layout elements in documents (PDFs or images) using YOLOv11 models and the Ultralytics library.
Upload a PDF or an image, select a model size, and click "Run" to see the detected layout elements.
- Finetuned models available at [Armaggheddon/yolo11-document-layout](https://huggingface.co/Armaggheddon/yolo11-document-layout)
- More available in the [GitHub Repository](https://github.com/Armaggheddon/yolo11_doc_layout)
"""
    )
    with gr.Row():
        with gr.Column():
            pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"], file_count="single")
            image_input = gr.Image(label="Upload Image", type="filepath")
            clear_button = gr.Button("Clear")
            run_button = gr.Button("Run", variant="primary")
        with gr.Column():
            outputs = gr.Gallery(label="Output Image")
            with gr.Group():
                model_name = gr.Dropdown(
                    list(AVAILABLE_MODELS.keys()),
                    value="yolo11n",
                    label="Model size",
                )
                conf = gr.Slider(0, 1, value=0.25, step=0.01, label="Confidence threshold")
                iou = gr.Slider(0, 1, value=0.45, step=0.01, label="IOU threshold")

    examples = gr.Examples(
        examples=[[str(p), "yolo11n"] for p in IMAGE_SAMPLES],
        inputs=[image_input, model_name],
        cache_examples=False,
        fn=process_input,
        outputs=outputs,
    )

    run_button.click(
        fn=process_input,
        inputs=[model_name, pdf_input, image_input, conf, iou],
        outputs=outputs,
    )

    clear_button.click(
        fn=lambda: (None, None, None),
        inputs=[],
        outputs=[pdf_input, image_input, outputs],
    )

demo.launch()