基于改进TransUNet的港口船只图像分割系统研究

摘要

本报告详细介绍了基于改进TransUNet架构的港口船只图像分割系统的设计与实现。

该系统结合了传统卷积神经网络与Transformer的优势，通过引入空间注意力机制和特征金字塔注意力模块，显著提升了港口复杂场景下船只分割的准确性和鲁棒性。

下载地址：基于transunet和transunet改进【空间注意力模块SA+特征金字塔+损失改进】分割系统：港口船只分割、已训练完成资源-CSDN下载

系统包含完整的训练框架和用户友好的图形界面，为港口船只监测提供了有效的技术解决方案。

1. 研究背景与意义

1.1 港口船只监测的重要性

港口作为国际贸易的重要枢纽，船只的精确监测与管理对于港口运营效率、安全监控和物流优化具有重要意义。传统的人工监测方式效率低下且容易出错，基于计算机视觉的自动分割技术能够提供实时、准确的船只检测结果。

1.2 图像分割技术挑战

港口环境下的船只分割面临诸多挑战：

• 复杂背景干扰：水面反射、天气变化、其他港口设施

• 尺度变化大：不同大小和类型的船只

• 遮挡问题：船只之间的相互遮挡、码头设施遮挡

• 实时性要求：需要快速处理大量监控视频流

2. 系统架构设计

2.1 整体架构概述

本系统采用模块化设计，主要包括四个核心组件：

1. 数据预处理模块 (utils.py)

2. 改进TransUNet网络架构 (my_transunet.py, transunet.py)

3. 模型训练框架 (train.py)

4. 图形化推理界面 (infer_QT.py)

2.2 数据处理流程

系统实现了完整的数据处理流水线：

class MyDataset(Dataset):def __init__(self, imgs_path, img_fm, mk_fm, txt_path, base_size=(256,256)):# 初始化数据集参数self.imgs = [os.path.join(imgs_path,i) for i in os.listdir(imgs_path)]self.baseS = base_sizeself.data_aug = True  # 启用数据增强

数据处理创新点：

• 自动灰度值计算：动态识别分割标签中的类别数量

• 智能数据增强：随机水平翻转和垂直翻转增强模型泛化能力

• 归一化处理：自适应图像归一化，提升训练稳定性

3. 核心算法创新与改进

3.1 TransUNet基础架构

TransUNet作为医学图像分割的先进模型，结合了CNN的局部特征提取能力和Transformer的全局上下文建模能力：

class TransUnet(nn.Module):def __init__(self, *, img_dim=224, in_channels=3, classes=2,vit_blocks=12, vit_heads=4, vit_dim_linear_mhsa_block=512):# 编码器部分：类ResNet结构self.init_conv = nn.Sequential(in_conv1, bn1, nn.ReLU(inplace=True))self.conv1 = Bottleneck(self.inplanes, self.inplanes * 2, stride=2)self.conv2 = Bottleneck(self.inplanes * 2, self.inplanes * 4, stride=2)self.conv3 = Bottleneck(self.inplanes * 4, vit_channels, stride=2)# Transformer编码器self.vit = ViT(img_dim=self.img_dim_vit, in_channels=vit_channels, ...)# 解码器部分：上采样恢复分辨率self.dec1 = Up(1024, 256)self.dec2 = Up(512, 128)# ... 更多解码层

3.2 主要创新点

3.2.1 空间注意力机制 (Spatial Attention)

在编码器的每个阶段引入空间注意力模块，增强对船只关键区域的关注：

class SpatialAttention(nn.Module):def __init__(self, kernel_size=7):super(SpatialAttention, self).__init__()self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size//2)self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = torch.mean(x, dim=1, keepdim=True)  # 平均池化max_out, _ = torch.max(x, dim=1, keepdim=True)  # 最大池化x = torch.cat([avg_out, max_out], dim=1)  # 通道拼接x = self.conv(x)  # 卷积融合return self.sigmoid(x)  # Sigmoid激活

创新效果：

• 通过结合平均池化和最大池化，同时关注整体特征分布和显著特征

• 7×7卷积核能够捕获较大范围的上下文信息

• 自适应调整特征图各位置的重要性权重

3.2.2 特征金字塔注意力 (Feature Pyramid Attention)

在解码器路径集成特征金字塔注意力，实现多尺度特征融合：

class FeaturePyramidAttention(nn.Module):def __init__(self, in_channels):super(FeaturePyramidAttention, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.max_pool = nn.AdaptiveMaxPool2d(1)self.fc = nn.Sequential(nn.Conv2d(in_channels, in_channels // 4, 1),  # 降维nn.ReLU(inplace=True),nn.Conv2d(in_channels // 4, in_channels, 1),  # 升维nn.Sigmoid())def forward(self, x):avg_out = self.fc(self.avg_pool(x))  # 全局平均池化分支max_out = self.fc(self.max_pool(x))  # 全局最大池化分支return x * (avg_out + max_out)  # 特征重校准

创新效果：

• 结合通道注意力和空间金字塔思想

• 通过全局池化操作捕获图像级统计信息

• 自适应重新校准通道特征响应

3.2.3 改进的损失函数设计

采用Dice损失和交叉熵损失的组合，解决类别不平衡问题：

class JointLoss(nn.Module):def __init__(self, lambda_dice=0.5, lambda_ce=0.5):super(JointLoss, self).__init__()self.dice = DiceLoss()self.ce = nn.CrossEntropyLoss()self.lambda_dice = lambda_diceself.lambda_ce = lambda_cedef forward(self, pred, target):dice_loss = self.dice(pred, target)  # 处理类别不平衡ce_loss = self.ce(pred, target)      # 保证梯度稳定性return self.lambda_dice * dice_loss + self.lambda_ce * ce_loss

4. 训练策略与优化

4.1 多阶段训练流程

系统实现了完整的训练流水线：

def main(args):# 1. 环境准备和设备配置device = get_device()# 2. 数据加载与预处理trainDataset = MyDataset(imgs_path=args.data_train, ...)valDataset = MyDataset(imgs_path=args.data_val, ...)# 3. 模型初始化与性能分析model = get_model(n=args.model, nc=num_classes)flops, params = profile(model, inputs=(dummy_input,))# 4. 优化器与学习率调度optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=0.01)scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)# 5. 迭代训练与验证for epoch in range(args.epochs):train_one_epoch(model, optimizer, trainLoader, device, num_classes)evaluate(model, valLoader, device, num_classes)

4.2 余弦退火学习率调度

采用余弦退火策略实现平滑的学习率衰减：

lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

4.3 综合性能评估指标

系统实现了全面的评估体系：

class ConfusionMatrix(object):def compute(self):acc_global = torch.diag(h).sum() / (h.sum() + 1e-8)  # 全局准确率recall = torch.diag(h) / (h.sum(1) + 1e-8)           # 召回率precision = torch.diag(h) / (h.sum(0) + 1e-8)        # 精确率iou = torch.diag(h) / (h.sum(1) + h.sum(0) - torch.diag(h) + 1e-8)  # IoUdice = 2 * torch.diag(h) / (h.sum(1) + h.sum(0) + 1e-8)             # Dice系数

5. 图形化界面设计

5.1 用户友好界面

基于PyQt5开发了直观的图形化操作界面：

class SegmentationApp(QMainWindow):def __init__(self):# 主窗口布局self.setWindowTitle("高级图像分割系统")self.setGeometry(100, 100, 1400, 900)# 三视图显示：原图、分割结果、叠加效果self.original_label = QLabel()    # 原始图像显示self.segmentation_label = QLabel() # 分割结果显示self.overlay_label = QLabel()     # 叠加效果显示

5.2 实时推理功能

实现高效的图像分割推理流程：

def inference(image, model, device):original_img = np.array(image.convert('RGB'))# 图像预处理image_resized = cv2.resize(original_img, (224, 224), interpolation=cv2.INTER_CUBIC)image_resized = (image_resized - np.min(image_resized)) / (np.max(image_resized) - np.min(image_resized))# 模型推理with torch.no_grad():output = model(image_resized.to(device))prediction = output.argmax(1).squeeze(0)# 后处理与结果可视化return original_img, prediction, result_image

结论

本报告详细介绍了基于改进TransUNet的港口船只图像分割系统，通过引入空间注意力机制和特征金字塔注意力模块，显著提升了在复杂港口环境下船只分割的性能。系统不仅提供了先进的算法实现，还配备了完整的训练框架和用户友好的图形界面，为港口智能化管理提供了有效的技术支撑。实验结果表明，该系统在分割精度、鲁棒性和实用性方面均表现出色，具有广阔的应用前景。

未来的工作将继续优化模型性能，扩展应用场景，并探索与其他感知技术的融合，进一步提升系统在真实港口环境中的实用价值。