• TransFuse跑自己的数据集

     原始链接如下

    https://github.com/Rayicer/TransFusehttps://github.com/Rayicer/TransFuse这个复现不难，下面给下关键步骤

    1.数据准备

     一级目录，数据分为训练（train）、验证（val）、测试（test），如果数据少的话验证和测试一样也行，但是一定都要有，因为代码里有个地方写死了，当然，看完代码比较熟悉的话可以自己改掉。

    另外，所有文件夹的名字务必和我一样，这个代码里也写死了，熟悉的可以自己改。

     二级目录

    三级目录，注意这个是二分类，标签的值是0和255

     

     

     2.数据加载

    复现GitHub的关键其实就是数据加载部分，这个项目原本是读取了npy的，直接改成读图片就好了，下面是我改的util/dataloader.py，这里需要注意的是图片的大小，原始的数据大小是256*192，我这里把我512*512的图resize到了256*192，这个规定的大小要更改有点麻烦，得自己去模型文件里探索下，留给大家自己来吧。

    1. import os
    2. from PIL import Image
    3. import torch
    4. import torch.utils.data as data
    5. import torchvision.transforms as transforms
    6. import numpy as np
    7. import matplotlib.pyplot as plt
    8. import albumentations as A
    9. import cv2
    10.  
    11.  
    12. class SkinDataset(data.Dataset):
    13.     """
    14.     dataloader for skin lesion segmentation tasks
    15.     """
    16.     def __init__(self, image_root, gt_root):
    17.         imgs_full_root = os.path.join(image_root, 'images')
    18.         labs_full_root = os.path.join(gt_root, 'labels')
    19.         image_names = os.listdir(imgs_full_root)
    20.         label_names = os.listdir(labs_full_root)
    21.         self.images = [os.path.join(imgs_full_root, im_name) for im_name in image_names]
    22.         self.gts = [os.path.join(labs_full_root, lab_name) for lab_name in label_names]
    23.         self.size = len(self.images)
    24.  
    25.         self.img_transform = transforms.Compose([
    26.             transforms.ToTensor(),
    27.             transforms.Normalize([0.485, 0.456, 0.406],
    28.                                  [0.229, 0.224, 0.225])
    29.         ])
    30.         self.gt_transform = transforms.Compose([
    31.             transforms.ToTensor()])
    32.         
    33.         self.transform = A.Compose(
    34.             [
    35.                 A.ShiftScaleRotate(shift_limit=0.15, scale_limit=0.15, rotate_limit=25, p=0.5, border_mode=0),
    36.                 A.ColorJitter(),
    37.                 A.HorizontalFlip(),
    38.                 A.VerticalFlip()
    39.             ]
    40.         )
    41.  
    42.     def __getitem__(self, index):
    43.         
    44.         image_idx = self.images[index]
    45.         image = cv2.imread(image_idx)
    46.         image = cv2.resize(image, (256, 192))
    47.         # image = image.transpose(2,1,0)
    48.         gt_idx = self.gts[index]
    49.         gt = cv2.imread(gt_idx, 0)
    50.         gt = cv2.resize(gt, (256, 192))
    51.         gt = gt / 255.0
    52.  
    53.         transformed = self.transform(image=image, mask=gt)
    54.         image = self.img_transform(transformed['image'])
    55.         # image = image / 255.0
    56.         gt = self.gt_transform(transformed['mask'])
    57.         return image, gt
    58.  
    59.     def __len__(self):
    60.         return self.size
    61.  
    62.  
    63. def get_loader(image_root, gt_root, batchsize, shuffle=True, num_workers=4, pin_memory=True):
    64.  
    65.     dataset = SkinDataset(image_root, gt_root)
    66.     data_loader = data.DataLoader(dataset=dataset,
    67.                                   batch_size=batchsize,
    68.                                   shuffle=shuffle,
    69.                                   num_workers=num_workers,
    70.                                   pin_memory=pin_memory)
    71.     return data_loader
    72.  
    73.  
    74. class test_dataset:
    75.     def __init__(self, image_root, gt_root):
    76.         imgs_full_root = os.path.join(image_root, 'images')
    77.         labs_full_root = os.path.join(gt_root, 'labels')
    78.         image_names = os.listdir(imgs_full_root)
    79.         label_names = os.listdir(labs_full_root)
    80.         self.images = [os.path.join(imgs_full_root, im_name) for im_name in image_names]
    81.         self.gts = [os.path.join(labs_full_root, lab_name) for lab_name in label_names]
    82.  
    83.         self.transform = transforms.Compose([
    84.             transforms.ToTensor(),
    85.             transforms.Normalize([0.485, 0.456, 0.406],
    86.                                  [0.229, 0.224, 0.225])
    87.             ])
    88.         self.gt_transform = transforms.ToTensor()
    89.         self.size = len(self.images)
    90.         self.index = 0
    91.  
    92.     def load_data(self):
    93.         image_idx = self.images[self.index]
    94.         image = cv2.imread(image_idx)
    95.         image = cv2.resize(image, (256, 192), interpolation=cv2.INTER_NEAREST)
    96.         image = self.transform(image).unsqueeze(0)
    97.         gt_idx = self.gts[self.index]
    98.         gt = cv2.imread(gt_idx, 0)
    99.         gt = cv2.resize(gt, (256, 192), interpolation=cv2.INTER_NEAREST)
    100.         gt = gt/255.0
    101.         self.index += 1
    102.  
    103.         return image, gt
    104.  
    105.  
    106. if __name__ == '__main__':
    107.     path = 'data/'
    108.     tt = SkinDataset(path+'data_train.npy', path+'mask_train.npy')
    109.  
    110.     for i in range(50):
    111.         img, gt = tt.__getitem__(i)
    112.  
    113.         img = torch.transpose(img, 0, 1)
    114.         img = torch.transpose(img, 1, 2)
    115.         img = img.numpy()
    116.         gt = gt.numpy()
    117.  
    118.         plt.imshow(img)
    119.         plt.savefig('vis/'+str(i)+".jpg")
    120.  
    121.         plt.imshow(gt[0])
    122.         plt.savefig('vis/'+str(i)+'_gt.jpg')

    3.训练

    train_isic.py主要改前面的参数部分，还有数据加载路径

    1. import torch
    2. from torch.autograd import Variable
    3. import argparse
    4. from datetime import datetime
    5. from lib.TransFuse import TransFuse_S
    6. from utils.dataloader import get_loader, test_dataset
    7. from utils.utils import AvgMeter
    8. import torch.nn.functional as F
    9. import numpy as np
    10. import matplotlib.pyplot as plt
    11. from test_isic import mean_dice_np, mean_iou_np
    12. import os
    13.  
    14.  
    15. def structure_loss(pred, mask):
    16.     weit = 1 + 5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask)
    17.     wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
    18.     wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))
    19.  
    20.     pred = torch.sigmoid(pred)
    21.     inter = ((pred * mask)*weit).sum(dim=(2, 3))
    22.     union = ((pred + mask)*weit).sum(dim=(2, 3))
    23.     wiou = 1 - (inter + 1)/(union - inter+1)
    24.     return (wbce + wiou).mean()
    25.  
    26.  
    27. def train(train_loader, model, optimizer, epoch, best_loss):
    28.     model.train()
    29.     loss_record2, loss_record3, loss_record4 = AvgMeter(), AvgMeter(), AvgMeter()
    30.     accum = 0
    31.     for i, pack in enumerate(train_loader, start=1):
    32.         # ---- data prepare ----
    33.         images, gts = pack
    34.         images = Variable(images).cuda()
    35.         gts = Variable(gts).cuda()
    36.  
    37.         # ---- forward ----
    38.         lateral_map_4, lateral_map_3, lateral_map_2 = model(images)
    39.  
    40.         # ---- loss function ----
    41.         loss4 = structure_loss(lateral_map_4, gts)
    42.         loss3 = structure_loss(lateral_map_3, gts)
    43.         loss2 = structure_loss(lateral_map_2, gts)
    44.  
    45.         loss = 0.5 * loss2 + 0.3 * loss3 + 0.2 * loss4
    46.  
    47.         # ---- backward ----
    48.         loss.backward() 
    49.         torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_norm)
    50.         optimizer.step()
    51.         optimizer.zero_grad()
    52.  
    53.         # ---- recording loss ----
    54.         loss_record2.update(loss2.data, opt.batchsize)
    55.         loss_record3.update(loss3.data, opt.batchsize)
    56.         loss_record4.update(loss4.data, opt.batchsize)
    57.  
    58.         # ---- train visualization ----
    59.         if i % 20 == 0 or i == total_step:
    60.             print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], '
    61.                   '[lateral-2: {:.4f}, lateral-3: {:0.4f}, lateral-4: {:0.4f}]'.  
    62.                   format(datetime.now(), epoch, opt.epoch, i, total_step,
    63.                          loss_record2.show(), loss_record3.show(), loss_record4.show()))
    64.  
    65.     print('lr: ', optimizer.param_groups[0]['lr'])
    66.     save_path = 'snapshots/{}/'.format(opt.train_save)
    67.     os.makedirs(save_path, exist_ok=True)
    68.     if (epoch+1) % 1 == 0:
    69.         meanloss = test(model, opt.test_path)
    70.         if meanloss < best_loss:
    71.             print('new best loss: ', meanloss)
    72.             best_loss = meanloss
    73.             torch.save(model.state_dict(), save_path + 'TransFuse-%d.pth' % epoch)
    74.             print('[Saving Snapshot:]', save_path + 'TransFuse-%d.pth'% epoch)
    75.     return best_loss
    76.  
    77.  
    78. def test(model, path):
    79.  
    80.     model.eval()
    81.     mean_loss = []
    82.  
    83.     for s in ['val', 'test']:
    84.         image_root = '{}/{}'.format(path, s)
    85.         gt_root = '{}/{}'.format(path, s)
    86.         test_loader = test_dataset(image_root, gt_root)
    87.  
    88.         dice_bank = []
    89.         iou_bank = []
    90.         loss_bank = []
    91.         acc_bank = []
    92.  
    93.         for i in range(test_loader.size):
    94.             image, gt = test_loader.load_data()
    95.             image = image.cuda()
    96.  
    97.             with torch.no_grad():
    98.                 _, _, res = model(image)
    99.             loss = structure_loss(res, torch.tensor(gt).unsqueeze(0).unsqueeze(0).cuda())
    100.  
    101.             res = res.sigmoid().data.cpu().numpy().squeeze()
    102.             gt = 1*(gt>0.5)            
    103.             res = 1*(res > 0.5)
    104.  
    105.             dice = mean_dice_np(gt, res)
    106.             iou = mean_iou_np(gt, res)
    107.             acc = np.sum(res == gt) / (res.shape[0]*res.shape[1])
    108.  
    109.             loss_bank.append(loss.item())
    110.             dice_bank.append(dice)
    111.             iou_bank.append(iou)
    112.             acc_bank.append(acc)
    113.             
    114.         print('{} Loss: {:.4f}, Dice: {:.4f}, IoU: {:.4f}, Acc: {:.4f}'.
    115.             format(s, np.mean(loss_bank), np.mean(dice_bank), np.mean(iou_bank), np.mean(acc_bank)))
    116.  
    117.         mean_loss.append(np.mean(loss_bank))
    118.  
    119.     return mean_loss[0] 
    120.  
    121.  
    122. if __name__ == '__main__':
    123.     parser = argparse.ArgumentParser()
    124.     parser.add_argument('--epoch', type=int, default=80, help='epoch number')
    125.     parser.add_argument('--lr', type=float, default=7e-5, help='learning rate')
    126.     parser.add_argument('--batchsize', type=int, default=8, help='training batch size')
    127.     parser.add_argument('--grad_norm', type=float, default=2.0, help='gradient clipping norm')
    128.     parser.add_argument('--train_path', type=str,
    129.                         default='./data/build512/', help='path to train dataset')
    130.     parser.add_argument('--test_path', type=str,
    131.                         default='./data/build512/', help='path to test dataset')
    132.     parser.add_argument('--train_save', type=str, default='TransFuse_S')
    133.     parser.add_argument('--beta1', type=float, default=0.5, help='beta1 of adam optimizer')
    134.     parser.add_argument('--beta2', type=float, default=0.999, help='beta2 of adam optimizer')
    135.  
    136.     opt = parser.parse_args()
    137.  
    138.     # ---- build models ----
    139.     model = TransFuse_S(pretrained=True).cuda()
    140.     params = model.parameters()
    141.  
    142.     optimizer = torch.optim.Adam(params, opt.lr, betas=(opt.beta1, opt.beta2))
    143.  
    144.  
    145.     image_root = '{}/train'.format(opt.train_path)
    146.     gt_root = '{}/train'.format(opt.train_path)
    147.  
    148.     train_loader = get_loader(image_root, gt_root, batchsize=opt.batchsize)
    149.     total_step = len(train_loader)
    150.  
    151.     print("#"*20, "Start Training", "#"*20)
    152.  
    153.     best_loss = 1e5
    154.     for epoch in range(1, opt.epoch + 1):
    155.         best_loss = train(train_loader, model, optimizer, epoch, best_loss)

    4.预测验证

    test_isic.py预测也主要是改了路径，细节自己对比下吧

    1. import torch
    2. import torch.nn.functional as F
    3. import numpy as np
    4. import os, argparse
    5. from lib.TransFuse import TransFuse_S
    6. from utils.dataloader import test_dataset
    7. import imageio
    8. import cv2
    9.  
    10.  
    11. def mean_iou_np(y_true, y_pred, **kwargs):
    12.     """
    13.     compute mean iou for binary segmentation map via numpy
    14.     """
    15.     axes = (0, 1) 
    16.     intersection = np.sum(np.abs(y_pred * y_true), axis=axes) 
    17.     mask_sum = np.sum(np.abs(y_true), axis=axes) + np.sum(np.abs(y_pred), axis=axes)
    18.     union = mask_sum  - intersection 
    19.     
    20.     smooth = .001
    21.     iou = (intersection + smooth) / (union + smooth)
    22.     return iou
    23.  
    24.  
    25. def mean_dice_np(y_true, y_pred, **kwargs):
    26.     """
    27.     compute mean dice for binary segmentation map via numpy
    28.     """
    29.     axes = (0, 1) # W,H axes of each image
    30.     intersection = np.sum(np.abs(y_pred * y_true), axis=axes) 
    31.     mask_sum = np.sum(np.abs(y_true), axis=axes) + np.sum(np.abs(y_pred), axis=axes)
    32.     
    33.     smooth = .001
    34.     dice = 2*(intersection + smooth)/(mask_sum + smooth)
    35.     return dice
    36.  
    37.  
    38. if __name__ == '__main__':
    39.     parser = argparse.ArgumentParser()
    40.     parser.add_argument('--ckpt_path', type=str, default='snapshots/TransFuse_S/TransFuse-6.pth')
    41.     parser.add_argument('--test_path', type=str,
    42.                         default='./data/build512', help='path to test dataset')
    43.     parser.add_argument('--save_path', type=str, default='./predict', help='path to save inference segmentation')
    44.  
    45.     opt = parser.parse_args()
    46.  
    47.     model = TransFuse_S().cuda()
    48.     model.load_state_dict(torch.load(opt.ckpt_path))
    49.     model.cuda()
    50.     model.eval()
    51.  
    52.     if opt.save_path is not None:
    53.         os.makedirs(opt.save_path, exist_ok=True)
    54.  
    55.     print('evaluating model: ', opt.ckpt_path)
    56.  
    57.     image_root = '{}/test'.format(opt.test_path)
    58.     gt_root = '{}/test'.format(opt.test_path)
    59.     test_loader = test_dataset(image_root, gt_root)
    60.  
    61.     dice_bank = []
    62.     iou_bank = []
    63.     acc_bank = []
    64.  
    65.     for i in range(test_loader.size):
    66.         image, gt = test_loader.load_data()
    67.         # gt = 1*(gt>0.5)
    68.         gt[gt > 0] = 1
    69.         image = image.cuda()
    70.  
    71.         with torch.no_grad():
    72.             _, _, res = model(image)
    73.  
    74.         res = res.sigmoid().data.cpu().numpy().squeeze()
    75.         # res = 1*(res > 0.5)
    76.         res[res > 0.5] = 1
    77.         res[res <= 0.5] = 0
    78.  
    79.         dice = mean_dice_np(gt, res)
    80.         iou = mean_iou_np(gt, res)
    81.         acc = np.sum(res == gt) / (res.shape[0]*res.shape[1])
    82.  
    83.         acc_bank.append(acc)
    84.         dice_bank.append(dice)
    85.         iou_bank.append(iou)
    86.  
    87.         if opt.save_path is not None:
    88.             res = 255 * res
    89.             gt = 255 * gt
    90.             cv2.imwrite(opt.save_path+'/'+str(i)+'_pred.jpg', res)
    91.             cv2.imwrite(opt.save_path+'/'+str(i)+'_gt.jpg', gt)
    92.  
    93.     print('Dice: {:.4f}, IoU: {:.4f}, Acc: {:.4f}'.
    94.         format(np.mean(dice_bank), np.mean(iou_bank), np.mean(acc_bank)))

    5.结果

     预测IOU大概在0.6左右，效果还需要自己调，这个项目里的学习率策略还可以加一下，还有数据增强

    积分下载：

    TransFuse复现-深度学习文档类资源-CSDN下载

  • 相关阅读:
    开源博客项目Blog .NET Core源码学习（16：App.Hosting项目结构分析-4）
    从C++CLI工程的依赖库引用问题看.Net加载程序集机制
    6G网络需求、架构及技术趋势
    rsync+inotify远程同步
    【PACS系统源码】与医院HIS系统双向数据交换，实现医学影像集成与影像后处理功能
    CentOS8安装部署DzzOffice协同办公平台
    【测试】最火的测试开发学习路线内容再次大更新，助力通关大厂测开
    并列句------六级
    json数据刨根究底
    C语言第五章第4节用do...while语句实现循环学习导案
  • 原文地址：https://blog.csdn.net/qq_20373723/article/details/126450108