更新 app.py

app.py

@@ -1,77 +1,76 @@
-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
-model = YOLO('Fracture_best.pt')
-
-def predict(img, conf, iou):
-    results = model.predict(img, conf=conf, iou=iou)
-    name = results[0].names
-    cls = results[0].boxes.cls
-    boneanomaly = 0
-    bonelesion = 0
-    fracture = 0
-    metal = 0
-    periostealreaction = 0
-    pronatorsign = 0
-    softtissue  = 0
-    text = 0
-    for i in cls:
-        if i == 0:
-            boneanomaly += 1
-        elif i == 1:
-            bonelesion += 1
-        elif i == 2:
-            fracture += 1
-        elif i == 3:
-            metal += 1
-        elif i == 4:
-            periostealreaction += 1
-        elif i == 5:
-            pronatorsign += 1
-        elif i==6:
-            softtissue += 1
-        elif i==7:
-            text += 1 
-        # 绘制柱状图
-    fig, ax = plt.subplots()
-    categories = ['boneanomaly', 'bonelesion', 'fracture', 'metal', 
-                  'periostealreaction', 'pronatorsign', 'softtissue', 'text']
-    counts = [boneanomaly, bonelesion, fracture, metal, periostealreaction, pronatorsign, softtissue, text]
-    ax.bar(categories, counts)
-    ax.set_title('Category-Count')
-    plt.ylim(0,5)
-    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, image_png
-
-
-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", '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()
+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
+model = YOLO('Fracture_best.pt')
+
+def predict(img, conf, iou):
+    results = model.predict(img, conf=conf, iou=iou)
+    name = results[0].names
+    cls = results[0].boxes.cls
+    boneanomaly = 0
+    bonelesion = 0
+    fracture = 0
+    metal = 0
+    periostealreaction = 0
+    pronatorsign = 0
+    softtissue  = 0
+    text = 0
+    for i in cls:
+        if i == 0:
+            boneanomaly += 1
+        elif i == 1:
+            bonelesion += 1
+        elif i == 2:
+            fracture += 1
+        elif i == 3:
+            metal += 1
+        elif i == 4:
+            periostealreaction += 1
+        elif i == 5:
+            pronatorsign += 1
+        elif i==6:
+            softtissue += 1
+        elif i==8:
+            text += 1 
+        # 绘制柱状图
+    fig, ax = plt.subplots()
+    categories = ['BMLY', 'bonelesion', 'fracture', 'metal', 'PN', 'PTS', 'SFTI', 'text']
+    counts = [boneanomaly, bonelesion, fracture, metal, periostealreaction, pronatorsign, softtissue, text]
+    ax.bar(categories, counts)
+    ax.set_title('Category-Count')
+    plt.ylim(0,5)
+    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, image_png
+
+
+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", '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()