基于YOLOv8/YOLOv7/YOLOv6/YOLOv5的草莓成熟度检测系统详解（深度学习模型+UI界面+Python代码+训练数据集）
·
摘要
本文详细介绍基于YOLOv8/YOLOv7/YOLOv6/YOLOv5的草莓成熟度检测系统的完整实现方案。系统通过深度学习技术自动识别草莓的成熟阶段（未成熟、半成熟、成熟），并配备用户友好的UI界面。文章包含系统架构设计、数据集构建、模型训练、优化策略、完整代码实现及部署方案。
目录
一、引言
1.1 研究背景与意义
草莓作为一种高价值水果，其成熟度直接影响果实品质、市场价格和采摘效率。传统的人工检测方法存在效率低、主观性强、劳动强度大等问题。随着计算机视觉和深度学习技术的发展，基于图像的果实成熟度自动检测成为研究热点。
1.2 YOLO系列算法优势
YOLO（You Only Look Once）系列算法作为单阶段目标检测的代表，具有以下优势：
-
实时性高：单次前向传播即可完成检测
-
精度优越：持续优化的网络结构提升检测精度
-
部署方便：支持多种硬件平台部署
-
开源生态：活跃的社区支持和持续更新
1.3 系统创新点
-
多版本YOLO模型对比实验
-
针对草莓特征的优化策略
-
完整的UI界面设计
-
数据增强策略优化
-
实际部署解决方案
二、系统总体设计
2.1 系统架构
text
草莓成熟度检测系统架构： 1. 数据采集模块 2. 数据预处理模块 3. 模型训练模块 4. 推理检测模块 5. UI界面模块 6. 结果分析模块
2.2 技术路线
python
# 技术路线示意图 1. 数据集收集与标注 2. 环境配置与依赖安装 3. 模型选择与配置 4. 训练与优化 5. 评估与测试 6. 界面开发 7. 系统集成
三、数据集构建与处理
3.1 数据集来源
3.1.1 公开数据集
-
Kaggle草莓数据集：包含不同成熟阶段的草莓图像
-
Roboflow公共数据集：多种光照条件下的草莓图像
-
自制数据集：实地采集的草莓图像
3.2.2 数据集结构
text
datasets/ ├── strawberry/ │ ├── images/ │ │ ├── train/ │ │ ├── val/ │ │ └── test/ │ └── labels/ │ ├── train/ │ ├── val/ │ └── test/ ├── data.yaml └── dataset_info.yaml
3.2 数据标注规范
使用LabelImg或CVAT工具进行标注，类别定义：
-
class 0: unripe (未成熟)
-
class 1: semi-ripe (半成熟)
-
class 2: ripe (成熟)
3.3 数据增强策略
python
# 数据增强配置示例
augmentations = {
'hsv_h': 0.015, # 色调增强
'hsv_s': 0.7, # 饱和度增强
'hsv_v': 0.4, # 明度增强
'rotation': 15, # 旋转角度
'translate': 0.2,# 平移
'scale': 0.5, # 缩放
'shear': 0.2, # 剪切
'fliplr': 0.5, # 水平翻转
'mosaic': 1.0, # Mosaic增强
'mixup': 0.2 # Mixup增强
}
四、环境配置与安装
4.1 硬件要求
-
GPU: NVIDIA GTX 1060 6GB或更高
-
RAM: 8GB或更高
-
存储: 50GB可用空间
4.2 软件环境
python
# requirements.txt torch>=1.7.0 torchvision>=0.8.0 ultralytics>=8.0.0 opencv-python>=4.5.0 pillow>=8.3.0 numpy>=1.19.5 pandas>=1.3.0 matplotlib>=3.4.0 seaborn>=0.11.0 pyqt5>=5.15.0 gradio>=3.0.0 streamlit>=1.0.0
4.3 安装步骤
bash
# 1. 创建虚拟环境 conda create -n strawberry-detection python=3.8 conda activate strawberry-detection # 2. 安装PyTorch pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 3. 安装YOLOv8 pip install ultralytics # 4. 安装其他依赖 pip install -r requirements.txt # 5. 验证安装 python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"
五、模型训练与优化
5.1 YOLOv8模型配置
python
# yolov8_strawberry.yaml # 完整模型配置文件 path: ./datasets/strawberry train: images/train val: images/val test: images/test # 类别数量 nc: 3 # 类别名称 names: ['unripe', 'semi-ripe', 'ripe'] # 模型参数 model: type: yolov8n # 可选择 yolov8n/s/m/l/x nc: 3 depth_multiple: 1.0 width_multiple: 1.0 # 训练参数 training: epochs: 100 batch_size: 16 imgsz: 640 workers: 4 device: 0 # GPU设备 optimizer: AdamW lr0: 0.001 lrf: 0.01 momentum: 0.937 weight_decay: 0.0005 warmup_epochs: 3.0 warmup_momentum: 0.8 warmup_bias_lr: 0.1 # 数据增强 augmentation: hsv_h: 0.015 hsv_s: 0.7 hsv_v: 0.4 degrees: 0.0 translate: 0.1 scale: 0.5 shear: 0.0 perspective: 0.0 flipud: 0.0 fliplr: 0.5 mosaic: 1.0 mixup: 0.0
5.2 训练代码实现
python
# train_strawberry.py
"""
草莓成熟度检测模型训练脚本
支持YOLOv5/v6/v7/v8
"""
import os
import sys
import argparse
import yaml
from pathlib import Path
import torch
from ultralytics import YOLO
def setup_training_config(args):
"""设置训练配置"""
config = {
'data': args.data,
'epochs': args.epochs,
'batch_size': args.batch_size,
'imgsz': args.img_size,
'device': args.device,
'workers': args.workers,
'patience': args.patience,
'save': True,
'save_period': args.save_period,
'cache': args.cache,
'name': args.name,
'exist_ok': args.exist_ok,
'pretrained': args.pretrained,
'optimizer': args.optimizer,
'lr0': args.lr0,
'lrf': args.lrf,
'momentum': args.momentum,
'weight_decay': args.weight_decay,
'warmup_epochs': args.warmup_epochs,
'warmup_momentum': args.warmup_momentum,
'box': args.box_loss_weight,
'cls': args.cls_loss_weight,
'dfl': args.dfl_loss_weight,
'hsv_h': args.hsv_h,
'hsv_s': args.hsv_s,
'hsv_v': args.hsv_v,
'degrees': args.rotate,
'translate': args.translate,
'scale': args.scale,
'shear': args.shear,
'perspective': args.perspective,
'flipud': args.flipud,
'fliplr': args.fliplr,
'mosaic': args.mosaic,
'mixup': args.mixup,
'copy_paste': args.copy_paste
}
return config
def train_yolov8(args):
"""训练YOLOv8模型"""
print(f"🚀 开始训练YOLOv8模型...")
print(f"📊 数据集: {args.data}")
print(f"⚙️ 模型类型: {args.model}")
# 加载模型
if args.pretrained:
model = YOLO(args.model)
else:
model = YOLO(args.model, task='detect')
# 训练模型
results = model.train(
data=args.data,
epochs=args.epochs,
batch=args.batch_size,
imgsz=args.img_size,
device=args.device,
workers=args.workers,
patience=args.patience,
project=args.project,
name=args.name,
exist_ok=args.exist_ok,
pretrained=args.pretrained,
optimizer=args.optimizer,
lr0=args.lr0,
lrf=args.lrf,
momentum=args.momentum,
weight_decay=args.weight_decay,
warmup_epochs=args.warmup_epochs,
warmup_momentum=args.warmup_momentum,
box=args.box_loss_weight,
cls=args.cls_loss_weight,
dfl=args.dfl_loss_weight,
hsv_h=args.hsv_h,
hsv_s=args.hsv_s,
hsv_v=args.hsv_v,
degrees=args.rotate,
translate=args.translate,
scale=args.scale,
shear=args.shear,
perspective=args.perspective,
flipud=args.flipud,
fliplr=args.fliplr,
mosaic=args.mosaic,
mixup=args.mixup,
copy_paste=args.copy_paste,
save=True,
save_period=args.save_period,
cache=args.cache,
verbose=True
)
return results
def train_yolov5(args):
"""训练YOLOv5模型"""
print(f"🚀 开始训练YOLOv5模型...")
# 导入YOLOv5
sys.path.append('./yolov5') # 假设yolov5代码在本地
import train as yolov5_train
# 设置训练参数
train_args = argparse.Namespace(
weights=args.model if args.pretrained else '',
cfg='' if args.pretrained else args.model,
data=args.data,
epochs=args.epochs,
batch_size=args.batch_size,
imgsz=args.img_size,
device=args.device,
workers=args.workers,
project=args.project,
name=args.name,
exist_ok=args.exist_ok,
optimizer=args.optimizer,
lr0=args.lr0,
lrf=args.lrf,
momentum=args.momentum,
weight_decay=args.weight_decay,
warmup_epochs=args.warmup_epochs,
warmup_momentum=args.warmup_momentum,
box=args.box_loss_weight,
cls=args.cls_loss_weight,
dfl=args.dfl_loss_weight,
save_period=args.save_period,
cache=args.cache
)
# 开始训练
yolov5_train.run(train_args)
def parse_args():
"""解析命令行参数"""
parser = argparse.ArgumentParser(description='草莓成熟度检测模型训练')
# 基本参数
parser.add_argument('--model-type', type=str, default='yolov8',
choices=['yolov5', 'yolov6', 'yolov7', 'yolov8'],
help='选择YOLO版本')
parser.add_argument('--model', type=str, default='yolov8n.pt',
help='模型路径或预训练模型名称')
parser.add_argument('--data', type=str, default='datasets/strawberry/data.yaml',
help='数据集配置文件路径')
parser.add_argument('--epochs', type=int, default=100,
help='训练轮数')
parser.add_argument('--batch-size', type=int, default=16,
help='批次大小')
parser.add_argument('--img-size', type=int, default=640,
help='输入图像大小')
# 设备参数
parser.add_argument('--device', type=str, default='0',
help='训练设备: cpu, 0, 0,1,2,3等')
parser.add_argument('--workers', type=int, default=4,
help='数据加载线程数')
# 优化器参数
parser.add_argument('--optimizer', type=str, default='AdamW',
choices=['SGD', 'Adam', 'AdamW', 'RMSprop'],
help='优化器选择')
parser.add_argument('--lr0', type=float, default=0.001,
help='初始学习率')
parser.add_argument('--lrf', type=float, default=0.01,
help='最终学习率因子')
parser.add_argument('--momentum', type=float, default=0.937,
help='动量')
parser.add_argument('--weight-decay', type=float, default=0.0005,
help='权重衰减')
# 损失权重
parser.add_argument('--box-loss-weight', type=float, default=7.5,
help='边界框损失权重')
parser.add_argument('--cls-loss-weight', type=float, default=0.5,
help='分类损失权重')
parser.add_argument('--dfl-loss-weight', type=float, default=1.5,
help='DFL损失权重')
# 数据增强参数
parser.add_argument('--hsv-h', type=float, default=0.015,
help='HSV色调增强')
parser.add_argument('--hsv-s', type=float, default=0.7,
help='HSV饱和度增强')
parser.add_argument('--hsv-v', type=float, default=0.4,
help='HSV明度增强')
parser.add_argument('--rotate', type=float, default=0.0,
help='旋转角度范围')
parser.add_argument('--translate', type=float, default=0.1,
help='平移范围')
parser.add_argument('--scale', type=float, default=0.5,
help='缩放范围')
parser.add_argument('--shear', type=float, default=0.0,
help='剪切范围')
parser.add_argument('--perspective', type=float, default=0.0,
help='透视变换')
parser.add_argument('--flipud', type=float, default=0.0,
help='垂直翻转概率')
parser.add_argument('--fliplr', type=float, default=0.5,
help='水平翻转概率')
parser.add_argument('--mosaic', type=float, default=1.0,
help='Mosaic增强概率')
parser.add_argument('--mixup', type=float, default=0.0,
help='Mixup增强概率')
parser.add_argument('--copy-paste', type=float, default=0.0,
help='复制粘贴增强概率')
# 训练控制
parser.add_argument('--warmup-epochs', type=float, default=3.0,
help='热身轮数')
parser.add_argument('--warmup-momentum', type=float, default=0.8,
help='热身动量')
parser.add_argument('--patience', type=int, default=50,
help='早停耐心值')
parser.add_argument('--save-period', type=int, default=-1,
help='保存周期')
parser.add_argument('--cache', type=str, default=False,
choices=['ram', 'disk', False],
help='缓存策略')
# 输出参数
parser.add_argument('--project', type=str, default='runs/train',
help='项目保存路径')
parser.add_argument('--name', type=str, default='strawberry_detection',
help='实验名称')
parser.add_argument('--exist-ok', action='store_true',
help='是否覆盖已存在的实验')
parser.add_argument('--pretrained', action='store_true', default=True,
help='是否使用预训练权重')
return parser.parse_args()
def main():
"""主函数"""
args = parse_args()
# 创建输出目录
os.makedirs(args.project, exist_ok=True)
# 根据模型类型选择训练函数
if args.model_type == 'yolov8':
results = train_yolov8(args)
elif args.model_type == 'yolov5':
results = train_yolov5(args)
elif args.model_type == 'yolov7':
# YOLOv7训练代码类似
print("YOLOv7训练需要单独的代码库")
# train_yolov7(args)
elif args.model_type == 'yolov6':
# YOLOv6训练代码类似
print("YOLOv6训练需要单独的代码库")
# train_yolov6(args)
print("✅ 训练完成!")
# 输出训练结果
if 'results' in locals():
print(f"📈 最佳mAP50: {results.best_map:.4f}")
print(f"📈 最佳mAP50-95: {results.best_map50_95:.4f}")
return 0
if __name__ == '__main__':
main()
5.3 模型评估与分析
python
# evaluate_model.py
"""
模型评估与分析脚本
"""
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
from ultralytics import YOLO
import cv2
import json
from sklearn.metrics import confusion_matrix, classification_report
import warnings
warnings.filterwarnings('ignore')
class StrawberryEvaluator:
"""草莓成熟度检测模型评估器"""
def __init__(self, model_path, data_yaml, device='cuda'):
"""
初始化评估器
Args:
model_path: 模型权重路径
data_yaml: 数据集配置文件路径
device: 运行设备
"""
self.model_path = model_path
self.data_yaml = data_yaml
self.device = device if torch.cuda.is_available() else 'cpu'
# 加载模型
self.model = YOLO(model_path)
self.model.to(self.device)
# 加载数据集配置
with open(data_yaml, 'r') as f:
self.data_config = yaml.safe_load(f)
# 类别信息
self.class_names = self.data_config['names']
self.num_classes = len(self.class_names)
# 存储评估结果
self.results = {}
self.predictions = []
self.ground_truths = []
def evaluate(self, split='val'):
"""
评估模型性能
Args:
split: 评估数据集划分 (val/test)
"""
print(f"📊 开始评估模型: {self.model_path}")
print(f"📁 数据集: {split}")
# 使用YOLO内置评估
metrics = self.model.val(
data=self.data_yaml,
split=split,
device=self.device,
conf=0.25,
iou=0.45,
verbose=True
)
# 保存评估结果
self.results = {
'mAP50': metrics.box.map50,
'mAP50-95': metrics.box.map,
'precision': metrics.box.p,
'recall': metrics.box.r,
'f1_score': 2 * (metrics.box.p * metrics.box.r) / (metrics.box.p + metrics.box.r + 1e-16)
}
# 打印结果
print("\n" + "="*50)
print("📈 评估结果汇总:")
print("="*50)
print(f"✅ mAP@0.5: {self.results['mAP50']:.4f}")
print(f"✅ mAP@0.5:0.95: {self.results['mAP50-95']:.4f}")
print(f"✅ 精确率: {self.results['precision']:.4f}")
print(f"✅ 召回率: {self.results['recall']:.4f}")
print(f"✅ F1分数: {self.results['f1_score']:.4f}")
print("="*50)
return self.results
def collect_predictions(self, split='val', max_samples=100):
"""
收集预测结果用于详细分析
Args:
split: 数据集划分
max_samples: 最大样本数
"""
print("📝 收集预测结果...")
# 获取数据集路径
data_dir = Path(self.data_config['path'])
images_dir = data_dir / 'images' / split
labels_dir = data_dir / 'labels' / split
# 获取图像列表
image_files = list(images_dir.glob('*.jpg')) + list(images_dir.glob('*.png'))
image_files = image_files[:max_samples]
self.predictions = []
self.ground_truths = []
for img_path in image_files:
# 加载图像
img = cv2.imread(str(img_path))
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 预测
results = self.model(img_rgb, conf=0.25, iou=0.45)
# 解析预测结果
pred_boxes = []
pred_classes = []
pred_scores = []
if results[0].boxes is not None:
boxes = results[0].boxes.xyxy.cpu().numpy()
classes = results[0].boxes.cls.cpu().numpy()
scores = results[0].boxes.conf.cpu().numpy()
for box, cls, score in zip(boxes, classes, scores):
pred_boxes.append(box.tolist())
pred_classes.append(int(cls))
pred_scores.append(float(score))
# 加载真实标注
label_path = labels_dir / f"{img_path.stem}.txt"
gt_boxes = []
gt_classes = []
if label_path.exists():
with open(label_path, 'r') as f:
for line in f.readlines():
parts = line.strip().split()
if len(parts) == 5:
cls_id = int(parts[0])
x_center = float(parts[1])
y_center = float(parts[2])
width = float(parts[3])
height = float(parts[4])
# 转换为xyxy格式
x1 = (x_center - width/2) * img.shape[1]
y1 = (y_center - height/2) * img.shape[0]
x2 = (x_center + width/2) * img.shape[1]
y2 = (y_center + height/2) * img.shape[0]
gt_boxes.append([x1, y1, x2, y2])
gt_classes.append(cls_id)
# 保存结果
self.predictions.append({
'image_path': str(img_path),
'pred_boxes': pred_boxes,
'pred_classes': pred_classes,
'pred_scores': pred_scores
})
self.ground_truths.append({
'image_path': str(img_path),
'gt_boxes': gt_boxes,
'gt_classes': gt_classes
})
print(f"✅ 已收集 {len(self.predictions)} 个样本的预测结果")
def plot_confusion_matrix(self, save_path='confusion_matrix.png'):
"""绘制混淆矩阵"""
print("📊 绘制混淆矩阵...")
# 收集所有预测和真实标签
all_preds = []
all_gts = []
for pred, gt in zip(self.predictions, self.ground_truths):
all_preds.extend(pred['pred_classes'])
all_gts.extend(gt['gt_classes'])
if not all_preds or not all_gts:
print("⚠️ 没有足够的数据绘制混淆矩阵")
return
# 计算混淆矩阵
cm = confusion_matrix(all_gts, all_preds, labels=range(self.num_classes))
# 绘制
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
xticklabels=self.class_names,
yticklabels=self.class_names)
plt.title('草莓成熟度检测混淆矩阵')
plt.xlabel('预测标签')
plt.ylabel('真实标签')
plt.tight_layout()
# 保存图像
plt.savefig(save_path, dpi=300, bbox_inches='tight')
plt.show()
print(f"✅ 混淆矩阵已保存至: {save_path}")
return cm
def plot_pr_curve(self, save_path='pr_curve.png'):
"""绘制PR曲线"""
print("📈 绘制PR曲线...")
# 这里简化实现，实际应用中可以使用更精确的PR曲线计算
from sklearn.metrics import precision_recall_curve
# 收集所有预测分数和真实标签
all_scores = []
all_gts_binary = []
for pred, gt in zip(self.predictions, self.ground_truths):
for p_cls, p_score in zip(pred['pred_classes'], pred['pred_scores']):
all_scores.append(p_score)
# 简化：假设我们只关心是否检测到目标
all_gts_binary.append(1 if len(gt['gt_classes']) > 0 else 0)
if not all_scores:
print("⚠️ 没有足够的数据绘制PR曲线")
return
# 计算PR曲线
precision, recall, thresholds = precision_recall_curve(all_gts_binary, all_scores)
# 绘制
plt.figure(figsize=(10, 6))
plt.plot(recall, precision, 'b-', linewidth=2)
plt.fill_between(recall, precision, alpha=0.2, color='blue')
plt.xlabel('召回率 (Recall)')
plt.ylabel('精确率 (Precision)')
plt.title('精确率-召回率曲线 (PR Curve)')
plt.grid(True, alpha=0.3)
plt.xlim([0, 1])
plt.ylim([0, 1])
# 计算AP
ap = np.trapz(precision, recall)
plt.text(0.6, 0.1, f'AP = {ap:.3f}', fontsize=12,
bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
plt.tight_layout()
plt.savefig(save_path, dpi=300, bbox_inches='tight')
plt.show()
print(f"✅ PR曲线已保存至: {save_path}")
return ap
def analyze_class_performance(self):
"""分析每个类别的性能"""
print("🔍 分析各类别性能...")
# 收集每个类别的统计信息
class_stats = {i: {'tp': 0, 'fp': 0, 'fn': 0} for i in range(self.num_classes)}
# 简化的性能分析
# 在实际应用中，需要实现更精确的匹配算法
for pred, gt in zip(self.predictions, self.ground_truths):
# 这里简化处理，实际需要实现IoU匹配
gt_classes_set = set(gt['gt_classes'])
pred_classes_set = set(pred['pred_classes'])
for cls_id in range(self.num_classes):
if cls_id in gt_classes_set and cls_id in pred_classes_set:
class_stats[cls_id]['tp'] += 1
elif cls_id in gt_classes_set and cls_id not in pred_classes_set:
class_stats[cls_id]['fn'] += 1
elif cls_id not in gt_classes_set and cls_id in pred_classes_set:
class_stats[cls_id]['fp'] += 1
# 计算每个类别的指标
performance_data = []
for cls_id, stats in class_stats.items():
tp = stats['tp']
fp = stats['fp']
fn = stats['fn']
precision = tp / (tp + fp + 1e-16)
recall = tp / (tp + fn + 1e-16)
f1 = 2 * precision * recall / (precision + recall + 1e-16)
performance_data.append({
'类别': self.class_names[cls_id],
'精确率': precision,
'召回率': recall,
'F1分数': f1,
'TP': tp,
'FP': fp,
'FN': fn
})
# 创建DataFrame
df_performance = pd.DataFrame(performance_data)
print("\n" + "="*60)
print("📊 各类别性能分析:")
print("="*60)
print(df_performance.to_string(index=False))
print("="*60)
# 可视化
plt.figure(figsize=(12, 6))
# 绘制柱状图
ax = df_performance[['类别', '精确率', '召回率', 'F1分数']].plot(
x='类别', kind='bar', figsize=(12, 6))
plt.title('草莓成熟度检测各类别性能对比')
plt.ylabel('分数')
plt.ylim([0, 1])
plt.legend(loc='lower right')
plt.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.savefig('class_performance.png', dpi=300, bbox_inches='tight')
plt.show()
return df_performance
def generate_report(self, output_dir='evaluation_report'):
"""生成完整评估报告"""
print("📄 生成评估报告...")
# 创建输出目录
output_path = Path(output_dir)
output_path.mkdir(exist_ok=True)
# 执行所有评估
self.evaluate()
self.collect_predictions()
# 生成各种图表
cm = self.plot_confusion_matrix(output_path / 'confusion_matrix.png')
ap = self.plot_pr_curve(output_path / 'pr_curve.png')
df_performance = self.analyze_class_performance()
# 生成HTML报告
report_html = f"""
<!DOCTYPE html>
<html>
<head>
<title>草莓成熟度检测模型评估报告</title>
<style>
body {{ font-family: Arial, sans-serif; margin: 40px; }}
.header {{ background-color: #4CAF50; color: white; padding: 20px; text-align: center; }}
.section {{ margin: 30px 0; padding: 20px; border: 1px solid #ddd; border-radius: 5px; }}
.metrics {{ display: grid; grid-template-columns: repeat(2, 1fr); gap: 20px; }}
.metric-card {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
h2 {{ color: #4CAF50; }}
table {{ width: 100%; border-collapse: collapse; }}
th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }}
th {{ background-color: #4CAF50; color: white; }}
img {{ max-width: 100%; height: auto; }}
</style>
</head>
<body>
<div class="header">
<h1>🍓 草莓成熟度检测模型评估报告</h1>
<p>生成时间: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
</div>
<div class="section">
<h2>📊 总体性能指标</h2>
<div class="metrics">
<div class="metric-card">
<h3>mAP@0.5</h3>
<p style="font-size: 24px; font-weight: bold;">{self.results.get('mAP50', 0):.4f}</p>
</div>
<div class="metric-card">
<h3>mAP@0.5:0.95</h3>
<p style="font-size: 24px; font-weight: bold;">{self.results.get('mAP50-95', 0):.4f}</p>
</div>
<div class="metric-card">
<h3>精确率</h3>
<p style="font-size: 24px; font-weight: bold;">{self.results.get('precision', 0):.4f}</p>
</div>
<div class="metric-card">
<h3>召回率</h3>
<p style="font-size: 24px; font-weight: bold;">{self.results.get('recall', 0):.4f}</p>
</div>
</div>
</div>
<div class="section">
<h2>📈 可视化分析</h2>
<div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 20px;">
<div>
<h3>混淆矩阵</h3>
<img src="confusion_matrix.png" alt="混淆矩阵">
</div>
<div>
<h3>PR曲线</h3>
<img src="pr_curve.png" alt="PR曲线">
</div>
</div>
</div>
<div class="section">
<h2>📋 详细性能分析</h2>
{df_performance.to_html(index=False) if df_performance is not None else ''}
</div>
<div class="section">
<h2>🔧 模型信息</h2>
<p><strong>模型路径:</strong> {self.model_path}</p>
<p><strong>数据集:</strong> {self.data_yaml}</p>
<p><strong>设备:</strong> {self.device}</p>
<p><strong>类别:</strong> {', '.join(self.class_names)}</p>
</div>
</body>
</html>
"""
# 保存HTML报告
report_path = output_path / 'evaluation_report.html'
with open(report_path, 'w', encoding='utf-8') as f:
f.write(report_html)
# 保存JSON格式的详细结果
detailed_results = {
'model_info': {
'path': self.model_path,
'classes': self.class_names
},
'overall_metrics': self.results,
'class_performance': df_performance.to_dict('records') if df_performance is not None else [],
'confusion_matrix': cm.tolist() if cm is not None else [],
'average_precision': ap if 'ap' in locals() else None
}
with open(output_path / 'detailed_results.json', 'w') as f:
json.dump(detailed_results, f, indent=2, ensure_ascii=False)
print(f"✅ 评估报告已生成至: {output_path}/")
print(f"📄 HTML报告: {report_path}")
print(f"📊 详细结果: {output_path}/detailed_results.json")
return output_path
# 使用示例
if __name__ == '__main__':
# 配置参数
MODEL_PATH = 'runs/train/strawberry_detection/weights/best.pt'
DATA_YAML = 'datasets/strawberry/data.yaml'
# 创建评估器
evaluator = StrawberryEvaluator(MODEL_PATH, DATA_YAML)
# 生成完整报告
report_dir = evaluator.generate_report('evaluation_results')
print("🎉 评估完成!")
5.4 模型优化策略
python
# 模型优化技术
optimization_techniques = {
'知识蒸馏': '使用大模型指导小模型训练',
'剪枝': '移除不重要的网络连接',
'量化': '降低模型精度减少计算量',
'数据增强优化': '针对草莓特征设计增强策略',
'迁移学习': '使用在ImageNet上预训练的权重',
'多尺度训练': '提高模型对不同大小目标的检测能力',
'标签平滑': '减少过拟合',
'混合精度训练': '加速训练过程'
}
六、UI界面设计与实现
6.1 PyQt5桌面应用
python
# strawberry_detection_ui.py
"""
草莓成熟度检测系统 - PyQt5 UI界面
"""
import sys
import os
from pathlib import Path
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
QHBoxLayout, QPushButton, QLabel, QFileDialog,
QMessageBox, QGroupBox, QComboBox, QSlider,
QProgressBar, QTextEdit, QTabWidget, QTableWidget,
QTableWidgetItem, QHeaderView, QSplitter)
from PyQt5.QtCore import Qt, QTimer, pyqtSignal, QThread, QSize
from PyQt5.QtGui import QImage, QPixmap, QFont, QPalette, QColor
import torch
from ultralytics import YOLO
import pandas as pd
from datetime import datetime
import json
class DetectionThread(QThread):
"""检测线程"""
detection_finished = pyqtSignal(list, np.ndarray)
progress_updated = pyqtSignal(int)
def __init__(self):
super().__init__()
self.model = None
self.image = None
self.conf_threshold = 0.25
self.iou_threshold = 0.45
self.is_running = False
def setup(self, model_path, image, conf_thresh=0.25, iou_thresh=0.45):
"""设置检测参数"""
self.model = YOLO(model_path)
self.image = image
self.conf_threshold = conf_thresh
self.iou_threshold = iou_thresh
def run(self):
"""运行检测"""
if self.image is None or self.model is None:
return
self.is_running = True
# 执行检测
results = self.model(self.image,
conf=self.conf_threshold,
iou=self.iou_threshold,
verbose=False)
# 解析结果
detections = []
result_image = self.image.copy()
if results[0].boxes is not None:
boxes = results[0].boxes.xyxy.cpu().numpy()
classes = results[0].boxes.cls.cpu().numpy()
scores = results[0].boxes.conf.cpu().numpy()
# 在图像上绘制结果
colors = [(0, 255, 0), (0, 165, 255), (0, 0, 255)] # 绿，橙，红
for box, cls, score in zip(boxes, classes, scores):
x1, y1, x2, y2 = map(int, box)
class_id = int(cls)
confidence = float(score)
# 绘制边界框
color = colors[class_id % len(colors)]
cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)
# 绘制标签
label = f"{['未成熟', '半成熟', '成熟'][class_id]}: {confidence:.2f}"
cv2.putText(result_image, label, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
# 添加到检测结果列表
detections.append({
'class_id': class_id,
'class_name': ['未成熟', '半成熟', '成熟'][class_id],
'confidence': confidence,
'bbox': [x1, y1, x2, y2],
'area': (x2 - x1) * (y2 - y1)
})
self.detection_finished.emit(detections, result_image)
self.is_running = False
def stop(self):
"""停止检测"""
self.is_running = False
class StrawberryDetectionUI(QMainWindow):
"""草莓成熟度检测主界面"""
def __init__(self):
super().__init__()
self.model = None
self.current_image = None
self.detection_thread = None
self.detection_history = []
self.init_ui()
self.load_default_model()
def init_ui(self):
"""初始化用户界面"""
self.setWindowTitle('🍓 草莓成熟度智能检测系统 v1.0')
self.setGeometry(100, 100, 1400, 900)
# 设置主窗口部件
central_widget = QWidget()
self.setCentralWidget(central_widget)
# 主布局
main_layout = QHBoxLayout(central_widget)
# 左侧控制面板
control_panel = self.create_control_panel()
main_layout.addWidget(control_panel, 1)
# 右侧主显示区域
display_panel = self.create_display_panel()
main_layout.addWidget(display_panel, 3)
# 设置样式
self.set_style()
def create_control_panel(self):
"""创建控制面板"""
control_panel = QGroupBox("控制面板")
control_layout = QVBoxLayout()
# 模型选择
model_group = QGroupBox("模型设置")
model_layout = QVBoxLayout()
self.model_combo = QComboBox()
self.model_combo.addItems(['YOLOv8n', 'YOLOv8s', 'YOLOv8m', 'YOLOv8l', 'YOLOv8x'])
self.model_combo.currentTextChanged.connect(self.on_model_changed)
self.btn_load_model = QPushButton("加载自定义模型")
self.btn_load_model.clicked.connect(self.load_model)
model_layout.addWidget(QLabel("选择模型:"))
model_layout.addWidget(self.model_combo)
model_layout.addWidget(self.btn_load_model)
model_group.setLayout(model_layout)
# 检测参数
param_group = QGroupBox("检测参数")
param_layout = QVBoxLayout()
# 置信度阈值
self.conf_label = QLabel("置信度阈值: 0.25")
self.conf_slider = QSlider(Qt.Horizontal)
self.conf_slider.setRange(1, 99)
self.conf_slider.setValue(25)
self.conf_slider.valueChanged.connect(self.update_conf_threshold)
# IoU阈值
self.iou_label = QLabel("IoU阈值: 0.45")
self.iou_slider = QSlider(Qt.Horizontal)
self.iou_slider.setRange(10, 90)
self.iou_slider.setValue(45)
self.iou_slider.valueChanged.connect(self.update_iou_threshold)
param_layout.addWidget(self.conf_label)
param_layout.addWidget(self.conf_slider)
param_layout.addWidget(self.iou_label)
param_layout.addWidget(self.iou_slider)
param_group.setLayout(param_layout)
# 操作按钮
btn_group = QGroupBox("操作")
btn_layout = QVBoxLayout()
self.btn_load_image = QPushButton("📷 加载图像")
self.btn_load_image.clicked.connect(self.load_image)
self.btn_load_folder = QPushButton("📁 加载文件夹")
self.btn_load_folder.clicked.connect(self.load_folder)
self.btn_camera = QPushButton("🎥 摄像头捕获")
self.btn_camera.clicked.connect(self.start_camera)
self.btn_detect = QPushButton("🔍 开始检测")
self.btn_detect.clicked.connect(self.start_detection)
self.btn_detect.setEnabled(False)
self.btn_export = QPushButton("💾 导出结果")
self.btn_export.clicked.connect(self.export_results)
self.btn_clear = QPushButton("🗑️ 清除")
self.btn_clear.clicked.connect(self.clear_results)
btn_layout.addWidget(self.btn_load_image)
btn_layout.addWidget(self.btn_load_folder)
btn_layout.addWidget(self.btn_camera)
btn_layout.addWidget(self.btn_detect)
btn_layout.addWidget(self.btn_export)
btn_layout.addWidget(self.btn_clear)
btn_group.setLayout(btn_layout)
# 统计信息
stats_group = QGroupBox("实时统计")
stats_layout = QVBoxLayout()
self.total_count_label = QLabel("检测总数: 0")
self.unripe_count_label = QLabel("未成熟: 0")
self.semiripe_count_label = QLabel("半成熟: 0")
self.ripe_count_label = QLabel("成熟: 0")
self.maturity_rate_label = QLabel("成熟度比例: -")
stats_layout.addWidget(self.total_count_label)
stats_layout.addWidget(self.unripe_count_label)
stats_layout.addWidget(self.semiripe_count_label)
stats_layout.addWidget(self.ripe_count_label)
stats_layout.addWidget(self.maturity_rate_label)
stats_group.setLayout(stats_layout)
# 添加到控制面板
control_layout.addWidget(model_group)
control_layout.addWidget(param_group)
control_layout.addWidget(btn_group)
control_layout.addWidget(stats_group)
control_layout.addStretch()
control_panel.setLayout(control_layout)
control_panel.setMaximumWidth(350)
return control_panel
def create_display_panel(self):
"""创建显示面板"""
display_panel = QWidget()
display_layout = QVBoxLayout(display_panel)
# 创建标签页
self.tab_widget = QTabWidget()
# 图像显示标签页
self.image_tab = QWidget()
image_layout = QVBoxLayout(self.image_tab)
# 原始图像显示
self.image_label = QLabel()
self.image_label.setAlignment(Qt.AlignCenter)
self.image_label.setMinimumSize(640, 480)
self.image_label.setText("等待加载图像...")
self.image_label.setStyleSheet("border: 2px dashed #aaa; background-color: #f0f0f0;")
# 结果图像显示
self.result_label = QLabel()
self.result_label.setAlignment(Qt.AlignCenter)
self.result_label.setMinimumSize(640, 480)
self.result_label.setText("检测结果将显示在这里...")
self.result_label.setStyleSheet("border: 2px dashed #aaa; background-color: #f0f0f0;")
# 使用分割器显示两个图像
splitter = QSplitter(Qt.Horizontal)
splitter.addWidget(self.create_image_container("原始图像", self.image_label))
splitter.addWidget(self.create_image_container("检测结果", self.result_label))
splitter.setSizes([600, 600])
image_layout.addWidget(splitter)
# 结果表格标签页
self.table_tab = QWidget()
table_layout = QVBoxLayout(self.table_tab)
self.result_table = QTableWidget()
self.result_table.setColumnCount(6)
self.result_table.setHorizontalHeaderLabels([
"序号", "类别", "置信度", "边界框", "面积", "时间"
])
self.result_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
table_layout.addWidget(self.result_table)
# 统计图表标签页
self.stats_tab = QWidget()
stats_layout = QVBoxLayout(self.stats_tab)
self.stats_label = QLabel("统计图表功能需要matplotlib支持")
self.stats_label.setAlignment(Qt.AlignCenter)
stats_layout.addWidget(self.stats_label)
# 添加标签页
self.tab_widget.addTab(self.image_tab, "📷 图像检测")
self.tab_widget.addTab(self.table_tab, "📊 检测结果")
self.tab_widget.addTab(self.stats_tab, "📈 统计分析")
display_layout.addWidget(self.tab_widget)
# 状态栏
self.status_bar = QLabel("就绪")
self.status_bar.setStyleSheet("padding: 5px; background-color: #e0e0e0;")
display_layout.addWidget(self.status_bar)
# 进度条
self.progress_bar = QProgressBar()
self.progress_bar.setVisible(False)
display_layout.addWidget(self.progress_bar)
return display_panel
def create_image_container(self, title, label):
"""创建图像显示容器"""
container = QWidget()
layout = QVBoxLayout(container)
layout.addWidget(QLabel(f"<b>{title}</b>"))
layout.addWidget(label)
return container
def set_style(self):
"""设置界面样式"""
self.setStyleSheet("""
QMainWindow {
background-color: #f5f5f5;
}
QGroupBox {
font-weight: bold;
border: 2px solid #4CAF50;
border-radius: 5px;
margin-top: 10px;
padding-top: 10px;
}
QGroupBox::title {
subcontrol-origin: margin;
left: 10px;
padding: 0 5px 0 5px;
}
QPushButton {
background-color: #4CAF50;
color: white;
border: none;
padding: 10px;
border-radius: 5px;
font-weight: bold;
min-height: 30px;
}
QPushButton:hover {
background-color: #45a049;
}
QPushButton:disabled {
background-color: #cccccc;
}
QLabel {
padding: 2px;
}
QTableWidget {
background-color: white;
alternate-background-color: #f9f9f9;
selection-background-color: #4CAF50;
}
QHeaderView::section {
background-color: #4CAF50;
color: white;
padding: 5px;
border: 1px solid #ddd;
}
QTabWidget::pane {
border: 1px solid #ddd;
background-color: white;
}
QTabBar::tab {
background-color: #e0e0e0;
padding: 8px 16px;
margin-right: 2px;
}
QTabBar::tab:selected {
background-color: #4CAF50;
color: white;
}
""")
def load_default_model(self):
"""加载默认模型"""
try:
model_name = self.model_combo.currentText()
model_path = f"models/{model_name.lower()}_strawberry.pt"
if os.path.exists(model_path):
self.model = YOLO(model_path)
self.status_bar.setText(f"✅ 模型加载成功: {model_name}")
else:
self.status_bar.setText("⚠️ 模型文件不存在，请先训练模型")
self.btn_detect.setEnabled(False)
except Exception as e:
QMessageBox.warning(self, "错误", f"加载模型失败: {str(e)}")
self.model = None
self.btn_detect.setEnabled(False)
def on_model_changed(self, model_name):
"""模型选择改变"""
self.load_default_model()
def load_model(self):
"""加载自定义模型"""
file_path, _ = QFileDialog.getOpenFileName(
self, "选择模型文件",
"",
"模型文件 (*.pt *.pth);;所有文件 (*.*)"
)
if file_path:
try:
self.model = YOLO(file_path)
self.status_bar.setText(f"✅ 自定义模型加载成功: {os.path.basename(file_path)}")
self.btn_detect.setEnabled(True)
except Exception as e:
QMessageBox.critical(self, "错误", f"加载模型失败: {str(e)}")
def load_image(self):
"""加载图像"""
file_path, _ = QFileDialog.getOpenFileName(
self, "选择图像",
"",
"图像文件 (*.jpg *.jpeg *.png *.bmp);;所有文件 (*.*)"
)
if file_path:
self.load_and_display_image(file_path)
def load_folder(self):
"""加载文件夹"""
folder_path = QFileDialog.getExistingDirectory(
self, "选择图像文件夹"
)
if folder_path:
# 这里可以实现批量处理功能
QMessageBox.information(self, "信息",
f"已选择文件夹: {folder_path}\n批量处理功能开发中...")
def load_and_display_image(self, image_path):
"""加载并显示图像"""
try:
# 使用OpenCV读取图像
image = cv2.imread(image_path)
if image is None:
raise ValueError("无法读取图像文件")
# 转换为RGB
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
self.current_image = image_rgb
# 显示图像
self.display_image(self.image_label, image_rgb)
# 更新状态
self.status_bar.setText(f"✅ 图像加载成功: {os.path.basename(image_path)}")
self.btn_detect.setEnabled(True)
# 清空之前的检测结果
self.result_label.clear()
self.result_label.setText("点击'开始检测'进行检测")
self.result_table.setRowCount(0)
except Exception as e:
QMessageBox.critical(self, "错误", f"加载图像失败: {str(e)}")
def display_image(self, label, image):
"""在QLabel中显示图像"""
if image is None:
return
# 调整图像大小以适应标签
h, w, ch = image.shape
bytes_per_line = ch * w
# 创建QImage
qimage = QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888)
# 缩放图像以适合标签
scaled_pixmap = QPixmap.fromImage(qimage).scaled(
label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation
)
label.setPixmap(scaled_pixmap)
def start_detection(self):
"""开始检测"""
if self.current_image is None:
QMessageBox.warning(self, "警告", "请先加载图像")
return
if self.model is None:
QMessageBox.warning(self, "警告", "请先加载模型")
return
# 创建检测线程
self.detection_thread = DetectionThread()
self.detection_thread.setup(
self.model.ckpt_path if hasattr(self.model, 'ckpt_path') else 'models/yolov8n_strawberry.pt',
self.current_image,
self.conf_slider.value() / 100,
self.iou_slider.value() / 100
)
# 连接信号
self.detection_thread.detection_finished.connect(self.on_detection_finished)
self.detection_thread.progress_updated.connect(self.progress_bar.setValue)
# 更新UI状态
self.btn_detect.setEnabled(False)
self.btn_detect.setText("检测中...")
self.progress_bar.setVisible(True)
self.progress_bar.setRange(0, 100)
self.status_bar.setText("正在检测...")
# 启动线程
self.detection_thread.start()
def on_detection_finished(self, detections, result_image):
"""检测完成处理"""
# 显示结果图像
self.display_image(self.result_label, result_image)
# 更新结果表格
self.update_results_table(detections)
# 更新统计信息
self.update_statistics(detections)
# 保存到历史记录
self.save_to_history(detections, result_image)
# 恢复UI状态
self.btn_detect.setEnabled(True)
self.btn_detect.setText("🔍 开始检测")
self.progress_bar.setVisible(False)
self.status_bar.setText(f"✅ 检测完成，共发现 {len(detections)} 个草莓")
# 切换到结果标签页
self.tab_widget.setCurrentIndex(1)
def update_results_table(self, detections):
"""更新结果表格"""
self.result_table.setRowCount(len(detections))
for i, detection in enumerate(detections):
# 序号
self.result_table.setItem(i, 0, QTableWidgetItem(str(i + 1)))
# 类别
class_item = QTableWidgetItem(detection['class_name'])
class_item.setBackground(self.get_class_color(detection['class_id']))
self.result_table.setItem(i, 1, class_item)
# 置信度
conf_item = QTableWidgetItem(f"{detection['confidence']:.4f}")
self.result_table.setItem(i, 2, conf_item)
# 边界框
bbox = detection['bbox']
bbox_item = QTableWidgetItem(f"[{bbox[0]}, {bbox[1]}, {bbox[2]}, {bbox[3]}]")
self.result_table.setItem(i, 3, bbox_item)
# 面积
area_item = QTableWidgetItem(str(detection['area']))
self.result_table.setItem(i, 4, area_item)
# 时间
time_item = QTableWidgetItem(datetime.now().strftime("%H:%M:%S"))
self.result_table.setItem(i, 5, time_item)
def get_class_color(self, class_id):
"""获取类别对应的颜色"""
colors = [
QColor(144, 238, 144), # 未成熟 - 浅绿
QColor(255, 165, 0), # 半成熟 - 橙色
QColor(255, 99, 71) # 成熟 - 红色
]
return colors[class_id % len(colors)]
def update_statistics(self, detections):
"""更新统计信息"""
# 计数
total = len(detections)
unripe = sum(1 for d in detections if d['class_id'] == 0)
semiripe = sum(1 for d in detections if d['class_id'] == 1)
ripe = sum(1 for d in detections if d['class_id'] == 2)
# 更新标签
self.total_count_label.setText(f"检测总数: {total}")
self.unripe_count_label.setText(f"未成熟: {unripe}")
self.semiripe_count_label.setText(f"半成熟: {semiripe}")
self.ripe_count_label.setText(f"成熟: {ripe}")
# 计算比例
if total > 0:
unripe_rate = unripe / total * 100
semiripe_rate = semiripe / total * 100
ripe_rate = ripe / total * 100
self.maturity_rate_label.setText(
f"成熟度比例: 未成熟({unripe_rate:.1f}%) | "
f"半成熟({semiripe_rate:.1f}%) | "
f"成熟({ripe_rate:.1f}%)"
)
else:
self.maturity_rate_label.setText("成熟度比例: -")
def save_to_history(self, detections, result_image):
"""保存到历史记录"""
history_entry = {
'timestamp': datetime.now().isoformat(),
'detections': detections,
'image_size': result_image.shape[:2],
'total_count': len(detections)
}
self.detection_history.append(history_entry)
def update_conf_threshold(self, value):
"""更新置信度阈值"""
conf = value / 100
self.conf_label.setText(f"置信度阈值: {conf:.2f}")
def update_iou_threshold(self, value):
"""更新IoU阈值"""
iou = value / 100
self.iou_label.setText(f"IoU阈值: {iou:.2f}")
def start_camera(self):
"""启动摄像头"""
QMessageBox.information(self, "信息", "摄像头功能开发中...")
def export_results(self):
"""导出结果"""
if not self.detection_history:
QMessageBox.warning(self, "警告", "没有可导出的检测结果")
return
# 选择导出文件路径
file_path, _ = QFileDialog.getSaveFileName(
self, "导出结果",
f"strawberry_detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
"JSON文件 (*.json);;CSV文件 (*.csv);;所有文件 (*.*)"
)
if file_path:
try:
if file_path.endswith('.json'):
# 导出为JSON
with open(file_path, 'w', encoding='utf-8') as f:
json.dump(self.detection_history, f, indent=2, ensure_ascii=False)
elif file_path.endswith('.csv'):
# 导出为CSV
self.export_to_csv(file_path)
else:
# 默认导出为JSON
with open(file_path + '.json', 'w', encoding='utf-8') as f:
json.dump(self.detection_history, f, indent=2, ensure_ascii=False)
QMessageBox.information(self, "成功", f"结果已导出到: {file_path}")
except Exception as e:
QMessageBox.critical(self, "错误", f"导出失败: {str(e)}")
def export_to_csv(self, file_path):
"""导出为CSV格式"""
# 准备数据
data = []
for entry in self.detection_history:
for detection in entry['detections']:
data.append({
'timestamp': entry['timestamp'],
'class_name': detection['class_name'],
'confidence': detection['confidence'],
'bbox': str(detection['bbox']),
'area': detection['area']
})
# 创建DataFrame并保存
df = pd.DataFrame(data)
df.to_csv(file_path, index=False, encoding='utf-8-sig')
def clear_results(self):
"""清除结果"""
reply = QMessageBox.question(
self, "确认",
"确定要清除所有结果吗？",
QMessageBox.Yes | QMessageBox.No,
QMessageBox.No
)
if reply == QMessageBox.Yes:
# 清空图像
self.image_label.clear()
self.image_label.setText("等待加载图像...")
self.result_label.clear()
self.result_label.setText("检测结果将显示在这里...")
# 清空表格
self.result_table.setRowCount(0)
# 清空统计
self.total_count_label.setText("检测总数: 0")
self.unripe_count_label.setText("未成熟: 0")
self.semiripe_count_label.setText("半成熟: 0")
self.ripe_count_label.setText("成熟: 0")
self.maturity_rate_label.setText("成熟度比例: -")
# 清空历史记录
self.detection_history.clear()
# 重置状态
self.current_image = None
self.btn_detect.setEnabled(False)
self.status_bar.setText("已清除所有结果")
def closeEvent(self, event):
"""关闭事件"""
if self.detection_thread and self.detection_thread.isRunning():
self.detection_thread.stop()
self.detection_thread.wait()
# 保存设置
self.save_settings()
event.accept()
def save_settings(self):
"""保存设置"""
settings = {
'last_model': self.model_combo.currentText(),
'conf_threshold': self.conf_slider.value(),
'iou_threshold': self.iou_slider.value(),
'window_geometry': {
'x': self.x(),
'y': self.y(),
'width': self.width(),
'height': self.height()
}
}
# 保存到文件
try:
with open('app_settings.json', 'w') as f:
json.dump(settings, f, indent=2)
except:
pass
def main():
"""主函数"""
app = QApplication(sys.argv)
app.setApplicationName("草莓成熟度检测系统")
app.setApplicationDisplayName("🍓 Strawberry Maturity Detection")
# 设置应用程序图标
app.setWindowIcon(QIcon("strawberry_icon.png") if os.path.exists("strawberry_icon.png") else QIcon())
# 创建并显示主窗口
window = StrawberryDetectionUI()
window.show()
# 运行应用程序
sys.exit(app.exec_())
if __name__ == '__main__':
main()
6.2 Streamlit Web应用
python
# streamlit_app.py
"""
草莓成熟度检测系统 - Streamlit Web界面
"""
import streamlit as st
import cv2
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from PIL import Image
import tempfile
import os
from datetime import datetime
import json
import torch
from ultralytics import YOLO
# 页面配置
st.set_page_config(
page_title="🍓 草莓成熟度检测系统",
page_icon="🍓",
layout="wide",
initial_sidebar_state="expanded"
)
# 应用标题
st.title("🍓 草莓成熟度智能检测系统")
st.markdown("---")
class StrawberryDetector:
"""草莓检测器类"""
def __init__(self, model_path='models/yolov8n_strawberry.pt'):
self.model_path = model_path
self.model = None
self.load_model()
def load_model(self):
"""加载模型"""
try:
if os.path.exists(self.model_path):
self.model = YOLO(self.model_path)
return True
else:
st.warning(f"模型文件不存在: {self.model_path}")
return False
except Exception as e:
st.error(f"加载模型失败: {str(e)}")
return False
def predict(self, image, conf_threshold=0.25, iou_threshold=0.45):
"""执行预测"""
if self.model is None:
return None, None
# 执行检测
results = self.model(image, conf=conf_threshold, iou=iou_threshold)
# 解析结果
detections = []
result_image = image.copy()
if results[0].boxes is not None:
boxes = results[0].boxes.xyxy.cpu().numpy()
classes = results[0].boxes.cls.cpu().numpy()
scores = results[0].boxes.conf.cpu().numpy()
# 类别颜色
colors = [(0, 255, 0), (0, 165, 255), (0, 0, 255)]
class_names = ['未成熟', '半成熟', '成熟']
for box, cls, score in zip(boxes, classes, scores):
x1, y1, x2, y2 = map(int, box)
class_id = int(cls)
confidence = float(score)
# 绘制边界框
color = colors[class_id % len(colors)]
cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)
# 绘制标签
label = f"{class_names[class_id]}: {confidence:.2f}"
cv2.putText(result_image, label, (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
# 添加到检测结果
detections.append({
'class_id': class_id,
'class_name': class_names[class_id],
'confidence': confidence,
'bbox': [x1, y1, x2, y2],
'area': (x2 - x1) * (y2 - y1)
})
return detections, result_image
def main():
"""主函数"""
# 初始化检测器
if 'detector' not in st.session_state:
st.session_state.detector = StrawberryDetector()
# 侧边栏
with st.sidebar:
st.header("⚙️ 设置")
# 模型选择
model_option = st.selectbox(
"选择模型",
["YOLOv8n", "YOLOv8s", "YOLOv8m", "YOLOv8l", "YOLOv8x"],
index=0
)
# 参数设置
conf_threshold = st.slider(
"置信度阈值",
min_value=0.0,
max_value=1.0,
value=0.25,
step=0.01,
help="检测结果的置信度阈值"
)
iou_threshold = st.slider(
"IoU阈值",
min_value=0.0,
max_value=1.0,
value=0.45,
step=0.01,
help="非极大值抑制的IoU阈值"
)
# 上传文件
st.header("📁 上传图像")
uploaded_file = st.file_uploader(
"选择图像文件",
type=['jpg', 'jpeg', 'png', 'bmp'],
help="上传草莓图像进行检测"
)
# 摄像头选项
use_camera = st.checkbox("使用摄像头")
if use_camera:
camera_image = st.camera_input("拍照")
if camera_image:
uploaded_file = camera_image
# 示例图像
st.header("🎯 示例")
if st.button("使用示例图像"):
# 这里可以加载示例图像
st.info("示例图像功能开发中...")
st.markdown("---")
st.markdown("### 📊 系统信息")
st.write(f"PyTorch版本: {torch.__version__}")
st.write(f"CUDA可用: {torch.cuda.is_available()}")
# 主内容区
col1, col2 = st.columns(2)
with col1:
st.header("📷 原始图像")
if uploaded_file is not None:
# 读取图像
file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 显示原始图像
st.image(image_rgb, caption="原始图像", use_column_width=True)
# 存储到session state
st.session_state.image = image_rgb
st.session_state.image_uploaded = True
with col2:
st.header("🔍 检测结果")
if 'image_uploaded' in st.session_state and st.session_state.image_uploaded:
if st.button("开始检测", type="primary"):
with st.spinner("正在检测中..."):
# 执行检测
detections, result_image = st.session_state.detector.predict(
st.session_state.image,
conf_threshold,
iou_threshold
)
if detections is not None:
# 显示结果图像
st.image(result_image, caption="检测结果", use_column_width=True)
# 存储结果
st.session_state.detections = detections
st.session_state.result_image = result_image
st.session_state.detection_done = True
st.success(f"检测完成！共发现 {len(detections)} 个草莓")
else:
st.error("检测失败！")
else:
st.info("请先上传图像")
# 显示检测结果
if 'detection_done' in st.session_state and st.session_state.detection_done:
st.markdown("---")
st.header("📊 检测结果分析")
# 创建结果表格
detections = st.session_state.detections
if detections:
# 转换为DataFrame
df = pd.DataFrame(detections)
# 显示统计信息
col1, col2, col3, col4 = st.columns(4)
with col1:
st.metric("总检测数", len(detections))
with col2:
unripe_count = len(df[df['class_id'] == 0])
st.metric("未成熟", unripe_count)
with col3:
semiripe_count = len(df[df['class_id'] == 1])
st.metric("半成熟", semiripe_count)
with col4:
ripe_count = len(df[df['class_id'] == 2])
st.metric("成熟", ripe_count)
# 显示详细结果表格
st.subheader("详细检测结果")
# 格式化表格数据
table_data = []
for i, det in enumerate(detections):
table_data.append({
'序号': i + 1,
'类别': det['class_name'],
'置信度': f"{det['confidence']:.4f}",
'边界框': str(det['bbox']),
'面积': det['area']
})
df_table = pd.DataFrame(table_data)
st.dataframe(df_table, use_container_width=True)
# 可视化分析
st.subheader("📈 可视化分析")
# 创建标签页
tab1, tab2, tab3 = st.tabs(["类别分布", "置信度分布", "面积分布"])
with tab1:
# 类别分布饼图
class_counts = df['class_name'].value_counts()
fig1 = go.Figure(data=[go.Pie(
labels=class_counts.index,
values=class_counts.values,
hole=.3,
marker_colors=['lightgreen', 'orange', 'red']
)])
fig1.update_layout(title="草莓成熟度分布")
st.plotly_chart(fig1, use_container_width=True)
with tab2:
# 置信度分布直方图
fig2 = px.histogram(
df,
x='confidence',
color='class_name',
nbins=20,
title='置信度分布',
color_discrete_map={
'未成熟': 'lightgreen',
'半成熟': 'orange',
'成熟': 'red'
}
)
st.plotly_chart(fig2, use_container_width=True)
with tab3:
# 面积分布散点图
fig3 = px.scatter(
df,
x='area',
y='confidence',
color='class_name',
size='confidence',
hover_data=['bbox'],
title='面积与置信度关系',
color_discrete_map={
'未成熟': 'lightgreen',
'半成熟': 'orange',
'成熟': 'red'
}
)
st.plotly_chart(fig3, use_container_width=True)
# 导出结果
st.subheader("💾 导出结果")
col1, col2 = st.columns(2)
with col1:
# 导出为CSV
if st.button("导出CSV"):
csv = df_table.to_csv(index=False)
st.download_button(
label="下载CSV文件",
data=csv,
file_name=f"strawberry_detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
mime="text/csv"
)
with col2:
# 导出为JSON
if st.button("导出JSON"):
json_data = json.dumps(detections, indent=2, ensure_ascii=False)
st.download_button(
label="下载JSON文件",
data=json_data,
file_name=f"strawberry_detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
mime="application/json"
)
else:
st.warning("未检测到草莓")
# 页脚
st.markdown("---")
st.markdown(
"""
<div style='text-align: center'>
<p>🍓 草莓成熟度检测系统 v1.0 | 基于YOLOv8深度学习技术</p>
<p>© 2024 农业智能检测实验室 | 技术支持: deeplearning.ai</p>
</div>
""",
unsafe_allow_html=True
)
if __name__ == '__main__':
main()
七、系统部署与优化
7.1 部署方案
python
# deployment.py
"""
系统部署脚本
支持多种部署方式
"""
import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
import tensorrt as trt
import numpy as np
from pathlib import Path
import cv2
import json
class StrawberryDeployer:
"""草莓检测模型部署器"""
def __init__(self, model_path):
self.model_path = model_path
self.model = None
self.load_model()
def load_model(self):
"""加载模型"""
from ultralytics import YOLO
self.model = YOLO(self.model_path)
def export_onnx(self, output_path, opset=12):
"""导出ONNX模型"""
print("🔄 正在导出ONNX模型...")
# 导出设置
export_params = {
'imgsz': 640,
'batch': 1,
'device': 'cpu',
'simplify': True,
'opset': opset,
'dynamic': False
}
# 执行导出
success = self.model.export(format='onnx', **export_params)
if success:
print(f"✅ ONNX模型已导出: {output_path}")
return True
else:
print("❌ ONNX模型导出失败")
return False
def export_tensorrt(self, output_path):
"""导出TensorRT模型"""
print("🔄 正在导出TensorRT模型...")
# 检查CUDA可用性
if not torch.cuda.is_available():
print("❌ CUDA不可用，无法导出TensorRT模型")
return False
# 导出设置
export_params = {
'imgsz': 640,
'batch': 1,
'device': 0,
'half': True, # FP16精度
'workspace': 4 # GB
}
# 执行导出
success = self.model.export(format='engine', **export_params)
if success:
print(f"✅ TensorRT模型已导出: {output_path}")
return True
else:
print("❌ TensorRT模型导出失败")
return False
def export_openvino(self, output_path):
"""导出OpenVINO模型"""
print("🔄 正在导出OpenVINO模型...")
# 导出设置
export_params = {
'imgsz': 640,
'batch': 1,
'device': 'cpu',
'half': False
}
# 执行导出
success = self.model.export(format='openvino', **export_params)
if success:
print(f"✅ OpenVINO模型已导出: {output_path}")
return True
else:
print("❌ OpenVINO模型导出失败")
return False
def optimize_for_mobile(self, output_path):
"""为移动设备优化模型"""
print("🔄 正在为移动设备优化模型...")
try:
# 转换为TorchScript
scripted_model = torch.jit.trace(self.model, torch.randn(1, 3, 640, 640))
# 保存模型
torch.jit.save(scripted_model, output_path)
print(f"✅ 移动端优化模型已保存: {output_path}")
return True
except Exception as e:
print(f"❌ 移动端优化失败: {str(e)}")
return False
def create_deployment_package(self, output_dir='deployment_package'):
"""创建完整的部署包"""
print("📦 创建部署包...")
# 创建输出目录
output_path = Path(output_dir)
output_path.mkdir(exist_ok=True)
# 导出各种格式的模型
formats = [
('onnx', self.export_onnx),
('engine', self.export_tensorrt),
('openvino', self.export_openvino),
('torchscript', self.optimize_for_mobile)
]
export_results = {}
for format_name, export_func in formats:
try:
output_file = output_path / f'model.{format_name}'
success = export_func(str(output_file))
export_results[format_name] = success
except Exception as e:
print(f"❌ {format_name.upper()}导出出错: {str(e)}")
export_results[format_name] = False
# 创建部署配置文件
config = {
'model_info': {
'original_model': self.model_path,
'export_time': str(datetime.now()),
'formats': export_results
},
'inference_config': {
'image_size': 640,
'confidence_threshold': 0.25,
'iou_threshold': 0.45,
'classes': ['unripe', 'semi-ripe', 'ripe']
},
'deployment_instructions': {
'onnx': '使用onnxruntime进行推理',
'tensorrt': '需要NVIDIA GPU和TensorRT环境',
'openvino': '使用OpenVINO推理引擎',
'torchscript': '适用于移动端和嵌入式设备'
}
}
# 保存配置文件
config_file = output_path / 'deployment_config.json'
with open(config_file, 'w') as f:
json.dump(config, f, indent=2)
print(f"✅ 部署包创建完成: {output_path}")
print(f"📄 配置文件: {config_file}")
return output_path
class InferenceEngine:
"""推理引擎基类"""
def __init__(self, model_path):
self.model_path = model_path
self.model = None
self.class_names = ['unripe', 'semi-ripe', 'ripe']
def preprocess(self, image):
"""图像预处理"""
# 调整大小
img_resized = cv2.resize(image, (640, 640))
# 归一化
img_normalized = img_resized / 255.0
# 转换通道顺序
img_transposed = np.transpose(img_normalized, (2, 0, 1))
# 添加批次维度
img_batched = np.expand_dims(img_transposed, axis=0).astype(np.float32)
return img_batched
def postprocess(self, outputs, conf_threshold=0.25, iou_threshold=0.45):
"""后处理"""
detections = []
# 这里需要根据具体的输出格式进行解析
# 简化的实现
if isinstance(outputs, np.ndarray):
# 假设输出是 [batch, num_boxes, 6] 格式
for detection in outputs[0]:
if detection[4] > conf_threshold: # 置信度
x1, y1, x2, y2 = detection[:4]
conf = detection[4]
cls_id = int(detection[5])
detections.append({
'bbox': [float(x1), float(y1), float(x2), float(y2)],
'confidence': float(conf),
'class_id': cls_id,
'class_name': self.class_names[cls_id] if cls_id < len(self.class_names) else 'unknown'
})
return detections
class ONNXInference(InferenceEngine):
"""ONNX推理引擎"""
def load_model(self):
"""加载ONNX模型"""
self.session = ort.InferenceSession(self.model_path)
self.input_name = self.session.get_inputs()[0].name
def inference(self, image):
"""执行推理"""
# 预处理
input_tensor = self.preprocess(image)
# 推理
outputs = self.session.run(None, {self.input_name: input_tensor})
# 后处理
detections = self.postprocess(outputs[0])
return detections
class TensorRTInference(InferenceEngine):
"""TensorRT推理引擎"""
def load_model(self):
"""加载TensorRT引擎"""
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
# 加载引擎
with open(self.model_path, 'rb') as f:
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
self.engine = runtime.deserialize_cuda_engine(f.read())
# 创建执行上下文
self.context = self.engine.create_execution_context()
# 分配内存
self.inputs, self.outputs, self.bindings = [], [], []
self.stream = cuda.Stream()
for binding in self.engine:
size = trt.volume(self.engine.get_binding_shape(binding))
dtype = trt.nptype(self.engine.get_binding_dtype(binding))
# 分配主机和设备内存
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
self.bindings.append(int(device_mem))
if self.engine.binding_is_input(binding):
self.inputs.append({'host': host_mem, 'device': device_mem})
else:
self.outputs.append({'host': host_mem, 'device': device_mem})
def inference(self, image):
"""执行推理"""
# 预处理
input_array = self.preprocess(image)
# 复制数据到设备
np.copyto(self.inputs[0]['host'], input_array.ravel())
cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
# 执行推理
self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
# 复制结果回主机
cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device'], self.stream)
self.stream.synchronize()
# 获取输出
output = self.outputs[0]['host']
# 重塑输出形状（需要根据实际模型调整）
output_shape = self.context.get_binding_shape(1) # 假设输出是第二个绑定
output = output.reshape(output_shape)
# 后处理
detections = self.postprocess(output)
return detections
def benchmark_inference(model_path, num_iterations=100, warmup=10):
"""推理性能基准测试"""
print("⚡ 开始性能基准测试...")
results = {}
# 测试不同推理引擎
engines = [
('PyTorch', YOLO),
('ONNX', ONNXInference),
('TensorRT', TensorRTInference) if torch.cuda.is_available() else None
]
# 创建测试图像
test_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
for engine_name, engine_class in engines:
if engine_class is None:
continue
print(f"\n测试 {engine_name} 引擎...")
try:
# 加载模型
if engine_name == 'PyTorch':
model = engine_class(model_path)
model.to('cuda' if torch.cuda.is_available() else 'cpu')
else:
model = engine_class(model_path)
model.load_model()
# Warmup
for _ in range(warmup):
if engine_name == 'PyTorch':
model(test_image, verbose=False)
else:
model.inference(test_image)
# 性能测试
import time
times = []
for i in range(num_iterations):
start_time = time.time()
if engine_name == 'PyTorch':
model(test_image, verbose=False)
else:
model.inference(test_image)
end_time = time.time()
times.append(end_time - start_time)
# 计算统计信息
avg_time = np.mean(times) * 1000 # 转换为毫秒
fps = 1000 / avg_time if avg_time > 0 else 0
std_time = np.std(times) * 1000
results[engine_name] = {
'avg_inference_time_ms': avg_time,
'fps': fps,
'std_time_ms': std_time,
'min_time_ms': np.min(times) * 1000,
'max_time_ms': np.max(times) * 1000
}
print(f" ✅ 平均推理时间: {avg_time:.2f} ms")
print(f" ✅ FPS: {fps:.2f}")
except Exception as e:
print(f" ❌ 测试失败: {str(e)}")
results[engine_name] = {'error': str(e)}
# 打印对比结果
print("\n" + "="*50)
print("📊 性能测试结果对比:")
print("="*50)
for engine_name, metrics in results.items():
if 'error' not in metrics:
print(f"\n{engine_name}:")
print(f" ├── 平均推理时间: {metrics['avg_inference_time_ms']:.2f} ms")
print(f" ├── FPS: {metrics['fps']:.2f}")
print(f" ├── 时间标准差: {metrics['std_time_ms']:.2f} ms")
print(f" └── 时间范围: [{metrics['min_time_ms']:.2f}, {metrics['max_time_ms']:.2f}] ms")
return results
# 使用示例
if __name__ == '__main__':
# 模型路径
MODEL_PATH = 'runs/train/strawberry_detection/weights/best.pt'
# 创建部署器
deployer = StrawberryDeployer(MODEL_PATH)
# 创建部署包
package_path = deployer.create_deployment_package()
# 性能基准测试
benchmark_results = benchmark_inference(MODEL_PATH)
print(f"\n🎉 部署完成!")
print(f"📦 部署包位置: {package_path}")
print(f"📊 性能报告已生成")
7.2 性能优化
python
# optimization_techniques.py
"""
性能优化技术实现
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.quantization import quantize_dynamic
import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic as quantize_onnx
class ModelOptimizer:
"""模型优化器"""
@staticmethod
def prune_model(model, pruning_rate=0.3):
"""模型剪枝"""
print("✂️ 执行模型剪枝...")
# 获取所有卷积层
conv_layers = []
for name, module in model.named_modules():
if isinstance(module, nn.Conv2d):
conv_layers.append((name, module))
# L1范数剪枝
for name, conv in conv_layers:
weight = conv.weight.data
# 计算权重绝对值
weight_abs = torch.abs(weight)
# 计算剪枝阈值
threshold = torch.quantile(weight_abs.view(-1), pruning_rate)
# 创建掩码
mask = weight_abs > threshold
# 应用剪枝
conv.weight.data *= mask.float()
print(f"✅ 模型剪枝完成，剪枝率: {pruning_rate*100}%")
return model
@staticmethod
def quantize_pytorch_model(model):
"""PyTorch模型动态量化"""
print("⚖️ 执行PyTorch模型量化...")
# 动态量化
quantized_model = quantize_dynamic(
model, # 原始模型
{nn.Linear, nn.Conv2d}, # 要量化的模块类型
dtype=torch.qint8 # 量化类型
)
print("✅ PyTorch模型量化完成")
return quantized_model
@staticmethod
def quantize_onnx_model(onnx_model_path, output_path):
"""ONNX模型量化"""
print("⚖️ 执行ONNX模型量化...")
try:
# 动态量化
quantized_model = quantize_onnx(
onnx_model_path,
output_path,
weight_type=quantize_onnx.QuantType.QInt8
)
print(f"✅ ONNX模型量化完成: {output_path}")
return True
except Exception as e:
print(f"❌ ONNX模型量化失败: {str(e)}")
return False
@staticmethod
def apply_knowledge_distillation(teacher_model, student_model, dataloader, epochs=10):
"""知识蒸馏"""
print("🎓 执行知识蒸馏...")
criterion_kd = nn.KLDivLoss() # KL散度损失
criterion_ce = nn.CrossEntropyLoss() # 交叉熵损失
optimizer = torch.optim.Adam(student_model.parameters(), lr=0.001)
# 温度参数
temperature = 4.0
alpha = 0.7 # 蒸馏损失权重
teacher_model.eval()
student_model.train()
for epoch in range(epochs):
total_loss = 0
for batch_idx, (images, labels) in enumerate(dataloader):
# 前向传播
with torch.no_grad():
teacher_logits = teacher_model(images)
student_logits = student_model(images)
# 计算损失
# 蒸馏损失
loss_kd = criterion_kd(
F.log_softmax(student_logits / temperature, dim=1),
F.softmax(teacher_logits / temperature, dim=1)
) * (temperature ** 2)
# 分类损失
loss_ce = criterion_ce(student_logits, labels)
# 总损失
loss = alpha * loss_kd + (1 - alpha) * loss_ce
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
print(f" Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}")
print("✅ 知识蒸馏完成")
return student_model
@staticmethod
def optimize_for_inference(model, input_shape=(1, 3, 640, 640)):
"""推理优化"""
print("⚡ 执行推理优化...")
# 设置为评估模式
model.eval()
# 开启推理模式优化
with torch.no_grad():
# 示例输入
example_input = torch.randn(input_shape)
# JIT编译
if hasattr(torch.jit, 'trace'):
try:
traced_model = torch.jit.trace(model, example_input)
traced_model = torch.jit.freeze(traced_model) # 冻结模型
traced_model = torch.jit.optimize_for_inference(traced_model)
print("✅ JIT编译优化完成")
return traced_model
except Exception as e:
print(f"⚠️ JIT编译失败: {str(e)}")
return model
class MemoryOptimizer:
"""内存优化器"""
@staticmethod
def optimize_memory_usage():
"""优化内存使用"""
import gc
print("🧠 优化内存使用...")
# 清理缓存
torch.cuda.empty_cache() if torch.cuda.is_available() else None
# 垃圾回收
gc.collect()
# 设置优化标志
torch.backends.cudnn.benchmark = True # 优化卷积算法
torch.backends.cudnn.deterministic = False # 允许非确定性算法以加速
print("✅ 内存优化完成")
@staticmethod
def check_memory_usage():
"""检查内存使用情况"""
if torch.cuda.is_available():
memory_allocated = torch.cuda.memory_allocated() / 1024**3 # GB
memory_cached = torch.cuda.memory_reserved() / 1024**3 # GB
print(f"🎯 GPU内存使用:")
print(f" ├── 已分配: {memory_allocated:.2f} GB")
print(f" └── 已缓存: {memory_cached:.2f} GB")
return {
'allocated_gb': memory_allocated,
'cached_gb': memory_cached
}
else:
print("ℹ️ CUDA不可用，无法检查GPU内存")
return None
class InferenceOptimizer:
"""推理优化器"""
def __init__(self, model, device='cuda'):
self.model = model
self.device = device
if self.device == 'cuda' and torch.cuda.is_available():
self.model.to('cuda')
def optimize_batch_inference(self, batch_size=32):
"""优化批量推理"""
print(f"📦 优化批量推理，批次大小: {batch_size}")
# 设置模型为评估模式
self.model.eval()
# 创建优化后的推理函数
def batch_predict(images):
with torch.no_grad():
# 批量处理
if len(images) > batch_size:
# 分批处理
all_detections = []
for i in range(0, len(images), batch_size):
batch = images[i:i+batch_size]
batch_tensor = torch.stack(batch).to(self.device)
# 推理
results = self.model(batch_tensor)
all_detections.extend(results)
return all_detections
else:
# 单批处理
images_tensor = torch.stack(images).to(self.device)
return self.model(images_tensor)
return batch_predict
def optimize_single_inference(self, use_half_precision=True):
"""优化单张图像推理"""
print("🖼️ 优化单张图像推理")
# 使用半精度浮点数
if use_half_precision and self.device == 'cuda':
self.model.half()
def single_predict(image):
with torch.no_grad():
# 转换为张量
if isinstance(image, np.ndarray):
image_tensor = torch.from_numpy(image).to(self.device)
if use_half_precision:
image_tensor = image_tensor.half()
else:
image_tensor = image_tensor.float()
else:
image_tensor = image.to(self.device)
# 添加批次维度
if len(image_tensor.shape) == 3:
image_tensor = image_tensor.unsqueeze(0)
# 推理
return self.model(image_tensor)
return single_predict
def create_async_inference(self, max_queue_size=100):
"""创建异步推理"""
print("🚀 创建异步推理管道")
import queue
import threading
class AsyncInference:
def __init__(self, model, device):
self.model = model
self.device = device
self.input_queue = queue.Queue(maxsize=max_queue_size)
self.output_queue = queue.Queue(maxsize=max_queue_size)
self.stop_event = threading.Event()
# 启动工作线程
self.worker_thread = threading.Thread(target=self._inference_worker)
self.worker_thread.start()
def _inference_worker(self):
"""推理工作线程"""
while not self.stop_event.is_set():
try:
# 从队列获取输入
input_data, callback = self.input_queue.get(timeout=0.1)
# 执行推理
with torch.no_grad():
result = self.model(input_data.to(self.device))
# 将结果放入输出队列
if callback:
callback(result)
self.input_queue.task_done()
except queue.Empty:
continue
except Exception as e:
print(f"推理错误: {str(e)}")
def predict_async(self, input_tensor, callback=None):
"""异步预测"""
self.input_queue.put((input_tensor, callback))
def stop(self):
"""停止异步推理"""
self.stop_event.set()
self.worker_thread.join()
return AsyncInference(self.model, self.device)
# 使用示例
if __name__ == '__main__':
# 加载模型
from ultralytics import YOLO
model = YOLO('runs/train/strawberry_detection/weights/best.pt')
# 创建优化器
optimizer = ModelOptimizer()
# 执行优化
print("开始优化模型...")
# 1. 剪枝
pruned_model = optimizer.prune_model(model.model, pruning_rate=0.2)
# 2. 量化
quantized_model = optimizer.quantize_pytorch_model(pruned_model)
# 3. 推理优化
optimized_model = optimizer.optimize_for_inference(quantized_model)
# 内存优化
MemoryOptimizer.optimize_memory_usage()
MemoryOptimizer.check_memory_usage()
# 推理优化
inference_optimizer = InferenceOptimizer(optimized_model)
# 获取优化后的推理函数
batch_predictor = inference_optimizer.optimize_batch_inference(batch_size=16)
single_predictor = inference_optimizer.optimize_single_inference(use_half_precision=True)
print("\n🎉 模型优化完成!")
print("📊 优化后的模型可以用于高效推理")更多推荐
2
0- 0
已为社区贡献731条内容
所有评论(0)