devbox
羊村懒王
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

Python Unet ++ :医学图像分割，医学细胞分割，Unet医学图像处理，语义分割_from lovaszsoftmax.pytorch import lovasz_losses as

作者：羊村懒王 | 2024-04-02 06:42:46

赞

踩

from lovaszsoftmax.pytorch import lovasz_losses as l

一，语义分割：分割领域前几年的发展

图像分割是机器视觉任务的一个重要基础任务，在图像分析、自动驾驶、视频监控等方面都有很重要的作用。图像分割可以被看成一个分类任务，需要给每个像素进行分类，所以就比图像分类任务更加复杂。此处主要介绍 Deep Learning-based 相关方法。

 

 

 

 主要介绍unet和unet++

 

二，数据介绍---医学细胞分割任务

原数据：

标签数据： 

 

 三，代码部分

模型包含以下文件：

archs.py为模型的主体部分：

  1. import torch
  2. from torch import nn
  3.  
  4. __all__ = ['UNet', 'NestedUNet']
  5.  
  6.  
  7. class VGGBlock(nn.Module):
  8.     def __init__(self, in_channels, middle_channels, out_channels):
  9.         super().__init__()
  10.         self.relu = nn.ReLU(inplace=True)
  11.         self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1)
  12.         self.bn1 = nn.BatchNorm2d(middle_channels)
  13.         self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1)
  14.         self.bn2 = nn.BatchNorm2d(out_channels)
  15.  
  16.     def forward(self, x):
  17.         out = self.conv1(x)
  18.         out = self.bn1(out)
  19.         out = self.relu(out)
  20.  
  21.         out = self.conv2(out)
  22.         out = self.bn2(out)
  23.         out = self.relu(out)
  24.  
  25.         return out
  26.  
  27.  
  28. class UNet(nn.Module):
  29.     def __init__(self, num_classes, input_channels=3, **kwargs):
  30.         super().__init__()
  31.  
  32.         nb_filter = [32, 64, 128, 256, 512]
  33.  
  34.         self.pool = nn.MaxPool2d(2, 2)
  35.         self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)#scale_factor:放大的倍数  插值
  36.  
  37.         self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
  38.         self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
  39.         self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
  40.         self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
  41.         self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])
  42.  
  43.         self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
  44.         self.conv2_2 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
  45.         self.conv1_3 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
  46.         self.conv0_4 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
  47.  
  48.         self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  49.  
  50.  
  51.     def forward(self, input):
  52.         x0_0 = self.conv0_0(input)
  53.         x1_0 = self.conv1_0(self.pool(x0_0))
  54.         x2_0 = self.conv2_0(self.pool(x1_0))
  55.         x3_0 = self.conv3_0(self.pool(x2_0))
  56.         x4_0 = self.conv4_0(self.pool(x3_0))
  57.  
  58.         x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
  59.         x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
  60.         x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
  61.         x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))
  62.  
  63.         output = self.final(x0_4)
  64.         return output
  65.  
  66.  
  67. class NestedUNet(nn.Module):
  68.     def __init__(self, num_classes, input_channels=3, deep_supervision=False, **kwargs):
  69.         super().__init__()
  70.  
  71.         nb_filter = [32, 64, 128, 256, 512]
  72.  
  73.         self.deep_supervision = deep_supervision
  74.  
  75.         self.pool = nn.MaxPool2d(2, 2)
  76.         self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
  77.  
  78.         self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
  79.         self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
  80.         self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
  81.         self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
  82.         self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])
  83.  
  84.         self.conv0_1 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
  85.         self.conv1_1 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
  86.         self.conv2_1 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
  87.         self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
  88.  
  89.         self.conv0_2 = VGGBlock(nb_filter[0]*2+nb_filter[1], nb_filter[0], nb_filter[0])
  90.         self.conv1_2 = VGGBlock(nb_filter[1]*2+nb_filter[2], nb_filter[1], nb_filter[1])
  91.         self.conv2_2 = VGGBlock(nb_filter[2]*2+nb_filter[3], nb_filter[2], nb_filter[2])
  92.  
  93.         self.conv0_3 = VGGBlock(nb_filter[0]*3+nb_filter[1], nb_filter[0], nb_filter[0])
  94.         self.conv1_3 = VGGBlock(nb_filter[1]*3+nb_filter[2], nb_filter[1], nb_filter[1])
  95.  
  96.         self.conv0_4 = VGGBlock(nb_filter[0]*4+nb_filter[1], nb_filter[0], nb_filter[0])
  97.  
  98.         if self.deep_supervision:
  99.             self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  100.             self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  101.             self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  102.             self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  103.         else:
  104.             self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
  105.  
  106.  
  107.     def forward(self, input):
  108.         print('input:',input.shape)
  109.         x0_0 = self.conv0_0(input)
  110.         print('x0_0:',x0_0.shape)
  111.         x1_0 = self.conv1_0(self.pool(x0_0))
  112.         print('x1_0:',x1_0.shape)
  113.         x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))
  114.         print('x0_1:',x0_1.shape)
  115.  
  116.         x2_0 = self.conv2_0(self.pool(x1_0))
  117.         print('x2_0:',x2_0.shape)
  118.         x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))
  119.         print('x1_1:',x1_1.shape)
  120.         x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))
  121.         print('x0_2:',x0_2.shape)
  122.  
  123.         x3_0 = self.conv3_0(self.pool(x2_0))
  124.         print('x3_0:',x3_0.shape)
  125.         x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))
  126.         print('x2_1:',x2_1.shape)
  127.         x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))
  128.         print('x1_2:',x1_2.shape)
  129.         x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))
  130.         print('x0_3:',x0_3.shape)
  131.         x4_0 = self.conv4_0(self.pool(x3_0))
  132.         print('x4_0:',x4_0.shape)
  133.         x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
  134.         print('x3_1:',x3_1.shape)
  135.         x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))
  136.         print('x2_2:',x2_2.shape)
  137.         x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))
  138.         print('x1_3:',x1_3.shape)
  139.         x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))
  140.         print('x0_4:',x0_4.shape)
  141.  
  142.         if self.deep_supervision:
  143.             output1 = self.final1(x0_1)
  144.             output2 = self.final2(x0_2)
  145.             output3 = self.final3(x0_3)
  146.             output4 = self.final4(x0_4)
  147.             return [output1, output2, output3, output4]
  148.  
  149.         else:
  150.             output = self.final(x0_4)
  151.             return output

dataset.py为数据的预处理部分

  1. import os
  2.  
  3. import cv2
  4. import numpy as np
  5. import torch
  6. import torch.utils.data
  7.  
  8.  
  9. class Dataset(torch.utils.data.Dataset):
  10.     def __init__(self, img_ids, img_dir, mask_dir, img_ext, mask_ext, num_classes, transform=None):
  11.         """
  12.         Args:
  13.             img_ids (list): Image ids.
  14.             img_dir: Image file directory.
  15.             mask_dir: Mask file directory.
  16.             img_ext (str): Image file extension.
  17.             mask_ext (str): Mask file extension.
  18.             num_classes (int): Number of classes.
  19.             transform (Compose, optional): Compose transforms of albumentations. Defaults to None.
  20.         
  21.         Note:
  22.             Make sure to put the files as the following structure:
  23.             <dataset name>
  24.             ├── images
  25.             |   ├── 0a7e06.jpg
  26.             │   ├── 0aab0a.jpg
  27.             │   ├── 0b1761.jpg
  28.             │   ├── ...
  29.             |
  30.             └── masks
  31.                 ├── 0
  32.                 |   ├── 0a7e06.png
  33.                 |   ├── 0aab0a.png
  34.                 |   ├── 0b1761.png
  35.                 |   ├── ...
  36.                 |
  37.                 ├── 1
  38.                 |   ├── 0a7e06.png
  39.                 |   ├── 0aab0a.png
  40.                 |   ├── 0b1761.png
  41.                 |   ├── ...
  42.                 ...
  43.         """
  44.         self.img_ids = img_ids
  45.         self.img_dir = img_dir
  46.         self.mask_dir = mask_dir
  47.         self.img_ext = img_ext
  48.         self.mask_ext = mask_ext
  49.         self.num_classes = num_classes
  50.         self.transform = transform
  51.  
  52.     def __len__(self):
  53.         return len(self.img_ids)
  54.  
  55.     def __getitem__(self, idx):
  56.         img_id = self.img_ids[idx]
  57.         
  58.         img = cv2.imread(os.path.join(self.img_dir, img_id + self.img_ext))
  59.  
  60.         mask = []
  61.         for i in range(self.num_classes):
  62.             mask.append(cv2.imread(os.path.join(self.mask_dir, str(i),
  63.                         img_id + self.mask_ext), cv2.IMREAD_GRAYSCALE)[..., None])
  64.         mask = np.dstack(mask)
  65.  
  66.         if self.transform is not None:
  67.             augmented = self.transform(image=img, mask=mask)#这个包比较方便，能把mask也一并做掉
  68.             img = augmented['image']#参考https://github.com/albumentations-team/albumentations
  69.             mask = augmented['mask']
  70.         
  71.         img = img.astype('float32') / 255
  72.         img = img.transpose(2, 0, 1)
  73.         mask = mask.astype('float32') / 255
  74.         mask = mask.transpose(2, 0, 1)
  75.         
  76.         return img, mask, {'img_id': img_id}

losses.py

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4.  
  5. try:
  6.     from LovaszSoftmax.pytorch.lovasz_losses import lovasz_hinge
  7. except ImportError:
  8.     pass
  9.  
  10. __all__ = ['BCEDiceLoss', 'LovaszHingeLoss']
  11.  
  12.  
  13. class BCEDiceLoss(nn.Module):
  14.     def __init__(self):
  15.         super().__init__()
  16.  
  17.     def forward(self, input, target):
  18.         bce = F.binary_cross_entropy_with_logits(input, target)
  19.         smooth = 1e-5
  20.         input = torch.sigmoid(input)
  21.         num = target.size(0)
  22.         input = input.view(num, -1)
  23.         target = target.view(num, -1)
  24.         intersection = (input * target)
  25.         dice = (2. * intersection.sum(1) + smooth) / (input.sum(1) + target.sum(1) + smooth)
  26.         dice = 1 - dice.sum() / num
  27.         return 0.5 * bce + dice
  28.  
  29.  
  30. class LovaszHingeLoss(nn.Module):
  31.     def __init__(self):
  32.         super().__init__()
  33.  
  34.     def forward(self, input, target):
  35.         input = input.squeeze(1)
  36.         target = target.squeeze(1)
  37.         loss = lovasz_hinge(input, target, per_image=True)
  38.  
  39.         return loss

metrics.py 模型效果评价指标

  1. import numpy as np
  2. import torch
  3. import torch.nn.functional as F
  4.  
  5.  
  6. def iou_score(output, target):
  7.     smooth = 1e-5
  8.  
  9.     if torch.is_tensor(output):
  10.         output = torch.sigmoid(output).data.cpu().numpy()
  11.     if torch.is_tensor(target):
  12.         target = target.data.cpu().numpy()
  13.     output_ = output > 0.5
  14.     target_ = target > 0.5
  15.     intersection = (output_ & target_).sum()
  16.     union = (output_ | target_).sum()
  17.  
  18.     return (intersection + smooth) / (union + smooth)
  19.  
  20.  
  21. def dice_coef(output, target):
  22.     smooth = 1e-5
  23.  
  24.     output = torch.sigmoid(output).view(-1).data.cpu().numpy()
  25.     target = target.view(-1).data.cpu().numpy()
  26.     intersection = (output * target).sum()
  27.  
  28.     return (2. * intersection + smooth) / \
  29.         (output.sum() + target.sum() + smooth)

preprocess.py 数据标签的合并处理，将同一张图的多个标签数据合并为一张

  1. import os
  2. from glob import glob
  3.  
  4. import cv2
  5. import numpy as np
  6. from tqdm import tqdm
  7.  
  8.  
  9. def main():
  10.     img_size = 96
  11.  
  12.     paths = glob('inputs/stage1_train/*')
  13.  
  14.     os.makedirs('inputs/dsb2018_%d/images' % img_size, exist_ok=True)
  15.     os.makedirs('inputs/dsb2018_%d/masks/0' % img_size, exist_ok=True)
  16.  
  17.     for i in tqdm(range(len(paths))):
  18.         path = paths[i]
  19.         img = cv2.imread(os.path.join(path, 'images',
  20.                          os.path.basename(path) + '.png'))
  21.         mask = np.zeros((img.shape[0], img.shape[1]))
  22.         for mask_path in glob(os.path.join(path, 'masks', '*')):
  23.             mask_ = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) > 127
  24.             mask[mask_] = 1
  25.         if len(img.shape) == 2:
  26.             img = np.tile(img[..., None], (1, 1, 3))
  27.         if img.shape[2] == 4:
  28.             img = img[..., :3]
  29.         img = cv2.resize(img, (img_size, img_size))
  30.         mask = cv2.resize(mask, (img_size, img_size))
  31.         cv2.imwrite(os.path.join('inputs/dsb2018_%d/images' % img_size,
  32.                     os.path.basename(path) + '.png'), img)
  33.         cv2.imwrite(os.path.join('inputs/dsb2018_%d/masks/0' % img_size,
  34.                     os.path.basename(path) + '.png'), (mask * 255).astype('uint8'))
  35.  
  36.  
  37. if __name__ == '__main__':
  38.     main()

 utils.py 其它设置

  1. import argparse
  2.  
  3.  
  4. def str2bool(v):
  5.     if v.lower() in ['true', 1]:
  6.         return True
  7.     elif v.lower() in ['false', 0]:
  8.         return False
  9.     else:
  10.         raise argparse.ArgumentTypeError('Boolean value expected.')
  11.  
  12.  
  13. def count_params(model):
  14.     return sum(p.numel() for p in model.parameters() if p.requires_grad)
  15.  
  16.  
  17. class AverageMeter(object):
  18.     """Computes and stores the average and current value"""
  19.  
  20.     def __init__(self):
  21.         self.reset()
  22.  
  23.     def reset(self):
  24.         self.val = 0
  25.         self.avg = 0
  26.         self.sum = 0
  27.         self.count = 0
  28.  
  29.     def update(self, val, n=1):
  30.         self.val = val
  31.         self.sum += val * n
  32.         self.count += n
  33.         self.avg = self.sum / self.count

train.py 模型训练

  1. import argparse
  2. import os
  3. from collections import OrderedDict
  4. from glob import glob
  5.  
  6. import pandas as pd
  7. import torch
  8. import torch.backends.cudnn as cudnn
  9. import torch.nn as nn
  10. import torch.optim as optim
  11. import yaml
  12. import albumentations as albu
  13. from albumentations.augmentations import transforms
  14. from albumentations.core.composition import Compose, OneOf
  15. from sklearn.model_selection import train_test_split
  16. from torch.optim import lr_scheduler
  17. from tqdm import tqdm
  18.  
  19. import archs
  20. import losses
  21. from dataset import Dataset
  22. from metrics import iou_score
  23. from utils import AverageMeter, str2bool
  24.  
  25. ARCH_NAMES = archs.__all__
  26. LOSS_NAMES = losses.__all__
  27. LOSS_NAMES.append('BCEWithLogitsLoss')
  28.  
  29. """
  31. 指定参数：
  32. --dataset dsb2018_96 
  33. --arch NestedUNet
  35. """
  36.  
  37. def parse_args():
  38.     parser = argparse.ArgumentParser()
  39.  
  40.     parser.add_argument('--name', default=None,
  41.                         help='model name: (default: arch+timestamp)')
  42.     parser.add_argument('--epochs', default=100, type=int, metavar='N',
  43.                         help='number of total epochs to run')
  44.     parser.add_argument('-b', '--batch_size', default=8, type=int,
  45.                         metavar='N', help='mini-batch size (default: 16)')
  46.     
  47.     # model
  48.     parser.add_argument('--arch', '-a', metavar='ARCH', default='NestedUNet',
  49.                         choices=ARCH_NAMES,
  50.                         help='model architecture: ' +
  51.                         ' | '.join(ARCH_NAMES) +
  52.                         ' (default: NestedUNet)')
  53.     parser.add_argument('--deep_supervision', default=False, type=str2bool)
  54.     parser.add_argument('--input_channels', default=3, type=int,
  55.                         help='input channels')
  56.     parser.add_argument('--num_classes', default=1, type=int,
  57.                         help='number of classes')
  58.     parser.add_argument('--input_w', default=96, type=int,
  59.                         help='image width')
  60.     parser.add_argument('--input_h', default=96, type=int,
  61.                         help='image height')
  62.     
  63.     # loss
  64.     parser.add_argument('--loss', default='BCEDiceLoss',
  65.                         choices=LOSS_NAMES,
  66.                         help='loss: ' +
  67.                         ' | '.join(LOSS_NAMES) +
  68.                         ' (default: BCEDiceLoss)')
  69.     
  70.     # dataset
  71.     parser.add_argument('--dataset', default='dsb2018_96',
  72.                         help='dataset name')
  73.     parser.add_argument('--img_ext', default='.png',
  74.                         help='image file extension')
  75.     parser.add_argument('--mask_ext', default='.png',
  76.                         help='mask file extension')
  77.  
  78.     # optimizer
  79.     parser.add_argument('--optimizer', default='SGD',
  80.                         choices=['Adam', 'SGD'],
  81.                         help='loss: ' +
  82.                         ' | '.join(['Adam', 'SGD']) +
  83.                         ' (default: Adam)')
  84.     parser.add_argument('--lr', '--learning_rate', default=1e-3, type=float,
  85.                         metavar='LR', help='initial learning rate')
  86.     parser.add_argument('--momentum', default=0.9, type=float,
  87.                         help='momentum')
  88.     parser.add_argument('--weight_decay', default=1e-4, type=float,
  89.                         help='weight decay')
  90.     parser.add_argument('--nesterov', default=False, type=str2bool,
  91.                         help='nesterov')
  92.  
  93.     # scheduler
  94.     parser.add_argument('--scheduler', default='CosineAnnealingLR',
  95.                         choices=['CosineAnnealingLR', 'ReduceLROnPlateau', 'MultiStepLR', 'ConstantLR'])
  96.     parser.add_argument('--min_lr', default=1e-5, type=float,
  97.                         help='minimum learning rate')
  98.     parser.add_argument('--factor', default=0.1, type=float)
  99.     parser.add_argument('--patience', default=2, type=int)
  100.     parser.add_argument('--milestones', default='1,2', type=str)
  101.     parser.add_argument('--gamma', default=2/3, type=float)
  102.     parser.add_argument('--early_stopping', default=-1, type=int,
  103.                         metavar='N', help='early stopping (default: -1)')
  104.     
  105.     parser.add_argument('--num_workers', default=0, type=int)
  106.  
  107.     config = parser.parse_args()
  108.  
  109.     return config
  110.  
  111.  
  112. def train(config, train_loader, model, criterion, optimizer):
  113.     avg_meters = {'loss': AverageMeter(),
  114.                   'iou': AverageMeter()}
  115.  
  116.     model.train()
  117.  
  118.     pbar = tqdm(total=len(train_loader))
  119.     for input, target, _ in train_loader:
  120.         input = input.cuda()
  121.         target = target.cuda()
  122.  
  123.         # compute output
  124.         if config['deep_supervision']:
  125.             outputs = model(input)
  126.             loss = 0
  127.             for output in outputs:
  128.                 loss += criterion(output, target)
  129.             loss /= len(outputs)
  130.             iou = iou_score(outputs[-1], target)
  131.         else:
  132.             output = model(input)
  133.             loss = criterion(output, target)
  134.             iou = iou_score(output, target)
  135.  
  136.         # compute gradient and do optimizing step
  137.         optimizer.zero_grad()
  138.         loss.backward()
  139.         optimizer.step()
  140.  
  141.         avg_meters['loss'].update(loss.item(), input.size(0))
  142.         avg_meters['iou'].update(iou, input.size(0))
  143.  
  144.         postfix = OrderedDict([
  145.             ('loss', avg_meters['loss'].avg),
  146.             ('iou', avg_meters['iou'].avg),
  147.         ])
  148.         pbar.set_postfix(postfix)
  149.         pbar.update(1)
  150.     pbar.close()
  151.  
  152.     return OrderedDict([('loss', avg_meters['loss'].avg),
  153.                         ('iou', avg_meters['iou'].avg)])
  154.  
  155.  
  156. def validate(config, val_loader, model, criterion):
  157.     avg_meters = {'loss': AverageMeter(),
  158.                   'iou': AverageMeter()}
  159.  
  160.     # switch to evaluate mode
  161.     model.eval()
  162.  
  163.     with torch.no_grad():
  164.         pbar = tqdm(total=len(val_loader))
  165.         for input, target, _ in val_loader:
  166.             input = input.cuda()
  167.             target = target.cuda()
  168.  
  169.             # compute output
  170.             if config['deep_supervision']:
  171.                 outputs = model(input)
  172.                 loss = 0
  173.                 for output in outputs:
  174.                     loss += criterion(output, target)
  175.                 loss /= len(outputs)
  176.                 iou = iou_score(outputs[-1], target)
  177.             else:
  178.                 output = model(input)
  179.                 loss = criterion(output, target)
  180.                 iou = iou_score(output, target)
  181.  
  182.             avg_meters['loss'].update(loss.item(), input.size(0))
  183.             avg_meters['iou'].update(iou, input.size(0))
  184.  
  185.             postfix = OrderedDict([
  186.                 ('loss', avg_meters['loss'].avg),
  187.                 ('iou', avg_meters['iou'].avg),
  188.             ])
  189.             pbar.set_postfix(postfix)
  190.             pbar.update(1)
  191.         pbar.close()
  192.  
  193.     return OrderedDict([('loss', avg_meters['loss'].avg),
  194.                         ('iou', avg_meters['iou'].avg)])
  195.  
  196.  
  197. def main():
  198.     config = vars(parse_args())
  199.  
  200.     if config['name'] is None:
  201.         if config['deep_supervision']:
  202.             config['name'] = '%s_%s_wDS' % (config['dataset'], config['arch'])
  203.         else:
  204.             config['name'] = '%s_%s_woDS' % (config['dataset'], config['arch'])
  205.     os.makedirs('models/%s' % config['name'], exist_ok=True)
  206.  
  207.     print('-' * 20)
  208.     for key in config:
  209.         print('%s: %s' % (key, config[key]))
  210.     print('-' * 20)
  211.  
  212.     with open('models/%s/config.yml' % config['name'], 'w') as f:
  213.         yaml.dump(config, f)
  214.  
  215.     # define loss function (criterion)
  216.     if config['loss'] == 'BCEWithLogitsLoss':
  217.         criterion = nn.BCEWithLogitsLoss().cuda()#WithLogits 就是先将输出结果经过sigmoid再交叉熵
  218.     else:
  219.         criterion = losses.__dict__[config['loss']]().cuda()
  220.  
  221.     cudnn.benchmark = True
  222.  
  223.     # create model
  224.     print("=> creating model %s" % config['arch'])
  225.     model = archs.__dict__[config['arch']](config['num_classes'],
  226.                                            config['input_channels'],
  227.                                            config['deep_supervision'])
  228.  
  229.     model = model.cuda()
  230.  
  231.     params = filter(lambda p: p.requires_grad, model.parameters())
  232.     if config['optimizer'] == 'Adam':
  233.         optimizer = optim.Adam(
  234.             params, lr=config['lr'], weight_decay=config['weight_decay'])
  235.     elif config['optimizer'] == 'SGD':
  236.         optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'],
  237.                               nesterov=config['nesterov'], weight_decay=config['weight_decay'])
  238.     else:
  239.         raise NotImplementedError
  240.  
  241.     if config['scheduler'] == 'CosineAnnealingLR':
  242.         scheduler = lr_scheduler.CosineAnnealingLR(
  243.             optimizer, T_max=config['epochs'], eta_min=config['min_lr'])
  244.     elif config['scheduler'] == 'ReduceLROnPlateau':
  245.         scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=config['factor'], patience=config['patience'],
  246.                                                    verbose=1, min_lr=config['min_lr'])
  247.     elif config['scheduler'] == 'MultiStepLR':
  248.         scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma'])
  249.     elif config['scheduler'] == 'ConstantLR':
  250.         scheduler = None
  251.     else:
  252.         raise NotImplementedError
  253.  
  254.     # Data loading code
  255.     img_ids = glob(os.path.join('inputs', config['dataset'], 'images', '*' + config['img_ext']))
  256.     img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
  257.  
  258.     train_img_ids, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41)
  259.     #数据增强：需要安装albumentations包
  260.     train_transform = Compose([
  261.         #角度旋转
  262.         albu.RandomRotate90(),
  263.         #图像翻转
  264.         albu.Flip(),
  265.         OneOf([
  266.             transforms.HueSaturationValue(),
  267.             transforms.RandomBrightness(),
  268.             transforms.RandomContrast(),
  269.         ], p=1),#按照归一化的概率选择执行哪一个
  270.         albu.Resize(config['input_h'], config['input_w']),
  271.         albu.Normalize(),
  272.     ])
  273.  
  274.     val_transform = Compose([
  275.         albu.Resize(config['input_h'], config['input_w']),
  276.         albu.Normalize(),
  277.     ])
  278.  
  279.     train_dataset = Dataset(
  280.         img_ids=train_img_ids,
  281.         img_dir=os.path.join('inputs', config['dataset'], 'images'),
  282.         mask_dir=os.path.join('inputs', config['dataset'], 'masks'),
  283.         img_ext=config['img_ext'],
  284.         mask_ext=config['mask_ext'],
  285.         num_classes=config['num_classes'],
  286.         transform=train_transform)
  287.     val_dataset = Dataset(
  288.         img_ids=val_img_ids,
  289.         img_dir=os.path.join('inputs', config['dataset'], 'images'),
  290.         mask_dir=os.path.join('inputs', config['dataset'], 'masks'),
  291.         img_ext=config['img_ext'],
  292.         mask_ext=config['mask_ext'],
  293.         num_classes=config['num_classes'],
  294.         transform=val_transform)
  295.  
  296.     train_loader = torch.utils.data.DataLoader(
  297.         train_dataset,
  298.         batch_size=config['batch_size'],
  299.         shuffle=True,
  300.         num_workers=config['num_workers'],
  301.         drop_last=True)#不能整除的batch是否就不要了
  302.     val_loader = torch.utils.data.DataLoader(
  303.         val_dataset,
  304.         batch_size=config['batch_size'],
  305.         shuffle=False,
  306.         num_workers=config['num_workers'],
  307.         drop_last=False)
  308.  
  309.     log = OrderedDict([
  310.         ('epoch', []),
  311.         ('lr', []),
  312.         ('loss', []),
  313.         ('iou', []),
  314.         ('val_loss', []),
  315.         ('val_iou', []),
  316.     ])
  317.  
  318.     best_iou = 0
  319.     trigger = 0
  320.     for epoch in range(config['epochs']):
  321.         print('Epoch [%d/%d]' % (epoch, config['epochs']))
  322.  
  323.         # train for one epoch
  324.         train_log = train(config, train_loader, model, criterion, optimizer)
  325.         # evaluate on validation set
  326.         val_log = validate(config, val_loader, model, criterion)
  327.  
  328.         if config['scheduler'] == 'CosineAnnealingLR':
  329.             scheduler.step()
  330.         elif config['scheduler'] == 'ReduceLROnPlateau':
  331.             scheduler.step(val_log['loss'])
  332.  
  333.         print('loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f'
  334.               % (train_log['loss'], train_log['iou'], val_log['loss'], val_log['iou']))
  335.  
  336.         log['epoch'].append(epoch)
  337.         log['lr'].append(config['lr'])
  338.         log['loss'].append(train_log['loss'])
  339.         log['iou'].append(train_log['iou'])
  340.         log['val_loss'].append(val_log['loss'])
  341.         log['val_iou'].append(val_log['iou'])
  342.  
  343.         pd.DataFrame(log).to_csv('models/%s/log.csv' %
  344.                                  config['name'], index=False)
  345.  
  346.         trigger += 1
  347.  
  348.         if val_log['iou'] > best_iou:
  349.             torch.save(model.state_dict(), 'models/%s/model.pth' %
  350.                        config['name'])
  351.             best_iou = val_log['iou']
  352.             print("=> saved best model")
  353.             trigger = 0
  354.  
  355.         # early stopping
  356.         if config['early_stopping'] >= 0 and trigger >= config['early_stopping']:
  357.             print("=> early stopping")
  358.             break
  359.  
  360.         torch.cuda.empty_cache()
  361.  
  362.  
  363. if __name__ == '__main__':
  364.     main()

四，模型结果：

 五：注意事项以及常见问题

安装增加模块albumentations，主要为数据增强模块，方便快捷

pip install albumentations

常见问题：

AttributeError: module ‘cv2’ has no attribute ‘gapi_wip_gst_GStreamerPipeline’

解决：opencv-python-headless和opencv-python的版本对应即可

pip install opencv-python-headless==4.2.0.32 -i https://pypi.tuna.tsinghua.edu.cn/simple

问题2

AttributeError: module ‘albumentations.augmentations.transforms’ has no attribute ‘RandomRotate90’

解决：直接导入import albumentations as albu

 

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/羊村懒王/article/detail/351112
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号