目录
1.引言
图像分割是计算机视觉领域的核心任务之一，在医学影像分析、自动驾驶、遥感图像处理等领域有着广泛应用。本文通过对比两种经典分割网络——UNet和DeepLabV3，分享一个完整的图像分割项目实现，包括数据预处理、模型训练、评估指标可视化和模型对比分析。
2.代码架构
本项目采用模块化设计，主要包含三个核心模块：
- 训练主程序：负责模型训练流程控制
- 数据处理与工具函数：实现数据加载、预处理和评估指标计算
- 结果可视化：训练过程指标的可视化对比
3.关键代码
1. 数据预处理
class MyDataset(Dataset):
def __init__(self, imgs_path, img_fm, mk_fm, txt_path='./runs/grayList.txt',
ct=False, base_size=(256,256)):
# 初始化数据路径和参数
self.imgs = [os.path.join(imgs_path,i) for i in os.listdir(imgs_path)]
self.baseS = base_size
self.img_format = img_fm
self.mask_format = mk_fm
self.window_CT = ct # CT图像专用窗宽窗位调整
self.data_aug = True
def __getitem__(self, index):
# 实现数据加载、归一化、尺寸调整和数据增强
image = np.array(Image.open(image_path).convert('RGB'))
mask = np.array(Image.open(mask_path).convert('L'))
# 灰度值标准化处理
for gray in np.unique(mask):
mask[mask==gray] = self.txt.index(str(gray))
# 数据增强
if self.data_aug:
if random.random() > 0.5:
image = np.ascontiguousarray(np.flip(image,axis=-1))
mask = np.ascontiguousarray(np.flip(mask,axis=-1))
return tensor(image).float(), tensor(mask).long()关键特点：
- 支持CT图像专用窗宽窗位调整
- 实现随机水平/垂直翻转的数据增强
- 自动处理多类别标签的灰度值映射
2. 模型训练与评估
训练流程采用标准的PyTorch训练范式，并添加了丰富的评估指标：
def train_one_epoch(model, optim, loader, device, numcl, f):
criterion = torch.nn.CrossEntropyLoss()
confmat = ConfusionMatrix(num_classes=numcl) # 混淆矩阵
for x, y in loader:
# 前向传播
o = model(x)
if f == 'deeplabv3':
o = model(x)['out'] # DeepLabV3特殊处理
# 损失计算与反向传播
loss = criterion(o, y)
optim.zero_grad()
loss.backward()
optim.step()
# 评估指标计算
with torch.no_grad():
model.eval()
o = torch.argmax(model(x), dim=1)
confmat.update(y, o)
# 返回学习率、损失和各项指标
return lr, loss.item(), miou, mdice, confmat评估指标包括：
- 像素准确率(Pixel Accuracy)
- 平均交并比(mIoU)
- Dice系数
- 精确率(Precision)
- 召回率(Recall)
- F1分数
3. 学习率调度
采用余弦退火学习率衰减策略，实现更稳定的训练过程：
# 自适应学习率衰减
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.模型对比分析
通过可视化工具可以直观比较不同模型的性能表现：
def plot_comparison(metrics, metric_type='f1', title=None):
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
# 绘制训练集指标
for idx, model_name in enumerate(model_names):
epochs = sorted(metrics['train'][model_name].keys())
values = [metrics['train'][model_name][e][metric_type] for e in epochs]
ax1.plot(epochs, values, label=model_name.upper())
# 绘制验证集指标
for idx, model_name in enumerate(model_names):
if model_name in metrics['val']:
epochs = sorted(metrics['val'][model_name].keys())
values = [metrics['val'][model_name][e][metric_type] for e in epochs]
ax2.plot(epochs, values, label=model_name.upper())可视化结果通常显示：
- UNet在小数据集上通常收敛更快
- DeepLabV3在大数据集上表现更优
- 两种模型在不同指标上的曲线
4.项目使用指南
-
数据准备：
- 按照
images和masks目录结构组织数据 - 对于医学CT图像，设置
--ct True启用窗宽窗位调整
- 按照
-
训练模型：
python train.py --model unet --epochs 100 --batch-size 16 --lr 0.001
python train.py --model deeplabv3 --epochs 100 --batch-size 8 --lr 0.00013.结果可视化：
- 自动生成训练曲线图
- 使用分析脚本比较不同模型：
本文实现了一个完整的图像分割项目框架，通过对比UNet和DeepLabV3两种经典网络，展示了不同架构在分割任务上的表现差异。项目提供了从数据预处理到结果可视化的完整流程，可以作为图像分割任务的开发模板。实验结果表明，没有绝对优越的模型，实际应用中需要根据具体任务需求和数据特点选择合适的架构。
5.脊椎分割实验
1.数据集
2.训练
这里仅仅训练了5个epoch用于测试代码的可行性，指标不做参考
因为unet的结果很好，这里只展示unet，而非deeplabv3
Namespace(base_size=(224, 224), batch_size=4, ct=False, data_test=False, data_train='./data/train/images', data_val='./data/val/images', epochs=5, img_f='Slice', lr=0.0001, lrf=0.01, mask_f='Labels_slice', model='unet', results='runs_unet')
结果保存在------> runs_unet 目录
当前训练设备: cuda
compute model output: 100%|██████████| 511/511 [00:00<00:00, 1488.37it/s]
数据标签灰度值: [0, 255]
网络输出个数: 2
数据加载使用线程数: 4
训练集个数:511 验证集个数:127
预处理数据可视化结果---->trainSetVis.jpg
选择的模型: unet
-----------train-----------
[epoch:0/4]
train: 100%|██████████| 128/128 [00:30<00:00, 4.17it/s, dice=0.504, iou=0.337, loss=0.177]
valid: 100%|██████████| 32/32 [00:05<00:00, 5.73it/s, dice=0.839, iou=0.723, loss=0.178]
train: 0%| | 0/128 [00:00<?, ?it/s]learning rate:0.00010000
train loss:0.3018 train mean iou:0.1173 train mean dice:0.1885
val loss:0.1834 val mean iou:0.7428 val mean dice:0.8523
[epoch:1/4]
train: 100%|██████████| 128/128 [00:30<00:00, 4.23it/s, dice=0.742, iou=0.59, loss=0.132]
valid: 100%|██████████| 32/32 [00:05<00:00, 5.68it/s, dice=0.857, iou=0.75, loss=0.131]
learning rate:0.00009055
train loss:0.1514 train mean iou:0.6923 train mean dice:0.8178
val loss:0.1339 val mean iou:0.7758 val mean dice:0.8736
[epoch:2/4]
train: 100%|██████████| 128/128 [00:30<00:00, 4.26it/s, dice=0.862, iou=0.757, loss=0.103]
valid: 100%|██████████| 32/32 [00:05<00:00, 5.77it/s, dice=0.874, iou=0.776, loss=0.103]
train: 0%| | 0/128 [00:00<?, ?it/s]learning rate:0.00006580
train loss:0.1149 train mean iou:0.7451 train mean dice:0.8539
val loss:0.1067 val mean iou:0.7897 val mean dice:0.8824
[epoch:3/4]
train: 100%|██████████| 128/128 [00:30<00:00, 4.23it/s, dice=0.891, iou=0.803, loss=0.0956]
valid: 100%|██████████| 32/32 [00:05<00:00, 5.74it/s, dice=0.888, iou=0.799, loss=0.0919]
learning rate:0.00003520
train loss:0.0979 train mean iou:0.7943 train mean dice:0.8854
val loss:0.0952 val mean iou:0.8121 val mean dice:0.8963
[epoch:4/4]
train: 100%|██████████| 128/128 [00:29<00:00, 4.30it/s, dice=0.908, iou=0.831, loss=0.0893]
valid: 100%|██████████| 32/32 [00:05<00:00, 5.67it/s, dice=0.904, iou=0.825, loss=0.0901]
learning rate:0.00001045
train loss:0.0900 train mean iou:0.8261 train mean dice:0.9048
val loss:0.0908 val mean iou:0.8434 val mean dice:0.9150
train finish!
验证集上表现最好的epoch为: 4
进程已结束，退出代码为 0
3.对比
可以自己将模型扩充多个，然后，这里填对应的json结果即可
file_paths = [
'./runs_unet/train_log_results.json',
'./runs_deeplabv3/train_log_results.json'
]
6.下载地址
下载地址：基于UNet与DeepLabV3的医学图像分割系统：训练、评估与可视化分析脊椎分割资源-CSDN文库
关于图像分类、分割网络的改进：图像分类网络改进_听风吹等浪起的博客-CSDN博客
扫一扫
1393
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
