• 深度补偿模型sparse-to-dense测试

    原始链接

    GitHub - fangchangma/sparse-to-dense.pytorch: ICRA 2018 "Sparse-to-Dense: Depth Prediction from Sparse Depth Samples and a Single Image" (PyTorch Implementation)

     1.数据说明

            第一列是场景原图，第二列是稀疏数据，第三列是稠密数据，第四列是该模型的预测结果

    该模型有三种训练模式：1.第一列RGB作为图像输入，第三列作为标签；2.第一列和第二列合并为4通道数据作为图像输入，第三列作为标签；3.第二列作为输入，第三列作为标签。

    其中第二列是通过第三列采样得到。采样方式有两种，在dataloaders/dense_to_sparse.py脚本中，如果要跑自己的数据集那么需要准备的是第一列和第三列数据。

     这个数据有30G左右，比较大，我下载了，网盘链接链接：https://pan.baidu.com/s/1SzQhDVZBJSy9gnMkr4UCMQ 
    提取码：tpql 
    --来自百度网盘超级会员V6的分享

    解压后如下：

             这里的h5文件其实就是数据的一种存储形式而已，内部结构和字典类似，在代码里有加载的函数（dataloaders/dataloader.py脚本中的h5_loader函数），包含了rgb和对应的depth数据。

    2.跑val数据（需要输入标签，rgbd模式）

            该项目通过设置evaluate模式来做评价，下载好数据、模型后，直接创建数据文件夹data放解压的数据即可。然后命令行输入python main.py --evaluate model_best.pth会自动创建结果文件夹results，以及产生一个拼接的长图comparison_7.png

     3.自建test数据测试（无需输入标签，rgbd模式）

             这个需要得到和val相同的效果，并且批量跑的情况下，需要在val的基础上改，改的地方比较多，下面我把改了的脚本都贴出来。

            （1）数据加载部分

            nyu_dataloader.py

    1. import numpy as np
    2. import dataloaders.transforms as transforms
    3. from dataloaders.dataloader import MyDataloader
    4.  
    5. iheight, iwidth = 480, 640 # raw image size
    6.  
    7. class NYUDataset(MyDataloader):
    8.     def __init__(self, root, type, sparsifier=None, modality='rgb'):
    9.         super(NYUDataset, self).__init__(root, type, sparsifier, modality)
    10.         self.output_size = (228, 304)
    11.  
    12.     def train_transform(self, rgb, depth):
    13.         s = np.random.uniform(1.0, 1.5) # random scaling
    14.         depth_np = depth / s
    15.         angle = np.random.uniform(-5.0, 5.0) # random rotation degrees
    16.         do_flip = np.random.uniform(0.0, 1.0) < 0.5 # random horizontal flip
    17.  
    18.         # perform 1st step of data augmentation
    19.         transform = transforms.Compose([
    20.             transforms.Resize(250.0 / iheight), # this is for computational efficiency, since rotation can be slow
    21.             transforms.Rotate(angle),
    22.             transforms.Resize(s),
    23.             transforms.CenterCrop(self.output_size),
    24.             transforms.HorizontalFlip(do_flip)
    25.         ])
    26.         rgb_np = transform(rgb)
    27.         rgb_np = self.color_jitter(rgb_np) # random color jittering
    28.         rgb_np = np.asfarray(rgb_np, dtype='float') / 255
    29.         depth_np = transform(depth_np)
    30.  
    31.         return rgb_np, depth_np
    32.  
    33.     def val_transform(self, rgb, depth):
    34.         depth_np = depth
    35.         transform = transforms.Compose([
    36.             transforms.Resize(240.0 / iheight),
    37.             transforms.CenterCrop(self.output_size),
    38.         ])
    39.         rgb_np = transform(rgb)
    40.         rgb_np = np.asfarray(rgb_np, dtype='float') / 255
    41.         depth_np = transform(depth_np)
    42.         return rgb_np, depth_np
    43.  
    44.     def test_transform(self, rgb, depth):
    45.         depth_np = depth
    46.         transform = transforms.Compose([
    47.             transforms.Resize(240.0 / iheight),
    48.             transforms.CenterCrop(self.output_size),
    49.         ])
    50.         rgb_np = transform(rgb)
    51.         rgb_np = np.asfarray(rgb_np, dtype='float') / 255
    52.         depth_np = transform(depth_np)
    53.  
    54.         return rgb_np, depth_np

            dataloader.py

    1. import os
    2. import os.path
    3. import numpy as np
    4. import torch.utils.data as data
    5. import h5py
    6. import dataloaders.transforms as transforms
    7.  
    8. IMG_EXTENSIONS = ['.h5',]
    9.  
    10. def is_image_file(filename):
    11.     return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
    12.  
    13. def find_classes(dir):
    14.     classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
    15.     classes.sort()
    16.     class_to_idx = {classes[i]: i for i in range(len(classes))}
    17.     return classes, class_to_idx
    18.  
    19. def make_dataset(dir, class_to_idx):
    20.     images = []
    21.     dir = os.path.expanduser(dir)
    22.     for target in sorted(os.listdir(dir)):
    23.         d = os.path.join(dir, target)
    24.         if not os.path.isdir(d):
    25.             continue
    26.         for root, _, fnames in sorted(os.walk(d)):
    27.             for fname in sorted(fnames):
    28.                 if is_image_file(fname):
    29.                     path = os.path.join(root, fname)
    30.                     item = (path, class_to_idx[target])
    31.                     images.append(item)
    32.     return images
    33.  
    34. def h5_loader(path):
    35.     h5f = h5py.File(path, "r")
    36.     rgb = np.array(h5f['rgb'])
    37.     rgb = np.transpose(rgb, (1, 2, 0))
    38.     depth = np.array(h5f['depth'])
    39.     return rgb, depth
    40.  
    41. # def rgb2grayscale(rgb):
    42. #     return rgb[:,:,0] * 0.2989 + rgb[:,:,1] * 0.587 + rgb[:,:,2] * 0.114
    43.  
    44. to_tensor = transforms.ToTensor()
    45.  
    46. class MyDataloader(data.Dataset):
    47.     modality_names = ['rgb', 'rgbd', 'd'] # , 'g', 'gd'
    48.     color_jitter = transforms.ColorJitter(0.4, 0.4, 0.4)
    49.  
    50.     def __init__(self, root, type, sparsifier=None, modality='rgb', loader=h5_loader):
    51.         classes, class_to_idx = find_classes(root)
    52.         imgs = make_dataset(root, class_to_idx)
    53.         assert len(imgs)>0, "Found 0 images in subfolders of: " + root + "\n"
    54.         print("Found {} images in {} folder.".format(len(imgs), type))
    55.         self.root = root
    56.         self.imgs = imgs
    57.         self.classes = classes
    58.         self.class_to_idx = class_to_idx
    59.         if type == 'train':
    60.             self.transform = self.train_transform
    61.         elif type == 'val':
    62.             self.transform = self.val_transform
    63.         elif type == 'test':
    64.             self.transform = self.test_transform
    65.         else:
    66.             raise (RuntimeError("Invalid dataset type: " + type + "\n"
    67.                                 "Supported dataset types are: train, val"))
    68.         self.loader = loader
    69.         self.sparsifier = sparsifier
    70.  
    71.         assert (modality in self.modality_names), "Invalid modality type: " + modality + "\n" + \
    72.                                 "Supported dataset types are: " + ''.join(self.modality_names)
    73.         self.modality = modality
    74.  
    75.         self.mark = type
    76.  
    77.     def train_transform(self, rgb, depth):
    78.         raise (RuntimeError("train_transform() is not implemented. "))
    79.  
    80.     def val_transform(rgb, depth):
    81.         raise (RuntimeError("val_transform() is not implemented."))
    82.  
    83.     def test_transform(rgb, depth):
    84.         raise (RuntimeError("test_transform() is not implemented."))
    85.  
    86.     def create_sparse_depth(self, rgb, depth):
    87.         if self.sparsifier is None:
    88.             return depth
    89.         else:
    90.             mask_keep = self.sparsifier.dense_to_sparse(rgb, depth)
    91.             sparse_depth = np.zeros(depth.shape)
    92.             sparse_depth[mask_keep] = depth[mask_keep]
    93.             return sparse_depth
    94.  
    95.     def create_rgbd(self, rgb, depth):
    96.         sparse_depth = self.create_sparse_depth(rgb, depth)
    97.         rgbd = np.append(rgb, np.expand_dims(sparse_depth, axis=2), axis=2)
    98.         return rgbd
    99.  
    100.     def __getraw__(self, index):
    101.         """
    102.         Args:
    103.             index (int): Index
    105.         Returns:
    106.             tuple: (rgb, depth) the raw data.
    107.         """
    108.         path, target = self.imgs[index]
    109.         rgb, depth = self.loader(path)
    110.         _, name = os.path.split(path)
    111.         name = name.split('.')[0]
    112.         return rgb, depth, name
    113.  
    114.     def __getitem__(self, index):
    115.         rgb, depth, name = self.__getraw__(index)
    116.         if self.transform is not None:
    117.             rgb_np, depth_np = self.transform(rgb, depth)
    118.         else:
    119.             raise(RuntimeError("transform not defined"))
    120.  
    121.         # color normalization
    122.         # rgb_tensor = normalize_rgb(rgb_tensor)
    123.         # rgb_np = normalize_np(rgb_np)
    124.  
    125.         if self.modality == 'rgb':
    126.             input_np = rgb_np
    127.         elif self.modality == 'rgbd':
    128.             input_np = self.create_rgbd(rgb_np, depth_np)
    129.         elif self.modality == 'd':
    130.             input_np = self.create_sparse_depth(rgb_np, depth_np)
    131.  
    132.         input_tensor = to_tensor(input_np)
    133.         while input_tensor.dim() < 3:
    134.             input_tensor = input_tensor.unsqueeze(0)
    135.  
    136.         if self.mark == 'test':
    137.             depth_tensor = name
    138.         else:
    139.             depth_tensor = to_tensor(depth_np)
    140.             depth_tensor = depth_tensor.unsqueeze(0)
    141.  
    142.         return input_tensor, depth_tensor
    143.  
    144.     def __len__(self):
    145.         return len(self.imgs)

            2.参数部分

            util.py

    1. import os
    2.  
    3. import cv2
    4. import torch
    5. import shutil
    6. import numpy as np
    7. import matplotlib.pyplot as plt
    8. from PIL import Image
    9.  
    10. cmap = plt.cm.viridis
    11.  
    12. def parse_command():
    13.     model_names = ['resnet18', 'resnet50']
    14.     loss_names = ['l1', 'l2']
    15.     data_names = ['nyudepthv2', 'kitti']
    16.     from dataloaders.dense_to_sparse import UniformSampling, SimulatedStereo
    17.     sparsifier_names = [x.name for x in [UniformSampling, SimulatedStereo]]
    18.     from models import Decoder
    19.     decoder_names = Decoder.names
    20.     from dataloaders.dataloader import MyDataloader
    21.     modality_names = MyDataloader.modality_names
    22.  
    23.     import argparse
    24.     parser = argparse.ArgumentParser(description='Sparse-to-Dense')
    25.     parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18', choices=model_names,
    26.                         help='model architecture: ' + ' | '.join(model_names) + ' (default: resnet18)')
    27.     parser.add_argument('--data', metavar='DATA', default='nyudepthv2',
    28.                         choices=data_names,
    29.                         help='dataset: ' + ' | '.join(data_names) + ' (default: nyudepthv2)')
    30.     parser.add_argument('--modality', '-m', metavar='MODALITY', default='rgb', choices=modality_names,
    31.                         help='modality: ' + ' | '.join(modality_names) + ' (default: rgb)')
    32.     parser.add_argument('-s', '--num-samples', default=0, type=int, metavar='N',
    33.                         help='number of sparse depth samples (default: 0)')
    34.     parser.add_argument('--max-depth', default=-1.0, type=float, metavar='D',
    35.                         help='cut-off depth of sparsifier, negative values means infinity (default: inf [m])')
    36.     parser.add_argument('--sparsifier', metavar='SPARSIFIER', default=UniformSampling.name, choices=sparsifier_names,
    37.                         help='sparsifier: ' + ' | '.join(sparsifier_names) + ' (default: ' + UniformSampling.name + ')')
    38.     parser.add_argument('--decoder', '-d', metavar='DECODER', default='deconv2', choices=decoder_names,
    39.                         help='decoder: ' + ' | '.join(decoder_names) + ' (default: deconv2)')
    40.     parser.add_argument('-j', '--workers', default=0, type=int, metavar='N',
    41.                         help='number of data loading workers (default: 10)')
    42.     parser.add_argument('--epochs', default=15, type=int, metavar='N',
    43.                         help='number of total epochs to run (default: 15)')
    44.     parser.add_argument('-c', '--criterion', metavar='LOSS', default='l1', choices=loss_names,
    45.                         help='loss function: ' + ' | '.join(loss_names) + ' (default: l1)')
    46.     parser.add_argument('-b', '--batch-size', default=2, type=int, help='mini-batch size (default: 8)')
    47.     parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,
    48.                         metavar='LR', help='initial learning rate (default 0.01)')
    49.     parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
    50.                         help='momentum')
    51.     parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
    52.                         metavar='W', help='weight decay (default: 1e-4)')
    53.     parser.add_argument('--print-freq', '-p', default=10, type=int,
    54.                         metavar='N', help='print frequency (default: 10)')
    55.     parser.add_argument('--resume', default='', type=str, metavar='PATH',
    56.                         help='path to latest checkpoint (default: none)')
    57.     parser.add_argument('-e', '--evaluate', dest='evaluate', type=str, default='',
    58.                         help='evaluate model on validation set')
    59.     parser.add_argument('-t', '--test', dest='test', type=str, default='',
    60.                         help='test model on test set')
    61.     parser.add_argument('--no-pretrain', dest='pretrained', action='store_false',
    62.                         help='not to use ImageNet pre-trained weights')
    63.     parser.set_defaults(pretrained=True)
    64.     args = parser.parse_args()
    65.     if args.modality == 'rgb' and args.num_samples != 0:
    66.         print("number of samples is forced to be 0 when input modality is rgb")
    67.         args.num_samples = 0
    68.     if args.modality == 'rgb' and args.max_depth != 0.0:
    69.         print("max depth is forced to be 0.0 when input modality is rgb/rgbd")
    70.         args.max_depth = 0.0
    71.     return args
    72.  
    73. def save_checkpoint(state, is_best, epoch, output_directory):
    74.     checkpoint_filename = os.path.join(output_directory, 'checkpoint-' + str(epoch) + '.pth.tar')
    75.     torch.save(state, checkpoint_filename)
    76.     if is_best:
    77.         best_filename = os.path.join(output_directory, 'model_best.pth.tar')
    78.         shutil.copyfile(checkpoint_filename, best_filename)
    79.     if epoch > 0:
    80.         prev_checkpoint_filename = os.path.join(output_directory, 'checkpoint-' + str(epoch-1) + '.pth.tar')
    81.         if os.path.exists(prev_checkpoint_filename):
    82.             os.remove(prev_checkpoint_filename)
    83.  
    84. def adjust_learning_rate(optimizer, epoch, lr_init):
    85.     """Sets the learning rate to the initial LR decayed by 10 every 5 epochs"""
    86.     lr = lr_init * (0.1 ** (epoch // 5))
    87.     for param_group in optimizer.param_groups:
    88.         param_group['lr'] = lr
    89.  
    90. def get_output_directory(args):
    91.     output_directory = os.path.join('results',
    92.         '{}.sparsifier={}.samples={}.modality={}.arch={}.decoder={}.criterion={}.lr={}.bs={}.pretrained={}'.
    93.         format(args.data, args.sparsifier, args.num_samples, args.modality, \
    94.             args.arch, args.decoder, args.criterion, args.lr, args.batch_size, \
    95.             args.pretrained))
    96.     return output_directory
    97.  
    98.  
    99. def colored_depthmap(depth, d_min=None, d_max=None):
    100.     if d_min is None:
    101.         d_min = np.min(depth)
    102.     if d_max is None:
    103.         d_max = np.max(depth)
    104.     depth_relative = (depth - d_min) / (d_max - d_min)
    105.     return 255 * cmap(depth_relative)[:,:,:3] # H, W, C
    106.  
    107.  
    108. def merge_into_row(input, depth_target, depth_pred):
    109.     rgb = 255 * np.transpose(np.squeeze(input.cpu().numpy()), (1,2,0)) # H, W, C
    110.     depth_target_cpu = np.squeeze(depth_target.cpu().numpy())
    111.     depth_pred_cpu = np.squeeze(depth_pred.data.cpu().numpy())
    112.  
    113.     d_min = min(np.min(depth_target_cpu), np.min(depth_pred_cpu))
    114.     d_max = max(np.max(depth_target_cpu), np.max(depth_pred_cpu))
    115.     depth_target_col = colored_depthmap(depth_target_cpu, d_min, d_max)
    116.     depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)
    117.     img_merge = np.hstack([rgb, depth_target_col, depth_pred_col])
    118.     
    119.     return img_merge
    120.  
    121.  
    122. def merge_into_row_with_gt(input, depth_input, depth_target, depth_pred):
    123.     rgb = 255 * np.transpose(np.squeeze(input.cpu().numpy()), (1,2,0)) # H, W, C
    124.     depth_input_cpu = np.squeeze(depth_input.cpu().numpy())
    125.     depth_target_cpu = np.squeeze(depth_target.cpu().numpy())
    126.     depth_pred_cpu = np.squeeze(depth_pred.data.cpu().numpy())
    127.  
    128.     d_min = min(np.min(depth_input_cpu), np.min(depth_target_cpu), np.min(depth_pred_cpu))
    129.     d_max = max(np.max(depth_input_cpu), np.max(depth_target_cpu), np.max(depth_pred_cpu))
    130.     depth_input_col = colored_depthmap(depth_input_cpu, d_min, d_max)
    131.     depth_target_col = colored_depthmap(depth_target_cpu, d_min, d_max)
    132.     depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)
    133.  
    134.     img_merge = np.hstack([rgb, depth_input_col, depth_target_col, depth_pred_col])
    135.  
    136.     return img_merge
    137.  
    138.  
    139. def add_row(img_merge, row):
    140.     return np.vstack([img_merge, row])
    141.  
    142.  
    143. def save_image(img_merge, filename):
    144.     img_merge = Image.fromarray(img_merge.astype('uint8'))
    145.     img_merge.save(filename)
    146.  
    147.  
    148. def strentch_img(pred):
    149.     depth_pred_cpu = np.squeeze(pred.data.cpu().numpy())
    150.     d_min = np.min(depth_pred_cpu)
    151.     d_max = np.max(depth_pred_cpu)
    152.     depth_pred_cpu = cv2.resize(depth_pred_cpu, (1280, 720))
    153.     depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)
    154.     return depth_pred_col

            （3）主函数部分

            main.py

    1. import os
    2. import time
    3. import csv
    4. import numpy as np
    5.  
    6. import torch
    7. import torch.backends.cudnn as cudnn
    8. import torch.optim
    9. cudnn.benchmark = True
    10.  
    11. from models import ResNet
    12. from metrics import AverageMeter, Result
    13. from dataloaders.dense_to_sparse import UniformSampling, SimulatedStereo
    14. import criteria
    15. import utils
    16. from PIL import Image
    17.  
    18. torch.nn.Module.dump_patches = True
    19. args = utils.parse_command()
    20. print(args)
    21.  
    22. fieldnames = ['mse', 'rmse', 'absrel', 'lg10', 'mae',
    23.                 'delta1', 'delta2', 'delta3',
    24.                 'data_time', 'gpu_time']
    25. best_result = Result()
    26. best_result.set_to_worst()
    27.  
    28. def create_data_loaders(args):
    29.     # Data loading code
    30.     print("=> creating data loaders ...")
    31.     traindir = os.path.join('data', args.data, 'train')
    32.     valdir = os.path.join('data', args.data, 'val')
    33.     testdir = os.path.join('data', args.data, 'test')
    34.  
    35.     train_loader = None
    36.     val_loader = None
    37.     test_loader = None
    38.  
    39.  
    40.     # sparsifier is a class for generating random sparse depth input from the ground truth
    41.     sparsifier = None
    42.     max_depth = args.max_depth if args.max_depth >= 0.0 else np.inf
    43.     if args.sparsifier == UniformSampling.name:
    44.         sparsifier = UniformSampling(num_samples=args.num_samples, max_depth=max_depth)
    45.     elif args.sparsifier == SimulatedStereo.name:
    46.         sparsifier = SimulatedStereo(num_samples=args.num_samples, max_depth=max_depth)
    47.  
    48.     if args.data == 'nyudepthv2':
    49.         from dataloaders.nyu_dataloader import NYUDataset
    50.         if args.evaluate:
    51.             val_dataset = NYUDataset(valdir, type='val',
    52.                 modality=args.modality, sparsifier=sparsifier)
    53.             # set batch size to be 1 for validation
    54.             val_loader = torch.utils.data.DataLoader(val_dataset,
    55.                                                      batch_size=1, shuffle=False, num_workers=args.workers,
    56.                                                      pin_memory=True)
    57.         elif args.test:
    58.             test_dataset = NYUDataset(testdir, type='test',
    59.                                      modality=args.modality, sparsifier=sparsifier)
    60.  
    61.             test_loader = torch.utils.data.DataLoader(test_dataset,
    62.                                                       batch_size=1, shuffle=False, num_workers=args.workers,
    63.                                                       pin_memory=True)
    64.         else:
    65.             train_dataset = NYUDataset(traindir, type='train',
    66.                 modality=args.modality, sparsifier=sparsifier)
    67.  
    68.     elif args.data == 'kitti':
    69.         from dataloaders.kitti_dataloader import KITTIDataset
    70.         if not args.evaluate:
    71.             train_dataset = KITTIDataset(traindir, type='train',
    72.                 modality=args.modality, sparsifier=sparsifier)
    73.         val_dataset = KITTIDataset(valdir, type='val',
    74.             modality=args.modality, sparsifier=sparsifier)
    75.  
    76.         # set batch size to be 1 for validation
    77.         val_loader = torch.utils.data.DataLoader(val_dataset,
    78.                                                  batch_size=1, shuffle=False, num_workers=args.workers,
    79.                                                  pin_memory=True)
    80.  
    81.     else:
    82.         raise RuntimeError('Dataset not found.' +
    83.                            'The dataset must be either of nyudepthv2 or kitti.')
    84.  
    85.  
    86.     # put construction of train loader here, for those who are interested in testing only
    87.     if not args.evaluate and not args.test:
    88.         train_loader = torch.utils.data.DataLoader(
    89.             train_dataset, batch_size=args.batch_size, shuffle=True,
    90.             num_workers=args.workers, pin_memory=True, sampler=None,
    91.             worker_init_fn=lambda work_id:np.random.seed(work_id))
    92.             # worker_init_fn ensures different sampling patterns for each data loading thread
    93.  
    94.     print("=> data loaders created.")
    95.     return train_loader, val_loader, test_loader
    96.  
    97. test_save_path = './results/'
    98. def main():
    99.     global args, best_result, output_directory, train_csv, test_csv
    100.  
    101.     # evaluation mode
    102.     start_epoch = 0
    103.     if args.evaluate:
    104.         assert os.path.isfile(args.evaluate), \
    105.         "=> no best model found at '{}'".format(args.evaluate)
    106.         print("=> loading best model '{}'".format(args.evaluate))
    107.         checkpoint = torch.load(args.evaluate)
    108.         output_directory = os.path.dirname(args.evaluate)
    109.         args = checkpoint['args']
    110.         start_epoch = checkpoint['epoch'] + 1
    111.         best_result = checkpoint['best_result']
    112.         model = checkpoint['model']
    113.         print("=> loaded best model (epoch {})".format(checkpoint['epoch']))
    114.         args.test = ''
    115.         args.evaluate = True
    116.         _, val_loader, _ = create_data_loaders(args)
    117.         validate(val_loader, model, checkpoint['epoch'], write_to_file=False)
    118.         return
    119.  
    120.     elif args.test:
    121.         assert os.path.isfile(args.test), \
    122.         "=> no best model found at '{}'".format(args.test)
    123.         print("=> loading best model '{}'".format(args.test))
    124.         checkpoint = torch.load(args.test)
    125.         output_directory = os.path.dirname(args.test)
    126.         args = checkpoint['args']
    127.         start_epoch = checkpoint['epoch'] + 1
    128.         best_result = checkpoint['best_result']
    129.         model = checkpoint['model']
    130.         print("=> loaded best model (epoch {})".format(checkpoint['epoch']))
    131.         args.test = True
    132.         _, _, test_loader = create_data_loaders(args)
    133.         test(test_loader, model, test_save_path)
    134.         return
    135.  
    136.     # optionally resume from a checkpoint
    137.     elif args.resume:
    138.         chkpt_path = args.resume
    139.         assert os.path.isfile(chkpt_path), \
    140.             "=> no checkpoint found at '{}'".format(chkpt_path)
    141.         print("=> loading checkpoint '{}'".format(chkpt_path))
    142.         checkpoint = torch.load(chkpt_path)
    143.         args = checkpoint['args']
    144.         start_epoch = checkpoint['epoch'] + 1
    145.         best_result = checkpoint['best_result']
    146.         model = checkpoint['model']
    147.         optimizer = checkpoint['optimizer']
    148.         output_directory = os.path.dirname(os.path.abspath(chkpt_path))
    149.         print("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))
    150.         train_loader, val_loader, test_loader = create_data_loaders(args)
    151.         args.resume = True
    152.  
    153.     # create new model
    154.     else:
    155.         train_loader, val_loader, test_loader = create_data_loaders(args)
    156.         print("=> creating Model ({}-{}) ...".format(args.arch, args.decoder))
    157.         in_channels = len(args.modality)
    158.         if args.arch == 'resnet50':
    159.             model = ResNet(layers=50, decoder=args.decoder, output_size=train_loader.dataset.output_size,
    160.                 in_channels=in_channels, pretrained=args.pretrained)
    161.         elif args.arch == 'resnet18':
    162.             model = ResNet(layers=18, decoder=args.decoder, output_size=train_loader.dataset.output_size,
    163.                 in_channels=in_channels, pretrained=args.pretrained)
    164.         print("=> model created.")
    165.         optimizer = torch.optim.SGD(model.parameters(), args.lr, \
    166.             momentum=args.momentum, weight_decay=args.weight_decay)
    167.  
    168.         # model = torch.nn.DataParallel(model).cuda() # for multi-gpu training
    169.         model = model.cuda()
    170.  
    171.     # define loss function (criterion) and optimizer
    172.     if args.criterion == 'l2':
    173.         criterion = criteria.MaskedMSELoss().cuda()
    174.     elif args.criterion == 'l1':
    175.         criterion = criteria.MaskedL1Loss().cuda()
    176.  
    177.     # create results folder, if not already exists
    178.     output_directory = utils.get_output_directory(args)
    179.     if not os.path.exists(output_directory):
    180.         os.makedirs(output_directory)
    181.     train_csv = os.path.join(output_directory, 'train.csv')
    182.     test_csv = os.path.join(output_directory, 'test.csv')
    183.     best_txt = os.path.join(output_directory, 'best.txt')
    184.  
    185.     # create new csv files with only header
    186.     if not args.resume:
    187.         with open(train_csv, 'w') as csvfile:
    188.             writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    189.             writer.writeheader()
    190.         with open(test_csv, 'w') as csvfile:
    191.             writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    192.             writer.writeheader()
    193.  
    194.     for epoch in range(start_epoch, args.epochs):
    195.         utils.adjust_learning_rate(optimizer, epoch, args.lr)
    196.         train(train_loader, model, criterion, optimizer, epoch) # train for one epoch
    197.         result, img_merge = validate(val_loader, model, epoch) # evaluate on validation set
    198.  
    199.         # remember best rmse and save checkpoint
    200.         is_best = result.rmse < best_result.rmse
    201.         if is_best:
    202.             best_result = result
    203.             with open(best_txt, 'w') as txtfile:
    204.                 txtfile.write("epoch={}\nmse={:.3f}\nrmse={:.3f}\nabsrel={:.3f}\nlg10={:.3f}\nmae={:.3f}\ndelta1={:.3f}\nt_gpu={:.4f}\n".
    205.                     format(epoch, result.mse, result.rmse, result.absrel, result.lg10, result.mae, result.delta1, result.gpu_time))
    206.             if img_merge is not None:
    207.                 img_filename = output_directory + '/comparison_best.png'
    208.                 utils.save_image(img_merge, img_filename)
    209.  
    210.         utils.save_checkpoint({
    211.             'args': args,
    212.             'epoch': epoch,
    213.             'arch': args.arch,
    214.             'model': model,
    215.             'best_result': best_result,
    216.             'optimizer' : optimizer,
    217.         }, is_best, epoch, output_directory)
    218.  
    219.  
    220. def train(train_loader, model, criterion, optimizer, epoch):
    221.     average_meter = AverageMeter()
    222.     model.train() # switch to train mode
    223.     end = time.time()
    224.     for i, (input, target) in enumerate(train_loader):
    225.  
    226.         input, target = input.cuda(), target.cuda()
    227.         torch.cuda.synchronize()
    228.         data_time = time.time() - end
    229.  
    230.         # compute pred
    231.         end = time.time()
    232.         pred = model(input)
    233.         loss = criterion(pred, target)
    234.         optimizer.zero_grad()
    235.         loss.backward() # compute gradient and do SGD step
    236.         optimizer.step()
    237.         torch.cuda.synchronize()
    238.         gpu_time = time.time() - end
    239.  
    240.         # measure accuracy and record loss
    241.         result = Result()
    242.         result.evaluate(pred.data, target.data)
    243.         average_meter.update(result, gpu_time, data_time, input.size(0))
    244.         end = time.time()
    245.  
    246.         if (i + 1) % args.print_freq == 0:
    247.             print('=> output: {}'.format(output_directory))
    248.             print('Train Epoch: {0} [{1}/{2}]\t'
    249.                   't_Data={data_time:.3f}({average.data_time:.3f}) '
    250.                   't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
    251.                   'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
    252.                   'MAE={result.mae:.2f}({average.mae:.2f}) '
    253.                   'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
    254.                   'REL={result.absrel:.3f}({average.absrel:.3f}) '
    255.                   'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
    256.                   epoch, i+1, len(train_loader), data_time=data_time,
    257.                   gpu_time=gpu_time, result=result, average=average_meter.average()))
    258.  
    259.     avg = average_meter.average()
    260.     with open(train_csv, 'a') as csvfile:
    261.         writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    262.         writer.writerow({'mse': avg.mse, 'rmse': avg.rmse, 'absrel': avg.absrel, 'lg10': avg.lg10,
    263.             'mae': avg.mae, 'delta1': avg.delta1, 'delta2': avg.delta2, 'delta3': avg.delta3,
    264.             'gpu_time': avg.gpu_time, 'data_time': avg.data_time})
    265.  
    266.  
    267. def validate(val_loader, model, epoch, write_to_file=True):
    268.     average_meter = AverageMeter()
    269.     model.eval() # switch to evaluate mode
    270.     end = time.time()
    271.     for i, (input, target) in enumerate(val_loader):
    272.         input, target = input.cuda(), target.cuda()
    273.         torch.cuda.synchronize()
    274.         data_time = time.time() - end
    275.  
    276.         # compute output
    277.         end = time.time()
    278.         with torch.no_grad():
    279.             pred = model(input)
    280.         torch.cuda.synchronize()
    281.         gpu_time = time.time() - end
    282.  
    283.         # measure accuracy and record loss
    284.         result = Result()
    285.         result.evaluate(pred.data, target.data)
    286.         average_meter.update(result, gpu_time, data_time, input.size(0))
    287.         end = time.time()
    288.  
    289.         # save 8 images for visualization
    290.         skip = 50
    291.         if args.modality == 'd':
    292.             img_merge = None
    293.         else:
    294.             if args.modality == 'rgb':
    295.                 rgb = input
    296.             elif args.modality == 'rgbd':
    297.                 rgb = input[:,:3,:,:]
    298.                 depth = input[:,3:,:,:]
    299.  
    300.             if i == 0:
    301.                 if args.modality == 'rgbd':
    302.                     img_merge = utils.merge_into_row_with_gt(rgb, depth, target, pred)
    303.                 else:
    304.                     img_merge = utils.merge_into_row(rgb, target, pred)
    305.             elif (i < 8*skip) and (i % skip == 0):
    306.                 if args.modality == 'rgbd':
    307.                     row = utils.merge_into_row_with_gt(rgb, depth, target, pred)
    308.                 else:
    309.                     row = utils.merge_into_row(rgb, target, pred)
    310.                 img_merge = utils.add_row(img_merge, row)
    311.             elif i == 8*skip:
    312.                 filename = output_directory + '/comparison_' + str(epoch) + '.png'
    313.                 utils.save_image(img_merge, filename)
    314.  
    315.         if (i+1) % args.print_freq == 0:
    316.             print('Test: [{0}/{1}]\t'
    317.                   't_GPU={gpu_time:.3f}({average.gpu_time:.3f})\n\t'
    318.                   'RMSE={result.rmse:.2f}({average.rmse:.2f}) '
    319.                   'MAE={result.mae:.2f}({average.mae:.2f}) '
    320.                   'Delta1={result.delta1:.3f}({average.delta1:.3f}) '
    321.                   'REL={result.absrel:.3f}({average.absrel:.3f}) '
    322.                   'Lg10={result.lg10:.3f}({average.lg10:.3f}) '.format(
    323.                    i+1, len(val_loader), gpu_time=gpu_time, result=result, average=average_meter.average()))
    324.  
    325.     avg = average_meter.average()
    326.  
    327.     print('\n*\n'
    328.         'RMSE={average.rmse:.3f}\n'
    329.         'MAE={average.mae:.3f}\n'
    330.         'Delta1={average.delta1:.3f}\n'
    331.         'REL={average.absrel:.3f}\n'
    332.         'Lg10={average.lg10:.3f}\n'
    333.         't_GPU={time:.3f}\n'.format(
    334.         average=avg, time=avg.gpu_time))
    335.  
    336.     if write_to_file:
    337.         with open(test_csv, 'a') as csvfile:
    338.             writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    339.             writer.writerow({'mse': avg.mse, 'rmse': avg.rmse, 'absrel': avg.absrel, 'lg10': avg.lg10,
    340.                 'mae': avg.mae, 'delta1': avg.delta1, 'delta2': avg.delta2, 'delta3': avg.delta3,
    341.                 'data_time': avg.data_time, 'gpu_time': avg.gpu_time})
    342.     return avg, img_merge
    343.  
    344.  
    345. def test(test_loader, model, save_path):
    346.     average_meter = AverageMeter()
    347.     model.eval() # switch to evaluate mode
    348.     for i, (input, target) in enumerate(test_loader):
    349.         input, name = input.cuda(), target
    350.         torch.cuda.synchronize()
    351.  
    352.         # compute output
    353.         end = time.time()
    354.         with torch.no_grad():
    355.             pred = model(input)
    356.  
    357.         torch.cuda.synchronize()
    358.         pred1 = utils.strentch_img(pred)
    359.         save_to_file = os.path.join(save_path, name[0] + '.png')
    360.         utils.save_image(pred1, save_to_file)
    361.  
    362.         save_to_tif = os.path.join(save_path, name[0] + '_ori.tiff')
    363.         depth_pred_cpu = np.squeeze(pred.data.cpu().numpy())
    364.         img = Image.fromarray(depth_pred_cpu)
    365.         img = img.resize((1280, 720))
    366.         img.save(save_to_tif)
    367.  
    368.  
    369. if __name__ == '__main__':
    370.     main()

            （4）测试

            改好上面后，创建test文件夹，放入数据

     接着命令行输入下面的命令

    python main.py --test model_best.pth

    白色的图是结果，彩色图是白色图可视化后的结果 ，存放位置在mian.py的第90行改（test_save_path = './results/'）

  • 相关阅读:
    WPF实现轮播图（图片、视屏）
    OpenMP编程-九点差分法求解泊松方程
    基于Python代码的相关性热力图，VIF共线性诊断图及残差四图的使用及解释
    QQ浏览器怎么才能设置默认搜索引擎为百度
    一百八十五、大数据离线数仓完整流程——步骤四、在Hive的DWD层建动态分区表并动态加载数据
    阿里云周宇：神龙计算平台智能运维体系建设
    【PyTorch深度学习项目实战100例】—— 基于Pytorch的中文问题相似度实战 | 第72例
    质量属性案例-架构真题（二十一）
    【JVM】触发 Full GC 的条件
    生成m3u8视频：批量剪辑与分割的完美结合
  • 原文地址：https://blog.csdn.net/qq_20373723/article/details/126815756