Spaces:
Running
Running
File size: 2,479 Bytes
416724a da0e90e 416724a da0e90e 416724a da0e90e 416724a da0e90e |
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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
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')
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, 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", '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()
|