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

摘要

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

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

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

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

1. 研究背景与意义

1.1 港口船只监测的重要性

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

1.2 图像分割技术挑战

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

复杂背景干扰：水面反射、天气变化、其他港口设施
尺度变化大：不同大小和类型的船只
遮挡问题：船只之间的相互遮挡、码头设施遮挡
实时性要求：需要快速处理大量监控视频流

2. 系统架构设计

2.1 整体架构概述

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

数据预处理模块 (utils.py)
改进TransUNet网络架构 (my_transunet.py, transunet.py)
模型训练框架 (train.py)
图形化推理界面 (infer_QT.py)

2.2 数据处理流程

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

python 复制代码

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_size
        self.data_aug = True  # 启用数据增强

数据处理创新点：

自动灰度值计算：动态识别分割标签中的类别数量
智能数据增强：随机水平翻转和垂直翻转增强模型泛化能力
归一化处理：自适应图像归一化，提升训练稳定性

3. 核心算法创新与改进

3.1 TransUNet基础架构

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

python 复制代码

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)

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

python 复制代码

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_dice
        self.lambda_ce = lambda_ce

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

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