PCB-DETECT / app.py
llzzyy233's picture
设置plt中文显示
6c08baf verified
raw
history blame
2.47 kB
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.family'] = '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()