modify main.py and model.py to load KITTI dataset

Author: AkariAsai

main.py

@@ -4,6 +4,7 @@
 import sys
 import csv
 import numpy as np
+import errno
 
 import torch
 import torch.nn as nn
@@ -13,6 +14,7 @@
 import torch.utils.data
 
 from nyu_dataloader import NYUDataset
+from kitti_dataloader import KITTIDataset
 from models import Decoder, ResNet
 from metrics import AverageMeter, Result
 from dense_to_sparse import UniformSampling, SimulatedStereo
@@ -21,7 +23,7 @@
 
 model_names = ['resnet18', 'resnet50']
 loss_names = ['l1', 'l2']
-data_names = ['nyudepthv2']
+data_names = ['nyudepthv2', 'kitti']
 sparsifier_names = [x.name for x in [UniformSampling, SimulatedStereo]]
 decoder_names = Decoder.names
 modality_names = NYUDataset.modality_names
@@ -134,15 +136,27 @@ def main():
     traindir = os.path.join('data', args.data, 'train')
     valdir = os.path.join('data', args.data, 'val')
 
-    train_dataset = NYUDataset(traindir, type='train',
-        modality=args.modality, sparsifier=sparsifier)
+    if args.data == 'nyudepthv2':
+        train_dataset = NYUDataset(traindir, type='train',
+            modality=args.modality, sparsifier=sparsifier)
+        val_dataset = NYUDataset(valdir, type='val',
+            modality=args.modality, sparsifier=sparsifier)
+
+    elif args.data == 'kitti':
+        train_dataset = KITTIDataset(traindir, type='train',
+            modality=args.modality, sparsifier=sparsifier)
+        val_dataset = KITTIDataset(valdir, type='val',
+            modality=args.modality, sparsifier=sparsifier)
+
+    else:
+        raise RuntimeError('Dataset not found.' +
+                           'The dataset must be either of nyudepthv2 or kitti.')
+
     train_loader = torch.utils.data.DataLoader(
         train_dataset, batch_size=args.batch_size, shuffle=True,
         num_workers=args.workers, pin_memory=True, sampler=None)
 
     # set batch size to be 1 for validation
-    val_dataset = NYUDataset(valdir, type='val',
-        modality=args.modality, sparsifier=sparsifier)
     val_loader = torch.utils.data.DataLoader(val_dataset,
         batch_size=1, shuffle=False, num_workers=args.workers, pin_memory=True)
 

models.py

@@ -9,31 +9,31 @@
 oheight, owidth = 228, 304
 
 class Unpool(nn.Module):
-    # Unpool: 2*2 unpooling with zero padding 
+    # Unpool: 2*2 unpooling with zero padding
     def __init__(self, num_channels, stride=2):
         super(Unpool, self).__init__()
 
         self.num_channels = num_channels
         self.stride = stride
 
         # create kernel [1, 0; 0, 0]
-        self.weights = torch.autograd.Variable(torch.zeros(num_channels, 1, stride, stride).cuda()) # currently not compatible with running on CPU 
+        self.weights = torch.autograd.Variable(torch.zeros(num_channels, 1, stride, stride).cuda()) # currently not compatible with running on CPU
         self.weights[:,:,0,0] = 1
 
     def forward(self, x):
         return F.conv_transpose2d(x, self.weights, stride=self.stride, groups=self.num_channels)
 
 def weights_init(m):
     # Initialize filters with Gaussian random weights
-    if isinstance(m, nn.Conv2d): 
+    if isinstance(m, nn.Conv2d):
         n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
         m.weight.data.normal_(0, math.sqrt(2. / n))
-        if m.bias is not None: 
+        if m.bias is not None:
             m.bias.data.zero_()
     elif isinstance(m, nn.ConvTranspose2d):
         n = m.kernel_size[0] * m.kernel_size[1] * m.in_channels
         m.weight.data.normal_(0, math.sqrt(2. / n))
-        if m.bias is not None: 
+        if m.bias is not None:
             m.bias.data.zero_()
     elif isinstance(m, nn.BatchNorm2d):
         m.weight.data.fill_(1)
@@ -63,13 +63,13 @@ class DeConv(Decoder):
     def __init__(self, in_channels, kernel_size):
         assert kernel_size>=2, "kernel_size out of range: {}".format(kernel_size)
         super(DeConv, self).__init__()
-        
+
         def convt(in_channels):
             stride = 2
             padding = (kernel_size - 1) // 2
             output_padding = kernel_size % 2
             assert -2 - 2*padding + kernel_size + output_padding == 0, "deconv parameters incorrect"
-            
+
             module_name = "deconv{}".format(kernel_size)
             return nn.Sequential(collections.OrderedDict([
                   (module_name, nn.ConvTranspose2d(in_channels,in_channels//2,kernel_size,
@@ -107,7 +107,7 @@ class UpProj(Decoder):
 
     class UpProjModule(nn.Module):
         # UpProj module has two branches, with a Unpool at the start and a ReLu at the end
-        #   upper branch: 5*5 conv -> batchnorm -> ReLU -> 3*3 conv -> batchnorm 
+        #   upper branch: 5*5 conv -> batchnorm -> ReLU -> 3*3 conv -> batchnorm
         #   bottom branch: 5*5 conv -> batchnorm
 
         def __init__(self, in_channels):
@@ -145,7 +145,7 @@ def __init__(self, in_channels):
 def choose_decoder(decoder, in_channels):
     # iheight, iwidth = 10, 8
     if decoder[:6] == 'deconv':
-        assert len(decoder)==7 
+        assert len(decoder)==7
         kernel_size = int(decoder[6])
         return DeConv(in_channels, kernel_size)
     elif decoder == "upproj":
@@ -161,10 +161,10 @@ def __init__(self, layers, decoder, in_channels=3, out_channels=1, pretrained=Tr
 
         if layers not in [18, 34, 50, 101, 152]:
             raise RuntimeError('Only 18, 34, 50, 101, and 152 layer model are defined for ResNet. Got {}'.format(layers))
-        
+
         super(ResNet, self).__init__()
         pretrained_model = torchvision.models.__dict__['resnet{}'.format(layers)](pretrained=pretrained)
-        
+
         if in_channels == 3:
             self.conv1 = pretrained_model._modules['conv1']
             self.bn1 = pretrained_model._modules['bn1']
@@ -173,7 +173,7 @@ def __init__(self, layers, decoder, in_channels=3, out_channels=1, pretrained=Tr
             self.bn1 = nn.BatchNorm2d(64)
             weights_init(self.conv1)
             weights_init(self.bn1)
-        
+
         self.relu = pretrained_model._modules['relu']
         self.maxpool = pretrained_model._modules['maxpool']
         self.layer1 = pretrained_model._modules['layer1']
@@ -187,6 +187,8 @@ def __init__(self, layers, decoder, in_channels=3, out_channels=1, pretrained=Tr
         # define number of intermediate channels
         if layers <= 34:
             num_channels = 512
+            # Need to modify owidth for ResNet18 model.
+            owidth = 912
         elif layers >= 50:
             num_channels = 2048