AI-exp-2/code/unet_segmentation.py

# 导入必要的库
import os  # 文件和路径管理
import time  # 计时功能
import logging  # 日志记录
import json  # 处理JSON数据

import numpy as np  # 数值运算
from PIL import Image  # 图像处理

import torch  # PyTorch深度学习框架
import torch.nn as nn  # 神经网络模块
import torch.optim as optim  # 优化器
from torch.utils.data import Dataset, DataLoader  # 数据集和加载器
import torchvision.utils as vutils  # 可视化工具

# COCO数据集相关库
from pycocotools import mask as maskUtils  # COCO掩码处理工具
import albumentations as A  # 数据增强库
from albumentations.pytorch import ToTensorV2  # Albumentations到Tensor的转换

torch.backends.cudnn.benchmark = True
# 固定随机种子，确保实验可复现
torch.manual_seed(42)
# 配置日志格式
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# ------------------ 定义 U-Net 模型（自定义） ------------------
class UNet(nn.Module):
    def __init__(self, in_channels=3, base_channels=64, out_channels=1):
        super(UNet, self).__init__()
        # 双卷积块定义
        def double_conv(in_c, out_c):
            return nn.Sequential(
                nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
                nn.BatchNorm2d(out_c),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_c, out_c, kernel_size=3, padding=1),
                nn.BatchNorm2d(out_c),
                nn.ReLU(inplace=True)
            )
        # 下采样路径
        self.enc1 = double_conv(in_channels, base_channels)
        self.enc2 = double_conv(base_channels, base_channels*2)
        self.enc3 = double_conv(base_channels*2, base_channels*4)
        self.enc4 = double_conv(base_channels*4, base_channels*8)
        self.pool = nn.MaxPool2d(2)  # 最大池化层
        # 中心层
        self.center = double_conv(base_channels*8, base_channels*16)
        # 上采样路径
        self.up4 = nn.ConvTranspose2d(base_channels*16, base_channels*8, kernel_size=2, stride=2)
        self.dec4 = double_conv(base_channels*16, base_channels*8)
        self.up3 = nn.ConvTranspose2d(base_channels*8, base_channels*4, kernel_size=2, stride=2)
        self.dec3 = double_conv(base_channels*8, base_channels*4)
        self.up2 = nn.ConvTranspose2d(base_channels*4, base_channels*2, kernel_size=2, stride=2)
        self.dec2 = double_conv(base_channels*4, base_channels*2)
        self.up1 = nn.ConvTranspose2d(base_channels*2, base_channels, kernel_size=2, stride=2)
        self.dec1 = double_conv(base_channels*2, base_channels)
        # 最终1x1卷积层
        self.final = nn.Conv2d(base_channels, out_channels, kernel_size=1)

    def forward(self, x):
        # 编码路径
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))
        # 中心层
        c = self.center(self.pool(e4))
        # 解码路径
        d4 = self.up4(c)
        d4 = torch.cat([d4, e4], dim=1)  # 特征拼接
        d4 = self.dec4(d4)
        d3 = self.up3(d4)
        d3 = torch.cat([d3, e3], dim=1)
        d3 = self.dec3(d3)
        d2 = self.up2(d3)
        d2 = torch.cat([d2, e2], dim=1)
        d2 = self.dec2(d2)
        d1 = self.up1(d2)
        d1 = torch.cat([d1, e1], dim=1)
        d1 = self.dec1(d1)
        return self.final(d1)  # 最终输出

# ------------------ Dice Loss 与 IoU 指标 ------------------
class DiceLoss(nn.Module):
    def __init__(self, eps=1e-6):
        super(DiceLoss, self).__init__()
        self.eps = eps  # 平滑项
    def forward(self, logits, targets):
        probs = torch.sigmoid(logits)  # 转换为概率
        num = 2 * (probs * targets).sum(dim=(2,3)) + self.eps  # 分子
        den = probs.sum(dim=(2,3)) + targets.sum(dim=(2,3)) + self.eps  # 分母
        return 1 - (num/den).mean()  # 返回Dice损失

# 计算IoU评分
def iou_score(preds, targets, eps=1e-6):
    preds = (preds > 0.5).float()  # 二值化预测结果
    inter = (preds * targets).sum(dim=(2,3))  # 交集
    union = preds.sum(dim=(2,3)) + targets.sum(dim=(2,3)) - inter  # 并集
    return ((inter+eps)/(union+eps)).mean().item()  # 返回IoU评分

# ------------------ COCO 分割数据集 ------------------
class CocoSegDataset(Dataset):
    def __init__(self, root_dir, annotation_file, transforms=None, mask_transforms=None):
        self.root_dir = root_dir  # 数据集根目录
        self.transforms = transforms  # 数据变换
        self.mask_transforms = mask_transforms  # 掩码变换
        with open(annotation_file, 'r') as f:
            self.coco = json.load(f)  # 加载COCO标注文件
        self.annotations = {}
        for ann in self.coco['annotations']:
            self.annotations.setdefault(ann['image_id'], []).append(ann)  # 建立图像ID到标注的映射
        self.image_ids = list(self.annotations.keys())  # 所有图像ID列表
    def __len__(self): return len(self.image_ids)  # 返回数据集大小
    def __getitem__(self, idx):
        img_id = self.image_ids[idx]  # 获取图像ID
        info = next(x for x in self.coco['images'] if x['id']==img_id)  # 获取图像信息
        img = Image.open(os.path.join(self.root_dir, info['file_name'])).convert('RGB')  # 读取图像
        h, w = info['height'], info['width']  # 获取图像尺寸
        mask = np.zeros((h,w), dtype=np.uint8)  # 初始化掩码
        for ann in self.annotations[img_id]:  # 生成掩码
            seg = ann['segmentation']
            if isinstance(seg, list): rle = maskUtils.merge(maskUtils.frPyObjects(seg,h,w))
            else: rle = seg
            mask += maskUtils.decode(rle)
        mask = (mask>0).astype(np.float32)  # 二值化掩码
        mask = Image.fromarray(mask)  # 转换为PIL图像
        # 应用数据增强
        if self.transforms and self.mask_transforms:
            aug = self.transforms(image=np.array(img), mask=np.array(mask))
            img_t = aug['image']; m_t = aug['mask'].unsqueeze(0)
        else:
            img_t = ToTensorV2()(image=np.array(img))['image']
            m_t   = ToTensorV2()(image=np.array(mask))['image']
        return img_t, m_t  # 返回处理后的图像和掩码

# ------------------ 主训练流程 ------------------
if __name__ == '__main__':
    # 路径配置
    train_dir, val_dir = '../data/train', '../data/valid'  # 训练集和验证集目录
    train_ann, val_ann = os.path.join(train_dir,'_annotations.coco.json'), os.path.join(val_dir,'_annotations.coco.json')  # 标注文件路径
    # 数据增强配置
    train_tf = A.Compose([
        A.Resize(256,256),  # 调整尺寸
        A.HorizontalFlip(0.5),  # 水平翻转
        A.RandomBrightnessContrast(0.2),  # 随机亮度对比度
        A.ShiftScaleRotate(0.5),  # 随机位移缩放旋转
        A.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225)),  # 归一化
        ToTensorV2()])  # 转换为Tensor
    val_tf   = A.Compose([
        A.Resize(256,256),  # 调整尺寸
        A.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225)),  # 归一化
        ToTensorV2()])  # 转换为Tensor
    # 数据加载
    train_ds = CocoSegDataset(train_dir,train_ann,train_tf,train_tf)  # 训练数据集
    val_ds   = CocoSegDataset(val_dir,  val_ann,  val_tf,  val_tf)  # 验证数据集
    train_ld = DataLoader(train_ds,batch_size=8,shuffle=True,num_workers=4)  # 训练数据加载器
    val_ld   = DataLoader(val_ds,  batch_size=8,shuffle=False,num_workers=4)  # 验证数据加载器
    logging.info(f"Train samples: {len(train_ds)}, Val samples: {len(val_ds)}")  # 输出数据集信息
    # 模型与训练配置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 设备选择
    print(f"当前设备：{device}")  # 输出当前设备信息
    model = UNet().to(device)  # 初始化模型并移动到指定设备
    opt = optim.AdamW(model.parameters(),lr=1e-3,weight_decay=1e-4)  # 优化器
    sched = optim.lr_scheduler.ReduceLROnPlateau(opt,'min',patience=3,factor=0.5)  # 学习率调度器
    bce = nn.BCEWithLogitsLoss(); dice = DiceLoss()  # 损失函数
    epochs=20; vis_dir='output/val_visuals'; os.makedirs(vis_dir,exist_ok=True)  # 训练参数和可视化目录
    # 训练循环
    for ep in range(1,epochs+1):
        model.train(); run_loss=0  # 设置为训练模式
        for imgs,msks in train_ld:
            imgs,msks=imgs.to(device),msks.to(device)  # 将数据移动到指定设备
            opt.zero_grad(); out=model(imgs)  # 前向传播
            l=(bce(out,msks)+dice(out,msks)); l.backward(); opt.step()  # 计算损失、反向传播
            run_loss+=l.item()*imgs.size(0)  # 累计损失
        tr_loss=run_loss/len(train_ds)  # 计算平均训练损失
        # 验证阶段
        model.eval(); v_loss=0; v_iou=0; v_dice=0
        with torch.no_grad():
            for imgs,msks in val_ld:
                imgs,msks=imgs.to(device),msks.to(device)  # 将数据移动到指定设备
                out=model(imgs)  # 前向传播
                v_loss+=(bce(out,msks)+dice(out,msks)).item()*imgs.size(0)  # 累计验证损失
                pr=torch.sigmoid(out)  # 转换为概率
                v_iou+=iou_score(pr,msks)*imgs.size(0)  # 累计IoU评分
                v_dice+=(1 - dice(out,msks)).item()*imgs.size(0)  # 累计Dice评分
        v_loss/=len(val_ds); v_iou/=len(val_ds); v_dice/=len(val_ds)  # 计算平均验证指标
        # 输出训练信息
        logging.info(f"Epoch {ep}/{epochs} - Tr:{tr_loss:.4f} Val:{v_loss:.4f} IoU:{v_iou:.4f} Dice:{v_dice:.4f}")
        # 可视化
        si,sm=next(iter(val_ld)); si=si.to(device)  # 获取示例图像
        with torch.no_grad(): sp=torch.sigmoid(model(si))  # 进行预测
        grid=vutils.make_grid(torch.cat([si.cpu(),sm.repeat(1,3,1,1).cpu(),sp.repeat(1,3,1,1).cpu()],0),nrow=si.size(0))  # 创建网格
        vpth=os.path.join(vis_dir,f'ep{ep}.png'); vutils.save_image(grid,vpth)  # 保存可视化结果
        logging.info(f"Saved visual: {vpth}")  # 输出保存信息
        sched.step(v_loss)  # 更新学习率
    # 保存模型权重
    torch.save(model.state_dict(),'unet_coco_segmentation.pth')
    logging.info('训练完成，模型已保存')