File size: 1,808 Bytes
2baba3c 7099d4e 2baba3c 7099d4e 2baba3c 7099d4e 2baba3c 7099d4e 0687234 2baba3c 0687234 1df5a00 2baba3c 7099d4e 2baba3c 730ccb9 2baba3c 8c016d4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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)
|