import os
import tempfile
import fitz  # PyMuPDF
import gradio as gr
from ultralytics import YOLO
from PIL import Image

# preload detection models
model_doc = YOLO("moldet_yolo11l_960_doc.pt")        # for PDF
model_img = YOLO("moldet_yolo11l_640_general.pt")    # general

def process_file(file):
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "input.png")

        # pdf or image
        if file.name.lower().endswith(".pdf"):
            pdf = fitz.open(file.name)
            page = pdf[0]
            pix = page.get_pixmap(dpi=300)
            pix.save(input_path)

            results = model_doc.predict(input_path, save=True, imgsz=960, conf=0.5)

        else:
            Image.open(file.name).save(input_path)

            results = model_img.predict(input_path, save=True, imgsz=640, conf=0.5)

        # results is a list, our results is in results[0] 
        # save_dir = results[0].save_dir
        # pred_path = os.path.join(save_dir, os.path.basename(input_path))
        result_img = results[0].plot()
        result_pil = Image.fromarray(result_img)

        # return pred_path
        # return result_pil,
        return result_img

# Gradio UI
demo = gr.Interface(
    fn=process_file,
    inputs=gr.File(label="Upload an image or a single PDF file", file_types=[".png", ".jpg", ".jpeg", ".pdf"]),
    outputs=gr.Image(type="pil", label="Detection results"),
    title="Molecule detection",
    description=(
            "Upload an image or a single PDF file to detect molecules using a YOLO-based deep learning model.\n"
            "The system processes the input and returns an annotated image with bounding boxes and labels around the detected molecular structures."
    )
)

if __name__ == "__main__":
    demo.launch(mcp_server=True)