Update app.py

app.py

@@ -2,26 +2,72 @@ import gradio as gr
 import torch
 from PIL import Image
 from ultralytics import YOLO
+import matplotlib.pyplot as plt
+import io
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
+plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
+model = YOLO(r'pcb-best.pt')
 
 
-model = YOLO(r'pcb-best.pt')
 def predict(img, conf, iou):
     results = model.predict(img, conf=conf, iou=iou)
+    name = results[0].names
+    cls = results[0].boxes.cls
+    copper = 0
+    mousebite = 0
+    open_defect = 0
+    pin_hole = 0
+    short = 0
+    spur = 0
+    for i in cls:
+        if i == 0:
+            copper += 1
+        elif i == 1:
+            mousebite += 1
+        elif i == 2:
+            open_defect += 1
+        elif i == 3:
+            pin_hole += 1
+        elif i == 4:
+            short += 1
+        elif i == 5:
+            spur += 1
+        # 绘制柱状图
+    fig, ax = plt.subplots()
+    categories = ['Copper', 'Mousebite', 'Open Defect', 'Pin Hole', 'Short', 'Spur']
+    counts = [copper, mousebite, open_defect, pin_hole, short, spur]
+    ax.bar(categories, counts)
+    ax.set_title('缺陷类别计数')
+    plt.ylim(0,5)
+    ax.set_xlabel('缺陷类别')
+    ax.set_ylabel('数目')
+    # 将图表保存为字节流
+    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
+        return im_rgb, image_png
 
 
 base_conf, base_iou = 0.25, 0.45
 title = "基于YOLO-V8的PCB电路板缺陷检测"
 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]])
-
-interface.launch()
+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", '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]])
+interface.launch()