最近在帮学弟学妹看毕业设计发现“带钢表面缺陷识别”这个选题热度一直很高。它确实是个很好的练手项目既有明确的工业应用背景又能串联起计算机视觉从数据处理到模型部署的完整链路。不过很多同学上手后才发现从论文到可运行的代码再到一个稳定可演示的系统中间有不少“坑”。今天我就结合自己的经验把这个项目的全流程拆解一遍希望能帮你理清思路高效完成毕设。1. 工业场景下的核心挑战理想很丰满现实很骨感别看任务描述就一句话“识别带钢表面的缺陷”真做起来你会发现原始数据就和实验室的MNIST、CIFAR-10完全不是一回事。主要挑战集中在以下几点图像质量不稳定这是最大的拦路虎。工厂环境复杂光照不均、钢板反光、相机抖动、灰尘干扰都会导致图像噪声大、对比度低。你可能拿到的数据集里有些图片亮得刺眼有些又暗得看不清纹理。缺陷形态多变且样本少划痕、孔洞、锈斑、夹杂等缺陷其大小、形状、位置千变万化。更头疼的是带钢生产质量其实很高缺陷样本本身就是“少数派”这就导致了严重的类别不平衡。你的数据集里可能90%都是正常图片剩下的10%才包含各类缺陷而且某些稀有缺陷的图片可能只有寥寥几张。背景干扰与正负样本模糊带钢表面的正常纹理如轧制纹路有时和轻微划痕很像机器容易“看走眼”。如何让模型学会区分“背景纹理”和“真实缺陷”是个细活儿。实时性要求虽然毕设可能不要求实时但了解这一点很重要。工业流水线速度很快你的算法必须在几十到几百毫秒内完成一张图的检测否则没有实用价值。2. 技术路线选择没有最好只有最合适面对这些挑战我们有哪些武器呢主要分两大流派传统图像处理 vs 深度学习传统方法如OpenCV滤波、形态学操作、阈值分割优点是原理直观、计算速度快、不依赖大量数据。对于对比度强烈、特征明显的缺陷比如一个大孔洞用边缘检测Canny或阈值分割可能很快就能搞定。但缺点也很明显鲁棒性差。光照一变参数可能就失效了对于纹理复杂、缺陷细微的情况手工设计特征如同大海捞针。深度学习方法主要是卷积神经网络CNN这是当前的主流和毕业设计的首选。它能自动从数据中学习多层次的特征表达对噪声、形变有一定的容忍度泛化能力更强。你的核心工作从“设计特征”变成了“准备数据、调参和部署”。深度学习模型内部的选择分类网络 vs 检测网络确定了用深度学习下一个问题就是用分类模型还是检测模型图像分类网络如ResNet, MobileNet, EfficientNet任务是把整张图分为“正常”或“某类缺陷”。这是最简单的切入点。如果你的数据已经是裁剪好的、只包含单个缺陷的小图块或者你只关心图片里“有没有缺陷”不关心位置那用分类模型就足够了。它的优点是结构简单、训练快容易上手。目标检测网络如YOLO系列, SSD, Faster R-CNN任务不仅要知道有没有缺陷还要用框Bounding Box标出缺陷在哪里。这对于需要定位缺陷位置、统计缺陷数量的场景是必须的。YOLO以其速度快著称非常适合向“实时检测”方向拓展。但相应地数据标注成本高需要画框模型也更复杂。给毕设的建议初期强烈建议从图像分类任务开始。先用分类网络比如轻量级的MobileNetV3跑通整个Pipeline解决“有无缺陷”的问题。这能让你快速验证想法积累对数据的感觉。等分类任务效果不错了如果时间和精力允许再升级到YOLOv5/v8做检测这会成为你毕设的一个显著亮点。3. 基于PyTorch的实战代码与思路详解接下来我们进入实战环节。假设我们选择MobileNetV3-small作为分类模型任务是将图像分为“正常(Normal)”、“划痕(Scratch)”、“孔洞(Hole)”三类。3.1 数据准备与增强策略数据是模型的“粮食”这一步没做好后面再怎么调参都白搭。数据加载使用PyTorch的Dataset和DataLoader。关键是要处理好类别不平衡。我们可以用WeightedRandomSampler给少数类样本更高的采样权重。import torch from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler from PIL import Image import os from torchvision import transforms import numpy as np class SteelDefectDataset(Dataset): def __init__(self, root_dir, transformNone): self.root_dir root_dir self.transform transform self.classes [normal, scratch, hole] self.class_to_idx {c: i for i, c in enumerate(self.classes)} self.samples [] # 存储(图片路径, 标签索引) # 遍历文件夹收集所有样本 for class_name in self.classes: class_dir os.path.join(root_dir, class_name) if not os.path.isdir(class_dir): continue for img_name in os.listdir(class_dir): if img_name.endswith((.jpg, .png, .bmp)): img_path os.path.join(class_dir, img_name) self.samples.append((img_path, self.class_to_idx[class_name])) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] image Image.open(img_path).convert(RGB) # 统一转为RGB三通道 if self.transform: image self.transform(image) return image, label # 计算每个类别的权重用于平衡采样 def get_class_weights(dataset): class_counts np.bincount([label for _, label in dataset.samples]) total len(dataset) num_classes len(dataset.classes) # 权重与样本数成反比 class_weights total / (num_classes * class_counts.astype(float)) return torch.tensor(class_weights, dtypetorch.float) # 定义数据增强和预处理 train_transform transforms.Compose([ transforms.Resize((256, 256)), # 统一缩放到256x256 transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转 transforms.RandomRotation(degrees10), # 随机旋转 transforms.ColorJitter(brightness0.2, contrast0.2), # 模拟光照变化 transforms.ToTensor(), # 转为Tensor并归一化到[0,1] transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) # ImageNet通用均值方差 ]) val_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 创建数据集和数据加载器 train_dataset SteelDefectDataset(data/train, transformtrain_transform) val_dataset SteelDefectDataset(data/val, transformval_transform) # 创建加权采样器解决类别不平衡 class_weights get_class_weights(train_dataset) sample_weights [class_weights[label] for _, label in train_dataset.samples] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_size32, samplersampler) # 训练集用采样器 val_loader DataLoader(val_dataset, batch_size32, shuffleFalse) # 验证集不用关键点数据增强是应对光照不均、样本少的利器。ColorJitter模拟光照变化RandomRotation和Flip增加缺陷的多样性。注意验证集不需要做随机增强。3.2 模型训练与验证这里我们使用预训练的MobileNetV3-small并微调其最后一层。import torch.nn as nn import torch.optim as optim from torchvision import models from tqdm import tqdm # 用于显示进度条 device torch.device(cuda if torch.cuda.is_available() else cpu) # 1. 加载预训练模型并替换分类头 model models.mobilenet_v3_small(pretrainedTrue) num_ftrs model.classifier[-1].in_features # 获取原模型最后一层输入特征数 model.classifier[-1] nn.Linear(num_ftrs, 3) # 替换为我们的3分类层 model model.to(device) # 2. 定义损失函数和优化器 # 由于使用了加权采样器这里用普通的CrossEntropyLoss即可。 # 如果没用采样器可以使用带权重的损失函数criterion nn.CrossEntropyLoss(weightclass_weights.to(device)) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001, weight_decay1e-4) # 加入权重衰减防止过拟合 scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) # 学习率衰减 # 3. 训练循环 num_epochs 30 best_val_acc 0.0 for epoch in range(num_epochs): # 训练阶段 model.train() running_loss 0.0 correct 0 total 0 pbar tqdm(train_loader, descfEpoch {epoch1}/{num_epochs} [Train]) for images, labels in pbar: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() pbar.set_postfix({Loss: loss.item(), Acc: correct/total}) train_loss running_loss / len(train_dataset) train_acc correct / total # 验证阶段 model.eval() val_running_loss 0.0 val_correct 0 val_total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) val_running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) val_total labels.size(0) val_correct (predicted labels).sum().item() val_loss val_running_loss / len(val_dataset) val_acc val_correct / val_total print(fEpoch {epoch1}: Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}) # 保存最佳模型 if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f - Best model saved with val_acc: {val_acc:.4f}) scheduler.step() # 更新学习率 print(Training Finished.)关键点使用预训练模型能极大加速收敛。训练时观察验证集准确率并保存最佳模型防止过拟合。tqdm进度条能让训练过程更直观。4. 模型部署从PyTorch到Flask API模型训练好了怎么让别人能用上呢我们需要把它变成一个服务。这里分两步模型导出和API封装。4.1 导出为ONNX格式ONNX是一种开放的模型格式可以让你的模型脱离PyTorch环境被多种推理引擎如ONNX Runtime, TensorRT调用速度往往更快。import torch.onnx # 加载训练好的模型 model.load_state_dict(torch.load(best_model.pth)) model.eval() # 创建一个示例输入张量尺寸需与训练时一致 dummy_input torch.randn(1, 3, 256, 256).to(device) # 导出模型 onnx_path steel_defect_mobilenetv3.onnx torch.onnx.export( model, # 要导出的模型 dummy_input, # 模型输入示例 onnx_path, # 保存路径 export_paramsTrue, # 导出模型参数 opset_version11, # ONNX算子集版本 do_constant_foldingTrue, # 优化常量 input_names[input], # 输入节点名 output_names[output], # 输出节点名 dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} # 支持动态批次 ) print(fModel has been exported to {onnx_path})4.2 使用Flask创建REST API现在我们用Flask这个轻量级Web框架创建一个接收图片并返回预测结果的API。# app.py from flask import Flask, request, jsonify import onnxruntime as ort import numpy as np from PIL import Image import io import torchvision.transforms as transforms app Flask(__name__) # 1. 加载ONNX模型 onnx_model_path steel_defect_mobilenetv3.onnx ort_session ort.InferenceSession(onnx_model_path) # 2. 定义与训练时相同的预处理 transform transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 3. 类别映射 idx_to_class {0: normal, 1: scratch, 2: hole} def preprocess_image(image_bytes): 将上传的图片字节流转换为模型需要的张量 image Image.open(io.BytesIO(image_bytes)).convert(RGB) image_tensor transform(image).unsqueeze(0) # 增加批次维度 return image_tensor.numpy() # 转为numpy供ONNX Runtime使用 app.route(/predict, methods[POST]) def predict(): # 4. 输入校验 if file not in request.files: return jsonify({error: No file part}), 400 file request.files[file] if file.filename : return jsonify({error: No selected file}), 400 # 5. 处理图片并进行推理 try: img_bytes file.read() input_array preprocess_image(img_bytes) # ONNX Runtime推理 ort_inputs {ort_session.get_inputs()[0].name: input_array} ort_outputs ort_session.run(None, ort_inputs) predictions ort_outputs[0] # 获取预测结果 predicted_idx np.argmax(predictions, axis1)[0] predicted_class idx_to_class[predicted_idx] confidence float(np.max(predictions)) return jsonify({ defect_type: predicted_class, confidence: confidence, all_probs: predictions[0].tolist() # 返回所有类别的概率便于分析 }) except Exception as e: # 6. 异常处理 app.logger.error(fPrediction error: {str(e)}) return jsonify({error: Failed to process image}), 500 if __name__ __main__: # 生产环境应使用WSGI服务器如gunicorn而非直接app.run app.run(host0.0.0.0, port5000, debugFalse) # 生产环境务必设置debugFalse运行这个Flask应用后你就可以通过发送HTTP POST请求到http://服务器IP:5000/predict并附上图片文件来获取缺陷识别结果了。5. 性能评估与安全考量一个完整的项目不能只关心准确率。性能评估准确率/召回率/F1-score对于类别不平衡的数据只看整体准确率会失真。一定要计算每个类别的精确率(Precision)、召回率(Recall)和F1-score绘制混淆矩阵看看模型到底在哪些类别上容易出错。推理速度使用time模块在CPU和GPU上分别测试ONNX模型处理单张图片和批量图片的耗时。这是评估能否满足实时性的关键。ONNX Runtime通常比原生PyTorch推理更快。mAP平均精度均值如果你最终做的是目标检测YOLO那么mAP是核心指标。它综合反映了模型在不同置信度阈值下的检测精度。安全性考量输入校验如上文代码所示必须检查请求中是否包含文件、文件是否为空、文件格式是否为允许的图片格式如jpg, png。防止恶意上传非图片文件导致服务崩溃。异常处理用try...except包裹核心推理代码捕获并记录可能出现的错误如图片损坏、预处理失败、模型推理错误给客户端返回友好的错误信息而不是内部堆栈跟踪。文件大小限制在Flask中可以通过app.config[MAX_CONTENT_LENGTH]限制上传文件大小防止DoS攻击。API幂等性对于预测接口多次发送同一张图片应该得到相同的结果。我们的模型是确定性的这一点自然满足。6. 生产环境避坑指南这些是你在实验室可能遇不到但一旦部署就可能“翻车”的点。避免过拟合小数据集工业数据获取难。除了数据增强还可以尝试迁移学习只微调网络最后几层而不是全部参数。早停法Early Stopping监控验证集损失当它不再下降时停止训练。交叉验证在小数据集上尤其有用能更稳健地评估模型性能。警惕GPU内存泄漏在训练或部署服务中如果长时间运行后内存占用不断增长可能是由于张量或变量没有正确释放。确保在推理完成后及时清理不必要的中间变量。在Flask中确保推理部分没有在全局累积数据。确保API的健壮性与幂等性如上文所述做好校验和异常处理。另外考虑使用队列如Redis来处理高并发请求避免一个耗时长的请求阻塞整个服务。模型版本管理当你优化模型后会产生新版本。在API设计时可以考虑加入模型版本号方便回滚和AB测试。监控与日志记录每一次预测的请求时间、耗时、结果和置信度。这不仅能帮你分析性能瓶颈还能在出现误检时快速定位问题数据。写在最后走完这一整套流程一个基本的带钢表面缺陷识别项目就算完成了。但这仅仅是开始。你可以思考如何将这个“单张图片分类”的项目进行扩展让它更贴近真实工业应用扩展至实时视频流检测利用OpenCV读取摄像头或视频流逐帧调用你的Flask API或直接加载模型进行推理。这里的关键是优化推理速度考虑使用TensorRT进一步加速和处理帧率匹配问题。集成到MES制造执行系统这需要你的系统能提供标准化的接口如RESTful API或消息队列将缺陷类型、位置、置信度、时间戳等信息按照MES要求的格式上报。同时可能还需要与数据库联动保存历史检测记录用于质量追溯和分析。主动学习与迭代优化将模型在实际场景中置信度低的样本难例保存下来人工标注后加入训练集重新训练模型形成闭环让模型在实际使用中越变越强。毕业设计的目的不仅是完成一个任务更是通过这个过程学会如何将一个工业问题拆解、用技术手段解决、并最终工程化落地的完整思维。希望这篇笔记能为你扫清一些障碍祝你毕设顺利