Create app.py

app.py

@@ -0,0 +1,54 @@
+import os
+import tempfile
+import fitz  # PyMuPDF
+import gradio as gr
+from ultralytics import YOLO
+from PIL import Image
+
+# 预加载两个模型
+model_doc = YOLO("moldet_yolo11l_960_doc.pt")        # PDF 专用
+model_img = YOLO("moldet_yolo11l_640_general.pt")    # 图片通用
+
+def process_file(file):
+    # 创建临时目录
+    with tempfile.TemporaryDirectory() as tmpdir:
+        input_path = os.path.join(tmpdir, "input.png")
+
+        # 判断文件类型
+        if file.name.lower().endswith(".pdf"):
+            # PDF -> 转第一页为图片
+            pdf = fitz.open(file.name)
+            page = pdf[0]
+            pix = page.get_pixmap(dpi=300)  # 高分辨率渲染
+            pix.save(input_path)
+
+            # 运行 PDF 模型
+            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)
+
+        # YOLO 会保存预测结果到 runs/predict/xxx 目录
+        save_dir = results[0].save_dir
+        pred_path = os.path.join(save_dir, os.path.basename(input_path))
+
+        return pred_path
+
+# Gradio UI
+demo = gr.Interface(
+    fn=process_file,
+    inputs=gr.File(label="上传图片或 PDF（单页）", file_types=[".png", ".jpg", ".jpeg", ".pdf"]),
+    outputs=gr.Image(type="filepath", label="检测结果"),
+    title="PDF / 图片目标检测",
+    description=(
+        "上传一张图片或单页 PDF，系统会自动选择对应的 YOLO 模型进行检测并输出结果。\n"
+        "PDF 使用 moldet_yolo11l_960_doc.pt，图片使用 moldet_yolo11l_640_general.pt"
+    )
+)
+
+if __name__ == "__main__":
+    demo.launch()