基于 Mask R-CNN+ PyQt5的药片缺陷分割检测系统

一、项目概述本项目实现了一个基于 Mask R-CNN 深度学习模型的药片缺陷检测系统能够自动识别和定位药片表面的污染contamination和裂纹crack两类缺陷。系统提供了完整的训练流程、GUI 可视化界面和推理工具适用于药品质量检测场景。二、系统架构2.1 核心技术栈深度学习框架: PyTorch torchvision模型架构: Mask R-CNN (ResNet-50-FPN backbone)数据增强: Albumentations界面开发: PyQt5预训练权重: COCO 数据集2.2 数据集结构magnesium/ ├── train/ │ ├── images/ # 420 张训练图片 │ └── labels/ # LabelMe JSON 标注 ├── val/ │ ├── images/ # 86 张验证图片 │ └── labels/ └── test/ ├── images/ # 7 张测试图片 └── labels/数据集包含两类缺陷污染 (contamination): 药片表面异物裂纹 (crack): 药片表面裂痕三、核心技术实现3.1 数据预处理LabelMe JSON 标注解析系统使用 OpenCV 将多边形标注转换为实例分割 maskdef _polygon_to_mask(self, points, img_shape): 将多边形顶点转换为二值 mask mask np.zeros(img_shape, dtypenp.uint8) points_array np.array(points, dtypenp.int32) cv2.fillPoly(mask, [points_array], 1) return mask3.2 数据增强策略使用 Albumentations 库进行多样化增强def get_train_transforms(): return A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.Rotate(limit15, p0.5), A.RandomBrightnessContrast(p0.3), A.GaussianBlur(blur_limit3, p0.2), A.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ToTensorV2() ], bbox_paramsA.BboxParams(formatpascal_voc, label_fields[labels]))3.3 模型设计Mask R-CNN 架构def get_model_maskrcnn(num_classes, pretrainedTrue, trainable_backbone_layers3): 构建 Mask R-CNN 模型 # 加载预训练模型 model torchvision.models.detection.maskrcnn_resnet50_fpn( pretrainedpretrained, trainable_backbone_layerstrainable_backbone_layers ) # 替换分类头3 类背景 污染 裂纹 in_features model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor FastRCNNPredictor(in_features, num_classes) # 替换 mask 预测头 in_features_mask model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor MaskRCNNPredictor( in_features_mask, 256, num_classes ) return model架构特点骨干网络ResNet-50 FPN特征金字塔网络RPN区域提议网络生成候选框RoI Align 提取精确特征双分支输出边界框分类 实例分割 mask四、 训练策略4.1 验证损失计算Mask R-CNN 在eval()模式下只返回预测结果无法获取损失。解决方案是在验证时保持train()模式但禁用梯度计算def validate_one_epoch(model, data_loader, device, epoch, total_epochs): 验证一个 epoch model.train() # 保持 train 模式以计算 loss total_loss 0.0 with torch.no_grad(): # 不计算梯度 for images, targets in data_loader: images [img.to(device) for img in images] targets [{k: v.to(device) for k, v in t.items()} for t in targets] # 计算损失模型在 train 模式下才返回 loss_dict loss_dict model(images, targets) losses sum(loss for loss in loss_dict.values()) total_loss losses.item() return total_loss / len(data_loader)4.2 训练配置优化器: SGD (momentum0.9, weight_decay0.0005) 学习率: 0.003 调度器: CosineAnnealingLR 批次大小: 2 训练轮数: 24 可训练 backbone 层数: 3五、 GUI 可视化训练系统实现了功能完整的 PyQt5 训练界面图5-1 训练界面系统界面包含参数配置区学习率、批次大小、训练轮数等超参数设置训练控制开始/停止按钮、进度条实时日志显示训练过程输出损失曲线动态更新训练和验证损失混淆矩阵实时展示分类性能核心实现使用多线程避免界面卡顿class TrainingThread(QThread): 训练线程 output_signal pyqtSignal(str) progress_signal pyqtSignal(int, int) loss_signal pyqtSignal(float, float) finished_signal pyqtSignal(bool, str) def run(self): # 启动训练子进程 self.process subprocess.Popen( self.command, stdoutsubprocess.PIPE, stderrsubprocess.STDOUT, textTrue, encodingutf-8 ) # 解析输出流发送信号更新界面 for line in iter(self.process.stdout.readline, ): if Epoch in line: self.progress_signal.emit(current_epoch, total_epochs) self.output_signal.emit(line)六、缺陷检测推理6.1 NMS 去重机制推理时使用非极大值抑制NMS去除重复检测框# 推理时使用 NMS 去除重复检测框 nms_keep torchvision.ops.nms( boxes, scores, iou_threshold0.5 ) boxes boxes[nms_keep] scores scores[nms_keep] labels labels[nms_keep] masks masks[nms_keep]七、可视化结果图7-1 污染检测边界框图7-2 污染检测mask系统在检测到药片表面的污染缺陷后同时输出红色边界框标注缺陷位置红色 mask 标注缺陷精确轮廓置信度分数0.99 表示 99% 确信度图7-3 污染检测轮廓图7-4 污染检测填充 mask使用skimage.measure.find_contours提取 mask 轮廓from skimage import measure # 提取轮廓 contours measure.find_contours(mask, 0.5) for contour in contours: plt.plot(contour[:, 1], contour[:, 0], linewidth2, colorred)图7-5 裂纹检测裂纹缺陷检测结果蓝色区域表示检测到的药片表面裂纹边界框和置信度 1.00 表示模型对该检测非常确定。八、训练结果分析8.1 损失曲线图8-1 损失曲线系统自动生成训练和验证损失曲线展示模型收敛过程def plot_loss_curve(train_losses, val_losses, output_path): 绘制损失曲线 plt.figure(figsize(10, 6)) epochs range(1, len(train_losses) 1) plt.plot(epochs, train_losses, b-, label训练损失, linewidth2) plt.plot(epochs, val_losses, r-, label验证损失, linewidth2) plt.xlabel(训练轮次 (Epoch), fontsize12) plt.ylabel(损失值 (Loss), fontsize12) plt.title(训练和验证损失曲线, fontsize14, fontweightbold) plt.legend() plt.grid(True, alpha0.3) plt.savefig(output_path, dpi300)九、性能指标最优训练结果lr0.003, cosine scheduler, 24 epochs训练时间: 16 分 36 秒最佳验证损失: 0.2740第 21 轮训练损失: 0.7505 → 0.3062模型参数量: 约 4400 万44M可训练参数: 约 3500 万十、错误分析系统实现了完整的错误分析功能def analyze_detection_errors(model, data_loader, device): 分析检测错误并生成混淆矩阵 error_stats { missed: 0, # 漏检真实缺陷未检测到 false_positive: 0, # 过检误报的缺陷 misclassified: 0 # 误判类别错误 } # 计算 IoU 匹配 for pred_box, gt_box in zip(pred_boxes, gt_boxes): iou compute_iou(pred_box, gt_box) if iou iou_threshold: if pred_label ! gt_label: error_stats[misclassified] 1 else: error_stats[false_positive] 1 return error_stats错误样本自动保存到output/error_analysis/目录便于后续分析优化。十一、项目文件结构Mask R-CNN/ ├── dataset.py # 数据集加载和预处理 ├── model.py # Mask R-CNN 模型定义 ├── train.py # 命令行训练脚本 ├── train_gui.py # GUI 训练界面PyQt5 ├── visualize.py # 推理和可视化工具 ├── utils.py # 工具函数checkpoint、seed 等 ├── requirements.txt # 依赖包列表 ├── magnesium/ # 数据集目录 └── output/ # 训练输出 ├── best_model.pth # 最佳模型权重 ├── final_model.pth # 最终模型权重 ├── loss_curve.png # 损失曲线图 ├── confusion_matrix.png # 混淆矩阵 ├── training_log.txt # 训练日志 └── error_analysis/ # 错误样本分析 ├── missed_detection/ ├── false_positive/ └── misclassified/十二、总结项目实现了一个功能完整的药片缺陷检测系统结合了深度学习的强大性能和工程化的实用性。通过 Mask R-CNN 模型系统能够精确定位和分割药片表面的污染和裂纹缺陷验证集损失达到 0.2740展现了良好的检测性能。GUI 界面和详细的错误分析功能使得系统易于使用和优化适合在实际药品质量检测场景中应用。十三、视频演示Mask R-CNN药片