import time
import cv2
import tkinter as tk
from tkinter import filedialog, ttk, messagebox
from PIL import Image, ImageTk
import numpy as np
from sklearn.metrics import confusion_matrix, precision_recall_curve, average_precision_score
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import pandas as pd
import subprocess
import os
import random
from threading import Thread # 用于训练时不阻塞界面
class ObjectDetectionApp:
def __init__(self):
self.window = tk.Tk()
self.window.title("YOLOv4训练&验证系统")
self.window.geometry("1600x1000")
# 模式变量（训练/验证/原评估）
self.mode = tk.StringVar(value="validation")
self.train_running = False # 训练状态标志
# 模型配置（可自定义）
self.base_cfg = "yolov4.cfg" # 基础配置文件
self.custom_cfg = "yolov4-custom.cfg" # 训练用自定义配置
self.weights_path = "yolov4.weights" # 默认权重/训练后权重
self.classes_file = "coco.names"
self.classes = self.load_classes()
self.COLORS = np.random.uniform(0, 255, size=(len(self.classes), 3))
self.voc_root = None # 新增：VOC数据集根路径
# 训练参数
self.train_params = {
"epochs": 100,
"batch": 8,
"subdivisions": 4,
"learning_rate": 0.001,
"data_file": "custom.data",
"names_file": self.classes_file
}
# 摄像头初始化
self.cap = None
self.image_flipped = True
# 评估相关变量（验证阶段使用）
self.predicted_labels = []
self.ground_truths = []
self.performance_data = {}
# 创建界面
self.create_widgets()
self.window.mainloop()
def create_widgets(self):
# -------------------- 顶部模式栏 --------------------
mode_frame = tk.Frame(self.window, padx=10, pady=5)
mode_frame.pack(fill=tk.X)
tk.Radiobutton(mode_frame, text="模型训练", variable=self.mode, value="training",
command=self.switch_mode).pack(side=tk.LEFT, padx=10)
tk.Radiobutton(mode_frame, text="实时验证", variable=self.mode, value="validation",
command=self.switch_mode).pack(side=tk.LEFT, padx=10)
# -------------------- 训练模块控件 --------------------
self.train_frame = tk.Frame(self.window, padx=10, pady=5)
# VOC数据集选择按钮
self.select_voc_btn = tk.Button(self.train_frame, text="选择VOC数据集",
command=self.select_voc_dataset)
self.select_voc_btn.pack(side=tk.LEFT, padx=10)
# 训练参数输入
param_frame = tk.Frame(self.train_frame)
tk.Label(param_frame, text="迭代次数:").grid(row=0, column=0, padx=5)
self.epochs_entry = ttk.Entry(param_frame, width=8)
self.epochs_entry.grid(row=0, column=1, padx=5)
self.epochs_entry.insert(0, "100")
tk.Label(param_frame, text="批次大小:").grid(row=1, column=0, padx=5)
self.batch_entry = ttk.Entry(param_frame, width=8)
self.batch_entry.grid(row=1, column=1, padx=5)
self.batch_entry.insert(0, "8")
param_frame.pack(side=tk.LEFT, padx=10)
# 数据集加载按钮
self.load_train_btn = tk.Button(self.train_frame, text="加载训练集",
command=self.load_train_dataset)
self.load_train_btn.pack(side=tk.LEFT, padx=10)
self.load_val_btn = tk.Button(self.train_frame, text="加载验证集",
command=self.load_val_dataset)
self.load_val_btn.pack(side=tk.LEFT, padx=10)
# 开始训练按钮（正确顺序：先定义，再布局）
self.start_train_btn = tk.Button(self.train_frame, text="开始训练",
command=self.start_training_thread)
self.start_train_btn.pack(side=tk.LEFT, padx=10)
# -------------------- 主显示区域 --------------------
self.main_frame = tk.Frame(self.window)
self.main_frame.pack(fill=tk.BOTH, expand=True)
# 视频/图像显示区
self.photo_label = tk.Label(self.main_frame, width=1000, height=600)
self.photo_label.pack(side=tk.LEFT, padx=10, pady=10)
# 右侧信息区
self.right_frame = tk.Frame(self.main_frame, width=400)
self.right_frame.pack(side=tk.RIGHT, padx=10, pady=10, fill=tk.Y)
# 日志文本框
self.log_text = tk.Text(self.right_frame, width=40, height=15)
self.log_text.pack(pady=5, fill=tk.X)
self.log_text.insert(tk.END, "系统日志：\n")
# 性能报告区
self.report_canvas = None
self.report_frame = tk.Frame(self.right_frame)
self.report_frame.pack(pady=5, fill=tk.BOTH, expand=True)
# 初始显示验证模式
self.switch_mode()
def select_voc_dataset(self):
"""选择VOC数据集根目录并执行划分"""
self.voc_root = filedialog.askdirectory(title="选择VOC数据集根目录（如VOCdevkit/VOC2007）")
if not self.voc_root:
return
# 检查所选目录下是否存在JPEGImages和Annotations目录
jpeg_dir = os.path.join(self.voc_root, "JPEGImages")
ann_dir = os.path.join(self.voc_root, "Annotations")
if not (os.path.isdir(jpeg_dir) and os.path.isdir(ann_dir)):
messagebox.showerror("错误", "所选目录不是正确的VOC数据集根目录（缺少JPEGImages或Annotations目录）")
return
try:
# 执行VOC数据集划分（75:25）
self.split_voc_dataset(self.voc_root)
self.log_text.insert(tk.END, "数据集划分完成！\n")
# 自动设置训练/验证集路径到参数中
self.train_params["train_images"] = os.path.join(self.voc_root, "ImageSets/Main/train.txt")
self.train_params["val_images"] = os.path.join(self.voc_root, "ImageSets/Main/val.txt")
self.log_text.insert(tk.END, f"训练集路径：{self.train_params['train_images']}\n")
self.log_text.insert(tk.END, f"验证集路径：{self.train_params['val_images']}\n")
except Exception as e:
messagebox.showerror("错误", f"数据集划分失败：{str(e)}")
self.log_text.insert(tk.END, f"划分错误：{str(e)}\n")
def split_voc_dataset(self, voc_root, train_ratio=0.75):
"""VOC数据集划分核心逻辑（集成到应用中）"""
img_dir = os.path.join(voc_root, "JPEGImages")
ann_dir = os.path.join(voc_root, "Annotations")
sets_dir = os.path.join(voc_root, "ImageSets", "Main")
os.makedirs(sets_dir, exist_ok=True)
img_files = [f for f in os.listdir(img_dir) if f.lower().endswith((".jpg", ".jpeg", ".png"))]
valid_ids = []
for img_file in img_files:
img_id = os.path.splitext(img_file)[0]
ann_path = os.path.join(ann_dir, f"{img_id}.xml")
if os.path.isfile(ann_path):
valid_ids.append(img_file) # 保存带扩展名的图像文件名
if not valid_ids:
messagebox.showerror("错误",
"未找到任何有效标注文件，确保JPEGImages和Annotations目录中的文件一一对应（除扩展名外文件名相同）")
return
random.shuffle(valid_ids)
total = len(valid_ids)
train_count = int(total * train_ratio)
train_ids = valid_ids[:train_count]
val_ids = valid_ids[train_count:]
train_txt = os.path.join(sets_dir, "train.txt")
val_txt = os.path.join(sets_dir, "val.txt")
# 写入完整路径（如：D:/pycharm/pythonProject/VOCdevkit/VOC2007/JPEGImages/1.jpg）
with open(train_txt, "w") as f:
for img_file in train_ids:
img_full_path = os.path.join(voc_root, "JPEGImages", img_file)
f.write(img_full_path + '\n')
with open(val_txt, "w") as f:
for img_file in val_ids:
img_full_path = os.path.join(voc_root, "JPEGImages", img_file)
f.write(img_full_path + '\n')
self.log_text.insert(tk.END, f"总样本数: {total}，训练集: {len(train_ids)}，验证集: {len(val_ids)}\n")
def switch_mode(self):
"""模式切换逻辑"""
current_mode = self.mode.get()
# 关闭摄像头
if self.cap:
self.cap.release()
self.cap = None
self.photo_label.config(image=None)
# 清空界面
self.log_text.delete(1.0, tk.END)
self.log_text.insert(tk.END, "系统日志：\n")
if current_mode == "training":
self.train_frame.pack(fill=tk.X, pady=5)
self.report_frame.pack_forget()
self.log_text.insert(tk.END, "切换到训练模式\n")
else: # 验证模式
self.train_frame.pack_forget()
self.report_frame.pack(fill=tk.BOTH, expand=True)
self.cap = cv2.VideoCapture(0) # 打开摄像头
self.update_validation_frame() # 启动实时验证
self.log_text.insert(tk.END, "切换到实时验证模式\n")
def load_train_dataset(self):
"""加载训练数据集（图像+标注）"""
img_dir = filedialog.askdirectory(title="选择训练图像文件夹")
label_dir = filedialog.askdirectory(title="选择训练标注文件夹")
if img_dir and label_dir:
self.train_params["train_images"] = img_dir
self.train_params["train_labels"] = label_dir
self.log_text.insert(tk.END, f"训练集加载完成：{img_dir}\n")
def load_val_dataset(self):
"""加载验证数据集"""
img_dir = filedialog.askdirectory(title="选择验证图像文件夹")
label_dir = filedialog.askdirectory(title="选择验证标注文件夹")
if img_dir and label_dir:
self.train_params["val_images"] = img_dir
self.train_params["val_labels"] = label_dir
self.log_text.insert(tk.END, f"验证集加载完成：{img_dir}\n")
def start_training_thread(self):
"""启动训练线程（防止界面阻塞）"""
if self.train_running:
messagebox.showwarning("提示", "训练已在进行中")
return
self.train_running = True
Thread(target=self.start_training, daemon=True).start()
def start_training(self):
"""执行训练流程（调用Darknet命令）"""
try:
# 生成训练配置文件
self.generate_training_config()
# 训练命令（示例，根据实际Darknet路径调整）
cmd = [
"darknet.exe", "detector", "train",
self.train_params["data_file"],
self.custom_cfg,
"yolov4.conv.137", # 预训练权重
"-map", # 计算mAP
"-gpus", "0",
"-batch", self.batch_entry.get(),
"-subdivisions", "4",
"-epochs", self.epochs_entry.get()
]
self.log_text.insert(tk.END, "开始训练...\n")
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
# 实时输出训练日志
for line in process.stdout:
self.log_text.insert(tk.END, line)
self.log_text.see(tk.END) # 自动滚动
process.wait()
self.train_running = False
self.log_text.insert(tk.END, "训练完成！最佳权重已保存\n")
# 训练完成后生成性能报告
self.generate_training_report()
except Exception as e:
self.log_text.insert(tk.END, f"训练错误：{str(e)}\n")
self.train_running = False
def generate_training_config(self):
"""生成自定义训练配置文件"""
with open(self.base_cfg, 'r') as f:
cfg_lines = f.readlines()
# 修改批次和subdivisions
for i, line in enumerate(cfg_lines):
if line.startswith("batch="):
cfg_lines[i] = f"batch={self.batch_entry.get()}\n"
if line.startswith("subdivisions="):
cfg_lines[i] = "subdivisions=4\n"
with open(self.custom_cfg, 'w') as f:
f.writelines(cfg_lines)
# 生成data文件，指定训练和验证的图像列表路径
data_content = f"""
train = {os.path.join(self.voc_root, "ImageSets/Main/train.txt")} # 训练集图像列表
valid = {os.path.join(self.voc_root, "ImageSets/Main/val.txt")} # 验证集图像列表
names = {self.classes_file}
backup = backup/
eval = coco
"""
with open(self.train_params["data_file"], 'w') as f:
f.write(data_content)
def generate_training_report(self):
"""生成训练性能报告"""
# 假设从训练日志或结果文件中读取数据
# 这里模拟加载验证集结果
y_true = ["car", "person", "car", "bike"]
y_pred = ["car", "person", "bike", "bike"]
confidences = [0.92, 0.85, 0.78, 0.91]
# 计算指标
cm = confusion_matrix(y_true, y_pred, labels=self.classes)
precision, recall, _ = precision_recall_curve(y_true, confidences, pos_label="car")
ap = average_precision_score([1 if x == "car" else 0 for x in y_true], confidences)
# 绘制报告
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
# 混淆矩阵
ax1.imshow(cm, cmap=plt.cm.Blues)
ax1.set_title("验证集混淆矩阵")
ax1.set_xticks(range(len(self.classes)))
ax1.set_xticklabels(self.classes, rotation=45)
ax1.set_yticks(range(len(self.classes)))
ax1.set_yticklabels(self.classes)
# PR曲线
ax2.plot(recall, precision)
ax2.set_title(f"PR曲线 (AP={ap:.2f})")
ax2.set_xlabel("召回率")
ax2.set_ylabel("精确率")
# 在界面显示
if self.report_canvas:
self.report_canvas.get_tk_widget().destroy()
self.report_canvas = FigureCanvasTkAgg(fig, master=self.report_frame)
self.report_canvas.draw()
self.report_canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
def update_validation_frame(self):
"""实时验证摄像头画面"""
if not self.cap or not self.mode.get() == "validation":
return
ret, frame = self.cap.read()
if ret:
if self.image_flipped:
frame = cv2.flip(frame, 1)
# 使用训练后的权重进行检测
frame, results = self.detect_objects(frame)
# 显示结果
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(frame_rgb)
tk_img = ImageTk.PhotoImage(image=pil_img)
self.photo_label.config(image=tk_img)
self.photo_label.image = tk_img
# 记录验证结果（可选）
self.record_validation_results(results)
self.window.after(30, self.update_validation_frame)
def detect_objects(self, frame):
"""使用当前权重进行目标检测"""
# 改为使用实例变量 self.net 存储网络
self.net = cv2.dnn.readNetFromDarknet(
self.custom_cfg if self.mode.get() == "training" else self.base_cfg,
self.weights_path
)
# 使用 self.net 获取层信息
layer_names = self.net.getLayerNames()
output_layers = [layer_names[i - 1] for i in self.net.getUnconnectedOutLayers()]
# 图像预处理
height, width = frame.shape[:2]
blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
# 前向传播（使用 self.net）
self.net.setInput(blob)
outs = self.net.forward(output_layers) # 注意：output_layers 已提前计算
# 后续代码保持不变...
class_ids = []
confidences = []
boxes = []
for out in outs:
for detection in out:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > 0.5: # 置信度阈值
center_x = int(detection[0] * width)
center_y = int(detection[1] * height)
w = int(detection[2] * width)
h = int(detection[3] * height)
# 计算边界框坐标
x = int(center_x - w / 2)
y = int(center_y - h / 2)
boxes.append([x, y, w, h])
confidences.append(float(confidence))
class_ids.append(class_id)
# 非极大值抑制
indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
results = []
for i in range(len(boxes)):
if i in indexes:
x, y, w, h = boxes[i]
label = str(self.classes[class_ids[i]])
confidence = confidences[i]
results.append((label, x, y, x + w, y + h, confidence))
# 绘制边界框和标签
color = self.COLORS[class_ids[i]]
cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
cv2.putText(frame, f"{label}: {confidence:.2f}", (x, y - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
return frame, results
def record_validation_results(self, results):
"""记录验证结果（预测标签和置信度）"""
# 从results中提取预测的标签（假设results格式为：(label, x1, y1, x2, y2, confidence)）
for result in results:
predicted_label = result[0] # 标签
confidence = result[5] # 置信度
self.predicted_labels.append(predicted_label)
# 如果有真实标签（例如从验证集标注文件读取），可添加类似逻辑：
# self.ground_truths.append(ground_truth_label)
def load_classes(self):
with open(self.classes_file, "r") as f:
return [line.strip() for line in f.readlines()]
def close_app(self):
if self.cap:
self.cap.release()
self.window.destroy()
if __name__ == "__main__":
app = ObjectDetectionApp()
代码检测
扫一扫
459
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
