Update app.py

app.py

@@ -5,66 +5,77 @@ from ultralytics import YOLO
 import matplotlib.pyplot as plt
 import io
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+
+# 修复代码：显式允许 YOLO 的自定义类
+from ultralytics.nn.tasks import DetectionModel
+torch.serialization.add_safe_globals([DetectionModel])
+
+# 加载模型（现在可以安全使用 weights_only=True）
 model = YOLO('detect-best.pt', weights_only=True)
 
 def predict(img, conf, iou):
     results = model.predict(img, conf=conf, iou=iou)
     name = results[0].names
     cls = results[0].boxes.cls
-    crazing = 0
-    inclusion = 0
-    patches = 0
-    pitted_surface = 0
-    rolled_inscale = 0
-    scratches = 0
+    
+    # 初始化计数器
+    counters = {
+        0: 'crazing',
+        1: 'inclusion',
+        2: 'patches',
+        3: 'pitted_surface',
+        4: 'rolled_inscale',
+        5: 'scratches'
+    }
+    counts = {v: 0 for v in counters.values()}
+    
+    # 统计类别
     for i in cls:
-        if i == 0:
-            crazing += 1
-        elif i == 1:
-            inclusion += 1
-        elif i == 2:
-            patches += 1
-        elif i == 3:
-            pitted_surface += 1
-        elif i == 4:
-            rolled_inscale += 1
-        elif i == 5:
-            scratches += 1
-        # 绘制柱状图
-    fig, ax = plt.subplots()
-    categories = ['crazing','inclusion', 'patches' ,'pitted_surface', 'rolled_inscale' ,'scratches']
-    counts = [crazing,inclusion, patches ,pitted_surface, rolled_inscale ,scratches]
-    ax.bar(categories, counts)
-    ax.set_title('Category-Count')
-    plt.ylim(0,5)
+        counts[counters[int(i)]] += 1
+    
+    # 绘制柱状图
+    fig, ax = plt.subplots(figsize=(10, 5))
+    ax.bar(counts.keys(), counts.values())
+    ax.set_title('Defect Category Distribution')
+    ax.set_ylim(0, max(counts.values()) + 1)
     plt.xticks(rotation=45, ha="right")
-    ax.set_xlabel('Category')
-    ax.set_ylabel('Count')
-    # 将图表保存为字节流
+    
+    # 转换为图像
     buf = io.BytesIO()
-    canvas = FigureCanvas(fig)
-    canvas.print_png(buf)
-    plt.close(fig)  # 关闭图形，释放资源
-
-    # 将字节流转换为PIL Image
-    image_png = Image.open(buf)
-    # 绘制并返回结果图片和类别计数图表
-
-    for i, r in enumerate(results):
-        # Plot results image
-        im_bgr = r.plot()  # BGR-order numpy array
-        im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
-
-        # Show results to screen (in supported environments)
-    return im_rgb
+    plt.savefig(buf, format='png', bbox_inches='tight')
+    plt.close()
+    chart_img = Image.open(buf)
+    
+    # 处理检测结果
+    im_bgr = results[0].plot()
+    det_img = Image.fromarray(im_bgr[..., ::-1])
+    
+    # 返回检测结果和统计图
+    return [det_img, chart_img]
 
+# 界面设置
 base_conf, base_iou = 0.25, 0.45
 title = "基于改进YOLOv8算法的工业瑕疵辅助检测系统"
 des = "鼠标点击上传图片即可检测缺陷，可通过鼠标调整预测置信度，还可点击网页最下方示例图片进行预测"
+
 interface = gr.Interface(
-    inputs=['image', gr.Slider(maximum=1, minimum=0, value=base_conf), gr.Slider(maximum=1, minimum=0, value=base_iou)],
-    outputs=["image"], fn=predict, title=title, description=des,
-    examples=[["example1.jpg", base_conf, base_iou],
-              ["example2.jpg", base_conf, base_iou],
-              ["example3.jpg", base_conf, base_iou]])
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="输入图片"),
+        gr.Slider(0, 1, value=base_conf, label="置信度阈值"),
+        gr.Slider(0, 1, value=base_iou, label="IoU阈值")
+    ],
+    outputs=[
+        gr.Image(label="检测结果"),
+        gr.Image(label="缺陷统计")
+    ],
+    title=title,
+    description=des,
+    examples=[
+        ["example1.jpg", base_conf, base_iou],
+        ["example2.jpg", base_conf, base_iou],
+        ["example3.jpg", base_conf, base_iou]
+    ]
+)
+
 interface.launch()