图像分割——图像自动识别汽车边缘
背景介绍
随着二手市场的兴起，产品信息的完整性和透明度成为了二手商品销售的核心。相较于新品，二手产品往往在信息完整性和透明度方面存在欠缺，这导致消费者往往难以对其质量产生信任。尤其在线上二手市场，由于消费者仅能通过有限的图片和描述来判断产品的可靠性，这一问题变得更加突出。
Carvana，一家专注于在线二手车销售的公司，精准捕捉到了消费者的这一心理需求，并在提供产品信息方面做了大量工作，极大地简化了购买流程。Carvana的一个亮点是其官网上对库存中的汽车提供16张详细图片，基于这些图片生成动态模型，允许消费者在购买前进行旋转查看，深入了解汽车的外观细节。然而，汽车图片中的明亮反光或与车身颜色相似的背景有时会导致图片处理出现误差。在这个问题中，如何准确识别汽车边缘成了一个重要课题。
在这个项目中，我们选择使用经典的UNet模型，并结合Pytorch技术实现对汽车边缘的精准识别。
方法介绍
图像分割是计算机视觉中非常重要的任务，是图像理解的基础。在本项目中的分割任务是语义分割：在普通分割的基础上，分类出每一块区域的语义，即在一张包含汽车的广告图中，将汽车部分分割出来。
U-Net是一种常用于医学图像分割的卷积神经网络架构，其特点是具有一个对称的“U”形结构，能够有效地捕捉图像中的局部和上下文信息。U-Net由收缩路径（下采样或编码器）和扩张路径（上采样或解码器）组成。
其中，收缩路径由卷积层和池化层交替组成，逐渐降低图像的分辨率，用来捕获上下文信息和提取特征；对称的扩张路径 由反卷积层和卷积层交替组成，将特征图还原到原始图像大小，同时与编码器相应层的特征图进行融合（跳跃连接），以恢复空间信息。最后，经过一系列的上采样和特征融合后，UNet 输出分割后的图像。
模型代码介绍
Train.py文件中的Main函数为程序训练的主要窗口，它首先通过构建数组转换组合对象，用于对训练集和评估集的预处理：
# 对训练集的预处理 A.Compose为albumentations数据强化的库函数
train_transform = A.Compose(
[
# 使用Resize函数将图片预处理为指定长宽
A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),
# 使用各种旋转翻转函数对图像进行随机旋转、翻转，增强模型的泛化能力
A.Rotate(limit=35, p=1.0),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.1),
# 对图像进行归一化，即将图像的像素值缩放到特定的范围
A.Normalize(
mean=[0.0, 0.0, 0.0],
std=[1.0, 1.0, 1.0],
max_pixel_value=255.0,
),
# 通过ToTensorV2函数将图像转换为张量(tensor)的格式
ToTensorV2(),
],
)
随后构建U-Net模型，并利用Adam优化算法构建一个优化器：
# 创建UNET模型，并设置损失函数以及优化器
model = UNET(in_channels=3, out_channels=1).to(DEVICE)
loss_fn = nn.BCEWithLogitsLoss()# binary cross entropy
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
初始化参数后开始epochs数次训练，并且在每个训练周期其中保存模型的权重参数以及准确率：
# 训练NUM_EPOCHS次并计算每次的正确率和损失值
for epoch in range(NUM_EPOCHS):
# 使用train_fn函数对模型进行训练
train_fn(train_loader, model, optimizer, loss_fn, scaler)
# 保存模型
checkpoint = {
"state_dict": model.state_dict(),
"optimizer":optimizer.state_dict(),
}
save_checkpoint(checkpoint)
# 检查正确率
check_accuracy(val_loader, model, device=DEVICE)
# 打印一些示例
save_predictions_as_imgs(
val_loader, model, folder="saved_images/", device=DEVICE
)
Main函数中调用的Tran_fn函数是对训练数据进行迭代并执行模型的训练过程，用于训练所构建模型的神经网络。在每个迭代中，遍历每个加载器中的数据和对应标签，并将其目标张量移动到内核(CPU或GPU)中，在GPU中启用混合精度训练，训练过程中计算正确率以及损失值。接下来将优化器梯度清零，反向传播更新参数，最后更新scaler对象调整精度缩放因子。
# 训练函数train_fn
def train_fn(loader, model, optimizer, loss_fn, scaler):
# loop为tqdm库下创建的训练进度条，用于可视化训练进度
loop = tqdm(loader)
# 迭代数据加载器中的每批数据，并将数据和对应标签放入内核中
for batch_idx, (data, targets) in enumerate(loop):
data = data.to(device=DEVICE)
targets = targets.float().unsqueeze(1).to(device=DEVICE)
# 使用混合精度训练的上下管理器，将计算转换为混合精度，提升训练和内存效率
with torch.cuda.amp.autocast():
# 使用模型向前传播
predictions = model(data)
# 计算损失值
loss = loss_fn(predictions, targets)
# 后向传播
optimizer.zero_grad() # 清除优化器梯度
scaler.scale(loss).backward() # 依据损失值进行反向传播
scaler.step(optimizer) # 根据缩放后的梯度更新模型参数
scaler.update() # 更新缩放器内部的缩放因子
# 更新进度条
loop.set_postfix(loss=loss.item())
Model.py为U-Net模型的主要函数，其定义了U-Net模型以及子模块DoubleConv。首先DoubleConv是由两个连续卷积层、标准化层和ReLU激活函数组成：该模块主要功能为提取上采样和下采样过程中的特征:
# 自定义卷积神经网络模块，DoubleConv继承nn.Module类
class DoubleConv(nn.Module):
# 构造函数，用于初始化模块
def __init__(self, in_channels, out_channels):
super(DoubleConv, self).__init__()
# 序列容器self.conv
self.conv = nn.Sequential(
# 创建一个二维卷积层
nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),
# 三个通道 1 100 10000 --> 1 1.2 1.3
# 创建一个二维批归一化层，进行通道归一化
nn.BatchNorm2d(out_channels),
# ReLU激活函数层
nn.ReLU(inplace=True),
# 下面再同上创建一个卷积块
nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
)
# 定义向前传播的过程
def forward(self, x):
return self.conv(x)
U-Net类为模型的主要定义，它由上采样(解码器)和下采样(编码器)组成。下采样部分通过多次运用DoubleConv进行特征的提取和降维，当达到一定维数后，再通过上采样部分使用转置卷积层和DoubleConv进行特征提取和升维。最后通过一个卷积层将特征图像映射到输出通道。
# UNET模型类，它是一个网络结构
class UNET(nn.Module):
# 初始化函数，继承自nn.Module类
def __init__(
self, in_channels=3, out_channels=1, features=[64, 128, 256, 512],
):
super(UNET, self).__init__()
self.ups = nn.ModuleList() # 用于存储上采样层的列表
self.downs = nn.ModuleList() # 用于存储下采样层的列表
self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # 最大池化层
# 下采样层
# 对于features每个特征数，创建一个DoubleConv模块，并将其添加到列表中
for feature in features:
self.downs.append(DoubleConv(in_channels, feature))
in_channels = feature
# 上采样层
# 对于features每个特征数，创建一个转置卷积层，并将对应的DoubleConv添加到列表中
for feature in reversed(features):
self.ups.append(
# 创建转置卷积层
nn.ConvTranspose2d(
feature*2, feature, kernel_size=2, stride=2,
)
)
self.ups.append(DoubleConv(feature*2, feature))
# UNET瓶颈部分，使用DoubleConv模块对特征进行处理
self.bottleneck = DoubleConv(features[-1], features[-1]*2)
# UNET最后一个卷积层，用于将特征图转换为输出的预测结果
self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
# 定义UNET模型前向传播
def forward(self, x):
skip_connections = []
# 通过self.down列表对数据进行下采样
for down in self.downs:
x = down(x)
# 将结果存储在skip中
skip_connections.append(x)
x = self.pool(x)
# 对输入的数据进行特征提取和降维
x = self.bottleneck(x)
# 翻转，准备上采样
skip_connections = skip_connections[::-1]
# 通过转置卷积层self.ups[idx]进行上采样
for idx in range(0, len(self.ups), 2):
x = self.ups[idx](x)
skip_connection = skip_connections[idx//2]
# 若上采样后的特征图形状与skip中的形状不符，则重新调整图像使得有相同大小
if x.shape != skip_connection.shape:
x = TF.resize(x, size=skip_connection.shape[2:])
# 重新生成新的特征图
concat_skip = torch.cat((skip_connection, x), dim=1)
# 赋值，准备进行下一次迭代
x = self.ups[idx+1](concat_skip)
return self.final_conv(x)
Utils.py文件主要用于加载模型或保存模型，并且有检查正确率的功能 Dataset.py主要用于加载数据集数据集。
训练过程
通过大量的训练集及其对应的图像掩模（mask）集，将包含汽车的广告图和与其对应的mask图作为一组进行训练。经过训练之后的模型在验证集上进行验证。下图为经过两轮训练后，在验证集上的准确率：
| 示例输入 | 输出 |
|---|---|
 |  |
训练效果
对于数据集中的图片测试结果如下：
 |  |  |
|---|---|---|
 |  |
将其与人工绘制的mask对比（变换前为人工绘制，变换后为模型运行结果）：
 |  |  |
|---|---|---|
 |  | |
最终效果
对于自选图片的测试结果如下：
 |  |  |
|---|---|---|
 |  |
发现该模型能比较好的生成图片的mask，但对于背景复杂的图片，准确率仍有待提高。
应用领域
生成汽车的Image Masking（图像掩模）在图像分割领域的应用有着广阔的前景。U-Net模型是此任务的一个经典、效果较好的解决方式。
在生成图像掩模进行图像分割上，有很多当下的应用：
自动驾驶和车辆识别
在自动驾驶中，准确识别和分割道路上的车辆对于安全导航至关重要。汽车掩模模型可以帮助识别和规避道路上的障碍物。
交通监控和管理
在交通监控系统中，可以用于统计道路上的车辆数量，分析交通流量。同时也可用于自动检测违章停车、闯红灯等交通违规行为。
在未来，或许能够参与解决更复杂、现实的问题：
智能交通系统
利用车辆图像分割数据来改善道路设计，提高交通效率和安全性；快速识别交通事故，优化紧急服务的响应。
无人机和机器人导航
在无人递送服务中，用于识别最佳路径和避障。在不熟悉或复杂的环境中探索和导航。
高级车辆维护和诊断
自动检测和评估车辆的损伤，如刮痕、凹陷等；辅助技术人员识别问题部位和进行维修工作。
转载自:https://juejin.cn/post/7314171116363350051