番茄新鲜度检测系统：基于YOLOv5/v8/v10的深度学习解决方案

摘要

本博客详细介绍了一个完整的番茄新鲜度检测系统，该系统采用YOLO系列算法（v5、v8、v10）实现高效准确的番茄新鲜度识别。系统包含完整的数据集处理、模型训练、性能评估以及用户友好的UI界面。本文将深入探讨技术细节，并提供完整的代码实现，帮助读者理解如何构建一个实用的农业产品检测系统。

1. 引言

1.1 背景与意义

在农业供应链和食品工业中，番茄新鲜度检测对质量控制、库存管理和消费者满意度至关重要。传统的人工检测方法效率低下、主观性强且成本高昂。深度学习技术，特别是目标检测算法，为自动化番茄新鲜度检测提供了创新解决方案。

1.2 系统概述

本系统采用最新的YOLO（You Only Look Once）系列算法，构建了一个端到端的番茄新鲜度检测系统。系统能够实时识别和分类番茄的新鲜度等级（新鲜、半新鲜、不新鲜），并提供直观的UI界面供用户使用。

2. 系统架构与设计

2.1 整体架构

text

番茄新鲜度检测系统架构： 1. 数据采集与预处理模块 2. 模型训练与优化模块 3. 推理检测模块 4. 用户界面模块 5. 结果分析与可视化模块

2.2 技术栈

• 深度学习框架：PyTorch

• 目标检测算法：YOLOv5, YOLOv8, YOLOv10

• 前端界面：Gradio/PyQt5

• 数据处理：OpenCV, Pillow, Albumentations

• 模型部署：ONNX, TensorRT

3. 数据集准备与处理

3.1 数据集结构

我们创建了一个包含三个类别的番茄数据集：

• 新鲜番茄（fresh_tomato）

• 半新鲜番茄（semi_fresh_tomato）

• 不新鲜番茄（rotten_tomato）

3.2 数据增强策略

python

import albumentations as A from albumentations.pytorch import ToTensorV2 def get_train_transforms(img_size=640): return A.Compose([ A.RandomResizedCrop(height=img_size, width=img_size, scale=(0.8, 1.0)), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.HueSaturationValue(p=0.2), A.Rotate(limit=30, p=0.5), A.Blur(blur_limit=3, p=0.1), A.CLAHE(p=0.1), ToTensorV2() ], bbox_params=A.BboxParams( format='yolo', label_fields=['class_labels'], min_visibility=0.3 )) def get_val_transforms(img_size=640): return A.Compose([ A.Resize(height=img_size, width=img_size), ToTensorV2() ], bbox_params=A.BboxParams( format='yolo', label_fields=['class_labels'], min_visibility=0.3 ))

3.3 数据集类实现

python

import torch from torch.utils.data import Dataset import cv2 import os import numpy as np class TomatoDataset(Dataset): def __init__(self, images_dir, labels_dir, transforms=None, class_names=None): self.images_dir = images_dir self.labels_dir = labels_dir self.transforms = transforms self.class_names = class_names or ['fresh_tomato', 'semi_fresh_tomato', 'rotten_tomato'] self.class_to_idx = {name: idx for idx, name in enumerate(self.class_names)} # 获取所有图像文件 self.image_files = [f for f in os.listdir(images_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] def __len__(self): return len(self.image_files) def __getitem__(self, idx): # 加载图像 img_path = os.path.join(self.images_dir, self.image_files[idx]) image = cv2.imread(img_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) h, w = image.shape[:2] # 加载标签 label_path = os.path.join(self.labels_dir, os.path.splitext(self.image_files[idx])[0] + '.txt') bboxes = [] class_labels = [] if os.path.exists(label_path): with open(label_path, 'r') as f: for line in f.readlines(): parts = line.strip().split() if len(parts) == 5: class_id = int(parts[0]) x_center = float(parts[1]) * w y_center = float(parts[2]) * h bbox_w = float(parts[3]) * w bbox_h = float(parts[4]) * h # 转换为[x_min, y_min, x_max, y_max] x_min = x_center - bbox_w / 2 y_min = y_center - bbox_h / 2 x_max = x_center + bbox_w / 2 y_max = y_center + bbox_h / 2 bboxes.append([x_min, y_min, x_max, y_max]) class_labels.append(class_id) # 应用数据增强 if self.transforms: transformed = self.transforms( image=image, bboxes=bboxes, class_labels=class_labels ) image = transformed['image'] bboxes = transformed['bboxes'] class_labels = transformed['class_labels'] # 转换边界框格式为YOLO格式 if len(bboxes) > 0: bboxes = torch.tensor(bboxes, dtype=torch.float32) class_labels = torch.tensor(class_labels, dtype=torch.long) # 转换为[x_center, y_center, width, height]格式 target_boxes = torch.zeros((len(bboxes), 4)) target_boxes[:, 0] = (bboxes[:, 0] + bboxes[:, 2]) / 2 # x_center target_boxes[:, 1] = (bboxes[:, 1] + bboxes[:, 3]) / 2 # y_center target_boxes[:, 2] = bboxes[:, 2] - bboxes[:, 0] # width target_boxes[:, 3] = bboxes[:, 3] - bboxes[:, 1] # height # 归一化 target_boxes[:, 0] /= image.shape[2] # 除以宽度 target_boxes[:, 1] /= image.shape[1] # 除以高度 target_boxes[:, 2] /= image.shape[2] target_boxes[:, 3] /= image.shape[1] targets = torch.cat([ class_labels.unsqueeze(1).float(), target_boxes ], dim=1) else: targets = torch.zeros((0, 5)) return image, targets

4. YOLO模型实现与训练

4.1 YOLOv5模型配置

python

# yolov5_tomato.yaml import yaml yolov5_config = { 'nc': 3, # 类别数量 'depth_multiple': 0.33, # 模型深度倍数 'width_multiple': 0.50, # 层宽度倍数 'anchors': [ [10,13, 16,30, 33,23], # P3/8 [30,61, 62,45, 59,119], # P4/16 [116,90, 156,198, 373,326] # P5/32 ], 'backbone': [ [-1, 1, 'Conv', [64, 6, 2, 2]], # 0-P1/2 [-1, 1, 'Conv', [128, 3, 2]], # 1-P2/4 [-1, 3, 'C3', [128]], [-1, 1, 'Conv', [256, 3, 2]], # 3-P3/8 [-1, 6, 'C3', [256]], [-1, 1, 'Conv', [512, 3, 2]], # 5-P4/16 [-1, 9, 'C3', [512]], [-1, 1, 'Conv', [1024, 3, 2]], # 7-P5/32 [-1, 3, 'C3', [1024]], [-1, 1, 'SPPF', [1024, 5]], # 9 ], 'head': [ [-1, 1, 'Conv', [512, 1, 1]], [-1, 1, 'nn.Upsample', [None, 2, 'nearest']], [[-1, 6], 1, 'Concat', [1]], # cat backbone P4 [-1, 3, 'C3', [512, False]], [-1, 1, 'Conv', [256, 1, 1]], [-1, 1, 'nn.Upsample', [None, 2, 'nearest']], [[-1, 4], 1, 'Concat', [1]], # cat backbone P3 [-1, 3, 'C3', [256, False]], [-1, 1, 'Conv', [256, 3, 2]], [[-1, 14], 1, 'Concat', [1]], # cat head P4 [-1, 3, 'C3', [512, False]], [-1, 1, 'Conv', [512, 3, 2]], [[-1, 10], 1, 'Concat', [1]], # cat head P5 [-1, 3, 'C3', [1024, False]], [[17, 20, 23], 1, 'Detect', ['nc', 'anchors']], # Detect(P3, P4, P5) ] } # 保存配置文件 with open('yolov5_tomato.yaml', 'w') as f: yaml.dump(yolov5_config, f)

4.2 YOLOv8模型训练

python

from ultralytics import YOLO import torch import os class YOLOv8TomatoTrainer: def __init__(self, model_size='m', num_classes=3, device='cuda' if torch.cuda.is_available() else 'cpu'): self.device = device self.num_classes = num_classes self.model_size = model_size self.model = None def initialize_model(self): """初始化YOLOv8模型""" # 可以选择预训练模型 model_name = f'yolov8{self.model_size}.pt' self.model = YOLO(model_name) # 修改模型为3个类别 self.model.model.nc = self.num_classes def prepare_data_config(self, data_dir, train_split=0.8): """准备数据配置文件""" data_config = { 'path': data_dir, 'train': 'images/train', 'val': 'images/val', 'test': 'images/test', 'nc': self.num_classes, 'names': ['fresh_tomato', 'semi_fresh_tomato', 'rotten_tomato'] } # 创建目录结构 for split in ['train', 'val', 'test']: os.makedirs(os.path.join(data_dir, 'images', split), exist_ok=True) os.makedirs(os.path.join(data_dir, 'labels', split), exist_ok=True) # 保存数据配置文件 import yaml config_path = os.path.join(data_dir, 'tomato_dataset.yaml') with open(config_path, 'w') as f: yaml.dump(data_config, f) return config_path def train(self, data_config, epochs=100, batch_size=16, img_size=640): """训练模型""" results = self.model.train( data=data_config, epochs=epochs, imgsz=img_size, batch=batch_size, device=self.device, workers=4, optimizer='AdamW', lr0=0.001, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, box=7.5, cls=0.5, dfl=1.5, close_mosaic=10, 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, copy_paste=0.0, save=True, save_period=-1, cache=False, name='yolov8_tomato_freshness', pretrained=True, verbose=True, seed=42, deterministic=True, single_cls=False, rect=False, cos_lr=False, label_smoothing=0.0, dropout=0.0, patience=100, freeze=None, save_dir='runs/detect', amp=True, fraction=1.0, profile=False, overlap_mask=True, mask_ratio=4, ) return results def evaluate(self, data_config): """评估模型""" metrics = self.model.val(data=data_config) return metrics def export(self, format='onnx'): """导出模型""" success = self.model.export(format=format) return success

4.3 YOLOv10模型实现

python

import torch import torch.nn as nn import torch.nn.functional as F class C2fV10(nn.Module): """YOLOv10中的C2f模块""" def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): super().__init__() self.c = int(c2 * e) self.cv1 = nn.Conv2d(c1, 2 * self.c, 1, 1, bias=False) self.cv2 = nn.Conv2d((2 + n) * self.c, c2, 1, bias=False) self.m = nn.ModuleList( BottleneckV10(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n) ) def forward(self, x): y = list(self.cv1(x).chunk(2, 1)) y.extend(m(y[-1]) for m in self.m) return self.cv2(torch.cat(y, 1)) class BottleneckV10(nn.Module): """YOLOv10瓶颈层""" def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5): super().__init__() c_ = int(c2 * e) self.cv1 = nn.Conv2d(c1, c_, k[0], 1, padding=k[0]//2, groups=g, bias=False) self.cv2 = nn.Conv2d(c_, c2, k[1], 1, padding=k[1]//2, groups=g, bias=False) self.add = shortcut and c1 == c2 def forward(self, x): return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) class PSA(nn.Module): """YOLOv10中的PSA注意力机制""" def __init__(self, c1, c2): super().__init__() self.cv1 = nn.Conv2d(c1, c2, 1) self.attention = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(c2, c2, 1), nn.Sigmoid() ) def forward(self, x): y = self.cv1(x) attention = self.attention(y) return x * attention + y class YOLOv10(nn.Module): """YOLOv10模型定义""" def __init__(self, nc=3): super().__init__() self.nc = nc # 骨干网络 self.backbone = nn.Sequential( nn.Conv2d(3, 64, 6, 2, 2), C2fV10(64, 64, 1), nn.Conv2d(64, 128, 3, 2, 1), C2fV10(128, 128, 2), nn.Conv2d(128, 256, 3, 2, 1), C2fV10(256, 256, 3), nn.Conv2d(256, 512, 3, 2, 1), C2fV10(512, 512, 1), PSA(512, 512), ) # 颈部网络 self.neck = nn.Sequential( nn.Conv2d(512, 256, 1), nn.Upsample(scale_factor=2), C2fV10(512, 256, 1), nn.Conv2d(256, 128, 1), nn.Upsample(scale_factor=2), C2fV10(256, 128, 1), ) # 检测头 self.detect = DetectV10(nc=nc) def forward(self, x): features = [] # 提取特征 for layer in self.backbone: x = layer(x) if isinstance(layer, (C2fV10, PSA)): features.append(x) # 特征金字塔 neck_features = [] x = features[-1] for layer in self.neck: x = layer(x) if isinstance(layer, C2fV10): neck_features.append(x) return self.detect(neck_features) class DetectV10(nn.Module): """YOLOv10检测头""" def __init__(self, nc=3, anchors=()): super().__init__() self.nc = nc self.no = nc + 4 # 每个锚点的输出维度 # 分类和回归层 self.cls = nn.ModuleList( nn.Sequential( nn.Conv2d(128, 128, 3, padding=1), nn.Conv2d(128, self.no, 1) ) for _ in range(3) ) self.reg = nn.ModuleList( nn.Sequential( nn.Conv2d(128, 128, 3, padding=1), nn.Conv2d(128, 4, 1) ) for _ in range(3) ) def forward(self, x): outputs = [] for i, feat in enumerate(x): cls_out = self.cls[i](feat) reg_out = self.reg[i](feat) output = torch.cat([reg_out, cls_out], dim=1) outputs.append(output) return outputs

5. 模型训练脚本

python

import torch from torch.utils.data import DataLoader import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR import torch.nn as nn import numpy as np from tqdm import tqdm import warnings warnings.filterwarnings('ignore') class YOLOTrainer: def __init__(self, model, train_dataset, val_dataset, device='cuda'): self.model = model.to(device) self.device = device self.train_dataset = train_dataset self.val_dataset = val_dataset # 超参数 self.epochs = 100 self.batch_size = 16 self.learning_rate = 0.001 self.weight_decay = 0.0005 self.momentum = 0.937 # 创建数据加载器 self.train_loader = DataLoader( train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4, pin_memory=True, collate_fn=self.collate_fn ) self.val_loader = DataLoader( val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=4, pin_memory=True, collate_fn=self.collate_fn ) # 优化器和调度器 self.optimizer = optim.AdamW( model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay ) self.scheduler = CosineAnnealingLR( self.optimizer, T_max=self.epochs, eta_min=0.0001 ) # 损失函数 self.loss_fn = YOLOLoss() def collate_fn(self, batch): """自定义批次处理函数""" images = [] targets = [] for img, target in batch: images.append(img) targets.append(target) images = torch.stack(images, dim=0) return images, targets def train_epoch(self, epoch): """训练一个epoch""" self.model.train() total_loss = 0 progress_bar = tqdm(self.train_loader, desc=f'Epoch {epoch+1}/{self.epochs}') for batch_idx, (images, targets) in enumerate(progress_bar): images = images.to(self.device) # 准备目标 target_list = [] for i, target in enumerate(targets): if len(target) > 0: target_with_img_idx = torch.cat([ torch.full((len(target), 1), i, device=self.device), target ], dim=1) target_list.append(target_with_img_idx) if len(target_list) > 0: targets_tensor = torch.cat(target_list, dim=0) else: targets_tensor = torch.zeros((0, 6), device=self.device) # 前向传播 self.optimizer.zero_grad() predictions = self.model(images) # 计算损失 loss = self.loss_fn(predictions, targets_tensor) # 反向传播 loss.backward() # 梯度裁剪 torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0) self.optimizer.step() total_loss += loss.item() progress_bar.set_postfix({'loss': loss.item()}) avg_loss = total_loss / len(self.train_loader) return avg_loss def validate(self): """验证模型""" self.model.eval() total_loss = 0 with torch.no_grad(): for images, targets in tqdm(self.val_loader, desc='Validating'): images = images.to(self.device) # 准备目标 target_list = [] for i, target in enumerate(targets): if len(target) > 0: target_with_img_idx = torch.cat([ torch.full((len(target), 1), i, device=self.device), target ], dim=1) target_list.append(target_with_img_idx) if len(target_list) > 0: targets_tensor = torch.cat(target_list, dim=0) else: targets_tensor = torch.zeros((0, 6), device=self.device) # 前向传播 predictions = self.model(images) # 计算损失 loss = self.loss_fn(predictions, targets_tensor) total_loss += loss.item() avg_loss = total_loss / len(self.val_loader) return avg_loss def train(self): """训练模型""" best_val_loss = float('inf') for epoch in range(self.epochs): # 训练 train_loss = self.train_epoch(epoch) # 验证 val_loss = self.validate() # 学习率调度 self.scheduler.step() # 保存最佳模型 if val_loss < best_val_loss: best_val_loss = val_loss torch.save({ 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'loss': val_loss, }, 'best_model.pth') print(f'Epoch {epoch+1}/{self.epochs}: ' f'Train Loss: {train_loss:.4f}, ' f'Val Loss: {val_loss:.4f}, ' f'LR: {self.scheduler.get_last_lr()[0]:.6f}') print(f'Best validation loss: {best_val_loss:.4f}') class YOLOLoss(nn.Module): """YOLO损失函数""" def __init__(self): super().__init__() self.mse_loss = nn.MSELoss() self.bce_loss = nn.BCEWithLogitsLoss() def forward(self, predictions, targets): """ 计算YOLO损失 predictions: 模型预测结果 targets: 真实标签 [batch_idx, class, x, y, w, h] """ # 这里简化了损失计算，实际YOLO损失更复杂 obj_loss = 0 box_loss = 0 cls_loss = 0 for pred in predictions: if len(pred) > 0 and len(targets) > 0: # 计算边界框损失 pred_boxes = pred[:, :4] target_boxes = targets[:, 2:6] box_loss += self.mse_loss(pred_boxes, target_boxes) # 计算类别损失 pred_cls = pred[:, 4:] target_cls = targets[:, 1] cls_loss += F.cross_entropy(pred_cls, target_cls.long()) total_loss = box_loss + cls_loss return total_loss

6. 用户界面实现

6.1 Gradio Web界面

python

import gradio as gr import torch import cv2 import numpy as np from PIL import Image import tempfile import os class TomatoFreshnessUI: def __init__(self, model_paths): """ 初始化UI界面 model_paths: 包含不同YOLO版本模型的字典 """ self.models = {} self.load_models(model_paths) def load_models(self, model_paths): """加载所有模型""" for model_name, model_path in model_paths.items(): if model_path.endswith('.pt'): if 'yolov5' in model_name: self.models[model_name] = self.load_yolov5(model_path) elif 'yolov8' in model_name: self.models[model_name] = self.load_yolov8(model_path) elif 'yolov10' in model_name: self.models[model_name] = self.load_yolov10(model_path) def load_yolov5(self, model_path): """加载YOLOv5模型""" try: from models.experimental import attempt_load model = attempt_load(model_path, device='cpu') return model except: # 简化版本 print(f"Loading YOLOv5 from {model_path}") return None def load_yolov8(self, model_path): """加载YOLOv8模型""" try: from ultralytics import YOLO model = YOLO(model_path) return model except: print(f"Loading YOLOv8 from {model_path}") return None def load_yolov10(self, model_path): """加载YOLOv10模型""" try: # 这里需要YOLOv10的具体实现 print(f"Loading YOLOv10 from {model_path}") return None except: return None def preprocess_image(self, image): """预处理图像""" if isinstance(image, str): image = cv2.imread(image) elif isinstance(image, np.ndarray): pass elif hasattr(image, 'pil'): image = np.array(image) # 调整大小为640x640 img_resized = cv2.resize(image, (640, 640)) # 归一化 img_normalized = img_resized / 255.0 # 转换为torch tensor img_tensor = torch.from_numpy(img_normalized).permute(2, 0, 1).float() return img_tensor.unsqueeze(0), image def predict(self, image, model_choice, confidence_threshold=0.5): """进行预测""" if model_choice not in self.models or self.models[model_choice] is None: return image, "模型未加载" model = self.models[model_choice] img_tensor, original_img = self.preprocess_image(image) # 执行推理 with torch.no_grad(): if 'yolov8' in model_choice: results = model(img_tensor, conf=confidence_threshold) detections = results[0] else: predictions = model(img_tensor) detections = self.process_predictions(predictions, confidence_threshold) # 绘制检测结果 result_img = self.draw_detections(original_img, detections, model_choice) # 统计结果 stats = self.calculate_statistics(detections) return result_img, stats def process_predictions(self, predictions, confidence_threshold): """处理预测结果""" detections = [] # 这里需要根据具体的模型输出格式进行处理 # 简化版本 return detections def draw_detections(self, image, detections, model_name): """在图像上绘制检测框""" img_with_boxes = image.copy() # 类别颜色映射 class_colors = { 0: (0, 255, 0), # 新鲜 - 绿色 1: (255, 255, 0), # 半新鲜 - 黄色 2: (255, 0, 0) # 不新鲜 - 红色 } class_names = ['新鲜番茄', '半新鲜番茄', '不新鲜番茄'] for detection in detections: if hasattr(detection, 'boxes'): # YOLOv8格式 boxes = detection.boxes for box in boxes: x1, y1, x2, y2 = box.xyxy[0].cpu().numpy() conf = box.conf[0].cpu().numpy() cls = int(box.cls[0].cpu().numpy()) color = class_colors.get(cls, (255, 255, 255)) label = f"{class_names[cls]}: {conf:.2f}" # 绘制边界框 cv2.rectangle(img_with_boxes, (int(x1), int(y1)), (int(x2), int(y2)), color, 2) # 绘制标签 cv2.putText(img_with_boxes, label, (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) return img_with_boxes def calculate_statistics(self, detections): """计算统计信息""" stats = { 'total': 0, 'fresh': 0, 'semi_fresh': 0, 'rotten': 0, 'freshness_rate': 0 } for detection in detections: if hasattr(detection, 'boxes'): boxes = detection.boxes for box in boxes: cls = int(box.cls[0].cpu().numpy()) stats['total'] += 1 if cls == 0: stats['fresh'] += 1 elif cls == 1: stats['semi_fresh'] += 1 elif cls == 2: stats['rotten'] += 1 if stats['total'] > 0: stats['freshness_rate'] = (stats['fresh'] + stats['semi_fresh'] * 0.5) / stats['total'] * 100 return stats def create_interface(self): """创建Gradio界面""" with gr.Blocks(title="番茄新鲜度检测系统", theme=gr.themes.Soft()) as demo: gr.Markdown("# 🍅 番茄新鲜度检测系统") gr.Markdown("使用YOLOv5/YOLOv8/YOLOv10模型检测番茄新鲜度") with gr.Row(): with gr.Column(): image_input = gr.Image(label="上传番茄图像", type="filepath") model_choice = gr.Dropdown( choices=["YOLOv5", "YOLOv8", "YOLOv10"], value="YOLOv8", label="选择模型" ) confidence_slider = gr.Slider( minimum=0.1, maximum=1.0, value=0.5, step=0.05, label="置信度阈值" ) predict_btn = gr.Button("开始检测", variant="primary") with gr.Column(): image_output = gr.Image(label="检测结果", type="numpy") stats_output = gr.JSON(label="检测统计") # 示例图像 gr.Markdown("### 示例图像") example_images = [ ["examples/fresh_tomatoes.jpg", "YOLOv8", 0.5], ["examples/mixed_tomatoes.jpg", "YOLOv8", 0.5], ["examples/rotten_tomatoes.jpg", "YOLOv8", 0.5] ] gr.Examples( examples=example_images, inputs=[image_input, model_choice, confidence_slider], outputs=[image_output, stats_output], fn=self.predict, cache_examples=False ) # 绑定事件 predict_btn.click( fn=self.predict, inputs=[image_input, model_choice, confidence_slider], outputs=[image_output, stats_output] ) # 添加说明 with gr.Accordion("使用说明", open=False): gr.Markdown(""" 1. 上传包含番茄的图像（支持JPG、PNG格式） 2. 选择使用的YOLO模型版本 3. 调整置信度阈值（推荐0.5） 4. 点击"开始检测"按钮 5. 查看检测结果和统计信息 **类别说明：** - 🟢 新鲜番茄：品质最佳，适合直接食用 - 🟡 半新鲜番茄：品质中等，适合烹饪使用 - 🔴 不新鲜番茄：品质较差，建议丢弃 """) return demo def main(): # 假设模型路径 model_paths = { "YOLOv5": "models/yolov5_tomato.pt", "YOLOv8": "models/yolov8_tomato.pt", "YOLOv10": "models/yolov10_tomato.pt" } # 创建UI实例 ui = TomatoFreshnessUI(model_paths) # 启动界面 demo = ui.create_interface() demo.launch( server_name="0.0.0.0", server_port=7860, share=True ) if __name__ == "__main__": main()

7. 模型评估与性能比较

7.1 评估指标计算

python

import numpy as np from sklearn.metrics import precision_recall_curve, average_precision_score import matplotlib.pyplot as plt class ModelEvaluator: def __init__(self, model, test_loader, device='cuda'): self.model = model self.test_loader = test_loader self.device = device self.results = {} def evaluate(self): """评估模型性能""" self.model.eval() all_predictions = [] all_targets = [] with torch.no_grad(): for images, targets in tqdm(self.test_loader, desc='Evaluating'): images = images.to(self.device) # 前向传播 predictions = self.model(images) # 处理预测结果 processed_preds = self.process_predictions(predictions) processed_targets = self.process_targets(targets) all_predictions.extend(processed_preds) all_targets.extend(processed_targets) # 计算各项指标 metrics = self.calculate_metrics(all_predictions, all_targets) return metrics def calculate_metrics(self, predictions, targets): """计算评估指标""" metrics = {} # 为每个类别计算指标 for class_id in range(3): # 3个类别 class_preds = [] class_targets = [] for pred, target in zip(predictions, targets): # 提取当前类别的预测和真实值 pred_mask = (pred['classes'] == class_id) target_mask = (target['classes'] == class_id) if len(pred['scores'][pred_mask]) > 0: class_preds.extend(pred['scores'][pred_mask].cpu().numpy()) class_targets.extend([1] * len(pred['scores'][pred_mask])) if len(target['classes'][target_mask]) > 0: # 添加负样本 class_preds.extend([0.0] * len(target['classes'][target_mask])) class_targets.extend([1] * len(target['classes'][target_mask])) if len(class_preds) > 0: # 计算AP ap = average_precision_score(class_targets, class_preds) metrics[f'AP_class_{class_id}'] = ap # 计算mAP aps = [metrics[f'AP_class_{i}'] for i in range(3) if f'AP_class_{i}' in metrics] if aps: metrics['mAP'] = np.mean(aps) return metrics def plot_pr_curve(self, predictions, targets, save_path='pr_curve.png'): """绘制PR曲线""" plt.figure(figsize=(10, 8)) for class_id, class_name in enumerate(['fresh', 'semi_fresh', 'rotten']): class_preds = [] class_targets = [] for pred, target in zip(predictions, targets): pred_mask = (pred['classes'] == class_id) target_mask = (target['classes'] == class_id) if len(pred['scores'][pred_mask]) > 0: class_preds.extend(pred['scores'][pred_mask].cpu().numpy()) class_targets.extend([1] * len(pred['scores'][pred_mask])) if len(target['classes'][target_mask]) > 0: class_preds.extend([0.0] * len(target['classes'][target_mask])) class_targets.extend([1] * len(target['classes'][target_mask])) if len(class_preds) > 0: precision, recall, _ = precision_recall_curve(class_targets, class_preds) ap = average_precision_score(class_targets, class_preds) plt.plot(recall, precision, lw=2, label=f'{class_name} (AP={ap:.3f})') plt.xlabel('Recall') plt.ylabel('Precision') plt.title('Precision-Recall Curve') plt.legend(loc='best') plt.grid(True) plt.savefig(save_path, dpi=300, bbox_inches='tight') plt.close()

7.2 性能比较

python

def compare_models(model_results): """比较不同模型的性能""" comparison_data = [] for model_name, results in model_results.items(): comparison_data.append({ 'Model': model_name, 'mAP': results.get('mAP', 0), 'AP_fresh': results.get('AP_class_0', 0), 'AP_semi_fresh': results.get('AP_class_1', 0), 'AP_rotten': results.get('AP_class_2', 0), 'Inference Time (ms)': results.get('inference_time', 0) }) # 创建性能对比表格 import pandas as pd df = pd.DataFrame(comparison_data) # 可视化对比 fig, axes = plt.subplots(1, 2, figsize=(15, 6)) # mAP对比 axes[0].bar(df['Model'], df['mAP']) axes[0].set_ylabel('mAP') axes[0].set_title('Model mAP Comparison') axes[0].tick_params(axis='x', rotation=45) # 推理时间对比 axes[1].bar(df['Model'], df['Inference Time (ms)']) axes[1].set_ylabel('Inference Time (ms)') axes[1].set_title('Model Inference Time Comparison') axes[1].tick_params(axis='x', rotation=45) plt.tight_layout() plt.savefig('model_comparison.png', dpi=300, bbox_inches='tight') return df

8. 部署与优化

8.1 ONNX导出与优化

python

def export_to_onnx(model, input_shape=(1, 3, 640, 640), onnx_path='model.onnx'): """导出模型为ONNX格式""" model.eval() # 创建示例输入 dummy_input = torch.randn(*input_shape) # 导出模型 torch.onnx.export( model, dummy_input, onnx_path, export_params=True, opset_version=12, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {0: 'batch_size'}, 'output': {0: 'batch_size'} } ) print(f"Model exported to {onnx_path}") # 验证ONNX模型 import onnx onnx_model = onnx.load(onnx_path) onnx.checker.check_model(onnx_model) return onnx_path def optimize_onnx_model(onnx_path, optimized_path='model_optimized.onnx'): """优化ONNX模型""" import onnxruntime as ort from onnxruntime.transformers import optimizer # 优化模型 optimized_model = optimizer.optimize_model( onnx_path, model_type='bert', # 这里需要根据模型类型调整 num_heads=8, hidden_size=512 ) optimized_model.save_model_to_file(optimized_path) return optimized_path

8.2 TensorRT加速

python

def convert_to_tensorrt(onnx_path, trt_path='model.trt', fp16=False): """转换为TensorRT引擎""" import tensorrt as trt logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) # 解析ONNX模型 with open(onnx_path, 'rb') as model: if not parser.parse(model.read()): for error in range(parser.num_errors): print(parser.get_error(error)) return None # 配置构建器 config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1GB if fp16 and builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) # 构建引擎 engine = builder.build_engine(network, config) # 保存引擎 with open(trt_path, 'wb') as f: f.write(engine.serialize()) return trt_path