• 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下载

  • 相关阅读:
    大学新生开学数码装备，2022年值得买的几款数码好物
    阿里P8大能倾力编撰的“Java 进阶面试手册”，助力跳槽外包毕业生秋招收获大厂offer
    区块链技术研究探讨
    【路径优化】基于A*算法的路径优化问题（Matlab代码实现）
    Ubuntu将图标放置于任务栏
    Compiling and Loading
    Redis为什么要使用SDS作为基本数据结构
    国内厨电行业数据分析
    2023NOIP A层联测19-多边形
    windows安装npm教程
  • 原文地址：https://blog.csdn.net/qq_20373723/article/details/126450108