基于Unet++在kaggle—2018dsb数据集上实现图像分割

1. 作者介绍

郭冠群，男，西安工程大学电子信息学院，2023级研究生
研究方向：机器视觉与人工智能
电子邮件：1347418097@qq.com

路治东，男，西安工程大学电子信息学院，2022级研究生，张宏伟人工智能课题组
研究方向：机器视觉与人工智能
电子邮件：2063079527@qq.com

2. 理论知识介绍

2.1 Unet++模型介绍

• Unet
  语义分割是将图像划分为有意义的区域，并标注每个区域所属的类别。语义分割网络是实现这一任务的工具，其中Unet模型通过跨阶段融合不同尺寸的特征图来实现这一目标。
  
• 特征图融合
  特征图融合的目的是结合浅层和深层特征，提升分割效果。浅层特征能提取图像的简单特征如边界和颜色，而深层特征提取图像的深层次语义信息。多个特征图的融合能够弥补单一特征层次信息的不足。
• Unet++
  Unet++通过嵌套的密集跳过路径连接编码器和解码器子网络，减少了特征映射之间的语义差距，从而提高了分割效果。在测试阶段，由于输入图像只进行前向传播，被剪掉的部分对前面输出没有影响，而在训练阶段，这些部分会帮助其他部分进行权重更新。
  

3. 实验过程

3.1 数据集介绍

• 数据集来源
  Kaggle—2018dsb数据集来自于2018年数据科学碗，其任务是从显微镜图像中分割细胞核。这对于推动医学发现具有重要意义，特别是在病理学、癌症研究和其他生命科学领域。
  
• 下载途径

百度网盘 链接：https://pan.baidu.com/s/1GXtZ0clE12oZKooF61siKQ
提取码：tsh7

• 数据集内容
  数据集包含显微镜下细胞图像及其对应的分割掩码。训练集用于训练模型，测试集用于评估模型性能。
  

3.2 代码实现

1. train.py

import os
import argparse
from glob import glob
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from tqdm import tqdm
import albumentations as albu
from albumentations.core.composition import Compose, OneOf
from sklearn.model_selection import train_test_split
import archs
import losses
from dataset import CustomDataset
from metrics import iou_score
from utils import AverageMeter, str2bool

class Config:
    @staticmethod
    def from_cmdline():
        parser = argparse.ArgumentParser(description='Training configuration')
        parser.add_argument('--name', default=None, help='Model name: (default: arch+timestamp)')
        parser.add_argument('--epochs', type=int, default=100, help='Number of total epochs to run')
        parser.add_argument('-b', '--batch_size', type=int, default=8, help='Mini-batch size (default: 16)')
        parser.add_argument('--arch', default='NestedUNet', choices=archs.__all__, help='Model architecture')
        parser.add_argument('--deep_supervision', type=str2bool, default=False, help='Use deep supervision if True')
        parser.add_argument('--input_channels', type=int, default=3, help='Number of input channels')
        parser.add_argument('--num_classes', type=int, default=1, help='Number of classes')
        parser.add_argument('--input_w', type=int, default=96, help='Input image width')
        parser.add_argument('--input_h', type=int, default=96, help='Input image height')
        parser.add_argument('--loss', default='BCEDiceLoss', choices=losses.__all__, help='Loss function')
        parser.add_argument('--dataset', default='dsb2018_96', help='Dataset name')
        parser.add_argument('--img_ext', default='.png', help='Image file extension')
        parser.add_argument('--mask_ext', default='.png', help='Mask file extension')
        parser.add_argument('--optimizer', default='SGD', choices=['Adam', 'SGD'], help='Optimizer type')
        parser.add_argument('--lr', '--learning_rate', type=float, default=1e-3, help='Initial learning rate')
        parser.add_argument('--momentum', type=float, default=0.9, help='Optimizer momentum')
        parser.add_argument('--weight_decay', type=float, default=1e-4, help='Weight decay rate')
        parser.add_argument('--nesterov', type=str2bool, default=False, help='Nesterov momentum')
        parser.add_argument('--scheduler', default='CosineAnnealingLR',
                            choices=['CosineAnnealingLR', 'ReduceLROnPlateau', 'MultiStepLR', 'ConstantLR'],
                            help='Learning rate scheduler')
        parser.add_argument('--min_lr', type=float, default=1e-5, help='Minimum learning rate')
        parser.add_argument('--factor', type=float, default=0.1, help='Factor for ReduceLROnPlateau')
        parser.add_argument('--patience', type=int, default=2, help='Patience for ReduceLROnPlateau')
        parser.add_argument('--milestones', type=str, default='1,2', help='Milestones for MultiStepLR')
        parser.add_argument('--gamma', type=float, default=2 / 3, help='Gamma for MultiStepLR')
        parser.add_argument('--early_stopping', type=int, default=-1, help='Early stopping threshold')
        parser.add_argument('--num_workers', type=int, default=0, help='Number of data loading workers')

        args = parser.parse_args()
        return vars(args)


class ModelManager:
    def __init__(self, config):
        self.config = config
        self.model = self.create_model().cuda()
        self.criterion = self.create_criterion().cuda()
        self.optimizer = self.create_optimizer()
        self.scheduler = self.create_scheduler()

    def create_model(self):
        return archs.__dict__[self.config['arch']](
            self.config['num_classes'],
            self.config['input_channels'],
            self.config['deep_supervision']
        )

    def create_criterion(self):
        if self.config['loss'] == 'BCEWithLogitsLoss':
            return nn.BCEWithLogitsLoss()
        else:
            return losses.__dict__[self.config['loss']]()

    def create_optimizer(self):
        params = filter(lambda p: p.requires_grad, self.model.parameters())
        if self.config['optimizer'] == 'Adam':
            return optim.Adam(params, lr=self.config['lr'], weight_decay=self.config['weight_decay'])
        elif self.config['optimizer'] == 'SGD':
            return optim.SGD(params, lr=self.config['lr'], momentum=self.config['momentum'],
                             nesterov=self.config['nesterov'], weight_decay=self.config['weight_decay'])

    def create_scheduler(self):
        if self.config['scheduler'] == 'CosineAnnealingLR':