HRNET模型实现钢板表面缺陷检测

深度学习实战训练营

于 2024-10-09 09:00:00 发布

阅读量223

点赞数 6

文章标签：深度学习

本文链接：https://blog.csdn.net/2401_87275147/article/details/142764865

版权

关于深度实战社区

我们是一个深度学习领域的独立工作室。团队成员有：中科大硕士、纽约大学硕士、浙江大学硕士、华东理工博士等，曾在腾讯、百度、德勤等担任算法工程师/产品经理。全网20多万+粉丝，拥有2篇国家级人工智能发明专利。

社区特色：深度实战算法创新

获取全部完整项目数据集、代码、视频教程，请进入官网：zzgcz.com。竞赛/论文/毕设项目辅导答疑，v：zzgcz_com

1. 项目简介

该项目旨在利用深度学习模型进行医学影像分割与图像处理。主要涉及使用UNet和HRNet两种经典的图像分割模型来进行图像特征提取与分割，从而实现对特定目标区域（如病灶、组织边界等）的精准识别。UNet模型通过其对称的编码器和解码器结构，使得模型能够捕捉到细粒度的特征信息，适用于较小样本量的医学图像数据。而HRNet模型通过引入高分辨率分支结构，使得模型在捕捉全局和局部特征时具有更高的表现力，尤其适用于处理复杂结构的医学图像。本项目的应用场景主要为医学图像分析与病灶分割，能够辅助医学专家提高诊断效率与准确度。项目通过多种可视化方案，如生成的分割图、特征可视化图等，帮助用户更好地理解模型的输出结果，并为后续的临床应用提供数据支持。
在这里插入图片描述

2.技术创新点摘要

模型架构设计与优化：本项目使用了经典的HRNet与UNet模型进行钢材缺陷检测与分割，并在此基础上引入多尺度特征融合策略。HRNet通过保持高分辨率的同时逐步进行多尺度特征聚合，能够在复杂背景下实现更高精度的缺陷边缘识别，而UNet利用其对称编码器-解码器结构，在分割任务中有效地捕获目标区域的细节信息。通过组合两种模型的优点，本项目能够在精度和效率之间取得良好平衡。
数据预处理与增强策略：该项目在数据预处理过程中应用了自定义的图像裁剪、旋转、翻转以及伪造缺陷的策略，以增强模型对不同缺陷类型的泛化能力。同时，为了应对类别不平衡问题，使用了数据重采样技术，通过控制训练样本的分布，进一步提升了模型对小样本缺陷的识别效果。
PaddleX优化加速方案：项目使用了百度飞桨（PaddleX）深度学习平台进行模型开发与训练，结合飞桨的高效计算库与分布式训练框架，在大规模数据集上实现了训练效率的提升。此外，利用飞桨特有的动态图模式进行模型调试，使得开发者可以更灵活地进行模型结构调整与创新设计。
高效的推理与部署方案：为了解决实际工业场景下的实时检测需求，项目在模型推理阶段使用了量化与裁剪技术，显著降低了模型的参数量与计算成本，使得模型能够部署在资源受限的边缘设备上，满足钢材生产线的在线缺陷检测需求。

3. 数据集与预处理

本项目使用的数据集来源于Kaggle举办的“Severstal: Steel Defect Detection”竞赛数据集，专用于钢材表面缺陷检测。数据集包含高分辨率的平板钢材图像及其对应的缺陷标注信息。图像数据被分为训练集和测试集：训练集包含若干张带有标注的钢材图像，而测试集则不提供标注。训练集中的每张图像可能包含四类不同的缺陷类型，包括划痕、锈斑、凹陷等，类别分布不均衡，且有的图像不包含缺陷（正常样本），因此类别不平衡是模型训练的一个主要挑战。

在数据预处理环节，首先对图像数据进行尺寸调整和格式转换，以便统一输入模型的尺寸规格（通常使用256×256或512×512的尺寸）。然后进行标准化（归一化）处理，使得每个像素值在[0,1]或[-1,1]区间内，确保输入数据具有相同的分布特征，有助于提升模型的收敛速度和稳定性。

为了提升模型的泛化能力，项目引入了多种数据增强策略，包括随机裁剪、水平与垂直翻转、旋转、亮度与对比度调整等。通过数据增强生成更多具有多样性特征的样本，从而有效缓解训练过程中过拟合的问题。此外，项目还使用了特征工程方法，如伪造小样本缺陷与随机噪声填充，进一步增加了模型对不同缺陷的辨识能力。

最终，通过系统化的数据预处理与特征工程流程，确保了模型在不同缺陷类型和复杂背景下都能保持较高的检测精度和稳定性。

4. 模型架构

模型结构的逻辑

本项目采用了两种经典的深度学习图像分割模型：HRNet（High-Resolution Network）和UNet。模型的整体结构设计如下：

UNet模型： UNet模型主要用于医学图像分割任务，其结构包含对称的编码器和解码器。编码器（下采样部分）逐步提取图像的特征信息，而解码器（上采样部分）逐步恢复图像空间分辨率，并结合编码器的特征信息，实现精细分割。其数学表示如下：
1. 编码器层： $C_i = \text{Conv}_{i} (C_{i-1})$ 其中，Convi 代表第i层的卷积操作，Ci 为第i层特征图。
2. 下采样操作： $D_i = \text{DownSample}(C_i)$ Di 为第i层下采样后的特征图。
3. 解码器层： $U_i = \text{UpSample}(D_{i+1}) + C_i$ 其中，Ui 为解码层的上采样特征图，通过跳跃连接将编码层的特征图与上采样后的特征图进行融合。
HRNet模型： HRNet通过并行的多分辨率卷积网络保持高分辨率的特征流，并在网络的每个阶段进行多尺度特征融合。模型主要由以下几部分组成：
1. 高分辨率分支（High-Resolution Stream）：每个高分辨率分支包含多个卷积层，保持输入图像的空间分辨率，并在每个阶段引入不同分辨率的特征图。
2. 特征聚合模块：使用跨尺度的特征融合操作： $F_i = \sum_{j=1}^{N} W_{ij} \times S_{j}$ 其中，Sj 表示第j尺度的特征图，Wij 表示跨尺度融合的权重矩阵。
3. 输出层：通过1×1卷积层将特征图映射到目标类别空间，生成最终的分割结果。

模型的整体训练流程

数据集划分： 数据集被分为训练集和验证集。训练集用于模型参数优化，而验证集用于监控模型性能。

数据预处理： 进行数据标准化（mean = [0.5, 0.5, 0.5]，std = [0.5, 0.5, 0.5]）、随机水平翻转等数据增强操作。

模型训练： 使用交叉熵损失函数（Cross-Entropy Loss）来度量模型的预测与真实标注之间的差异。优化算法采用Adam优化器，并使用动态学习率调整策略来提升模型的收敛效果。

评估指标： 使用**mIoU（Mean Intersection over Union）**作为评估指标，mIoU的计算公式如下：

$\text{mIoU} = \frac{1}{N} \sum_{i=1}^{N} \frac{A_i \cap B_i}{A_i \cup B_i}$

其中，Ai 和 Bi 分别表示第i个类别的预测区域和真实区域。

模型保存与可视化： 在每个训练周期后保存模型，并使用VisualDL工具对训练过程中的损失、学习率及mIoU进行可视化分析，以便调整模型参数和结构。

5. 核心代码详细讲解

1. 数据预处理与特征工程（ `prepare_dataset.py` ）

核心代码功能：该模块主要用于将Kaggle提供的训练数据进行预处理，并生成对应的Mask（分割标注图）文件，为模型训练做好数据准备。

import os
import numpy as np
import pandas as pd
import shutil
import cv2

导入库： os 用于文件路径操作，numpy 进行数值计算，pandas 读取数据文件，shutil 用于文件操作，cv2 处理图像数据。

def name_and_mask(start_idx):

定义函数 name_and_mask ：该函数用于根据数据索引生成每个图像对应的文件名及其Mask（标注）图像。

    img_names = [str(i).split("_")[0] for i in train_df.iloc[col:col + 4, 0].values]

获取图像名称： 从 train.csv 文件中读取图像文件名，提取前缀作为该图像的唯一标识符。

    if not (img_names[0] == img_names[1] == img_names[2] == img_names[3]):raise ValueError

数据一致性检查： 确保同一图像的标注数据是4行连续的（每种缺陷类型一行）。如果前缀不一致，则说明数据格式有误。

    labels = train_df.iloc[col:col + 4, 1]
    mask = np.zeros((256, 1600), dtype=np.uint8)

创建初始的 mask 图： 256×1600 是图像分辨率，初始值全为0（表示无缺陷区域）。

    for idx, label in enumerate(labels.values):if label is not np.nan:
            label = label.split(" ")
            positions = map(int, label[0::2])
            length = map(int, label[1::2])

解析标注数据： positions 表示缺陷在图像中的起始位置，length 表示该缺陷的长度。

            for pos, le in zip(positions, length):
                mask_label[pos - 1:pos + le - 1] = idx + 1

生成分割标注图： 将缺陷区域在 mask_label 中标记为相应的类别ID（从1开始）。

    mask[:, :] = mask_label.reshape(256, 1600, order='F')return img_names[0], mask

转换 mask_label 为二维矩阵： 重新排列标注数据，使其与图像大小匹配，order='F' 代表按列填充数据。

2. 模型推理与评估（ `infer.py` ）

核心代码功能：该模块用于加载训练好的模型，进行缺陷分割预测，并生成可视化的预测结果图。

import paddlex as pdx
import os
import cv2

导入库： paddlex 是飞桨（PaddlePaddle）的高层API库，用于深度学习任务，os 进行路径操作，cv2 处理图像文件。

pic_index = 25
test_path = 'steel/test_list.txt'
f = open(test_path)
lines = f.readlines()

获取测试图像文件名： 从 test_list.txt 中读取第 pic_index 行的测试图像路径，准备进行模型推理。

imgname = os.path.basename(lines[pic_index].split(' ')[0])
labelname = os.path.basename(lines[pic_index].split(' ')[1]).replace('\n', '')
f.close()

提取图像和标签文件名： 解析测试图像路径，并去除标签文件名中的换行符，保证路径格式一致。

image_path = os.path.join('steel/JPEGImages', imgname)
model = pdx.load_model('output/hrnet/best_model')

加载模型： 使用 pdx.load_model 加载训练好的HRNet模型，模型保存在 output/hrnet/best_model 中。

result = model.predict(image_path)
pdx.seg.visualize(image_path, result, weight=0.4, save_dir='output/predict')

模型推理与可视化：
- 推理： 使用 model.predict 对图像进行分割预测。
- 可视化： 通过 pdx.seg.visualize 将预测结果叠加在原图上，weight=0.4 表示预测图透明度为40%，保存结果至 output/predict 目录中。

label_path = os.path.join('steel/Annotations', labelname)
mask = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
gt = {'label_map': mask}
pdx.seg.visualize(image_path, gt, weight=0.4, save_dir='output/gt')

可视化真实标签： 使用 cv2.imread 加载真实标签 mask，并生成ground truth图像，保存至 output/gt 目录中。

6. 模型优缺点评价

优点：

高分辨率特征保留： HRNet模型在不同分辨率的特征流中保持高分辨率信息，有效捕捉了钢材表面细微的缺陷特征，提升了分割的精度，特别是在复杂背景中表现优异。
多尺度特征融合： HRNet模型通过多尺度融合策略实现了不同分辨率特征的高效聚合，从而能够在处理复杂形状和边缘信息时保持全局和局部的细节。
轻量化与可移植性： UNet模型结构简单，参数量较少，适合在资源受限的场景中部署，如边缘设备或实时检测中使用。
有效的数据增强策略： 项目中通过多种数据增强（随机翻转、旋转等）提升了模型对不同场景和缺陷形态的适应性，缓解了类别不平衡的问题。