Adds requirements file

demo.py

@@ -1,11 +1,12 @@
-import argparse
+from fire import Fire
+import numpy as np
+from PIL import Image
+from scipy.misc import imresize
 import torch
 import torch.nn.parallel
 
 from models import modules, net, resnet, densenet, senet
-import numpy as np
 import loaddata_demo as loaddata
-import pdb
 
 import matplotlib.image
 import matplotlib.pyplot as plt
@@ -29,23 +30,34 @@ def define_model(is_resnet, is_densenet, is_senet):
     return model
 
 
-def main():
+def main(image_path):
     model = define_model(is_resnet=False, is_densenet=False, is_senet=True)
-    model = torch.nn.DataParallel(model).cuda()
-    model.load_state_dict(torch.load('./pretrained_model/model_senet'))
+    model = torch.nn.DataParallel(model)
+    if torch.cuda.is_available():
+        model = model.cuda()
+    model.load_state_dict(torch.load(
+        f='./pretrained_model/model_senet',
+        map_location=None if torch.cuda.is_available() else 'cpu'))
     model.eval()
 
-    nyu2_loader = loaddata.readNyu2('data/demo/img_nyu2.png')
+    nyu2_loader = loaddata.readNyu2(image_path)
 
     test(nyu2_loader, model)
 
 
 def test(nyu2_loader, model):
     for i, image in enumerate(nyu2_loader):     
-        image = torch.autograd.Variable(image, volatile=True).cuda()
+        image = torch.autograd.Variable(image, volatile=True)
+        if torch.cuda.is_available():
+            image = image.cuda()
         out = model(image)
-
-        matplotlib.image.imsave('data/demo/out.png', out.view(out.size(2),out.size(3)).data.cpu().numpy())
+
+        out = out.view(out.size(2), out.size(3)).data.cpu().numpy()
+        input_shape = image.data.cpu().numpy().shape[2:4]
+        out = imresize(arr=out, size=input_shape)
+        Image.fromarray(out.astype(np.uint8)).save('data/demo/out.png')
+
+        # matplotlib.image.imsave('data/demo/out.png', out)
 
 if __name__ == '__main__':
-    main()
+    Fire(main)

loaddata.py

@@ -1,9 +1,7 @@
 import pandas as pd
-import numpy as np
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms, utils
 from PIL import Image
-import random
 from nyu_transform import *
 
 

loaddata_demo.py

@@ -1,9 +1,6 @@
-import pandas as pd
-import numpy as np
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms, utils
 from PIL import Image
-import random
 from demo_transform import *
 
 

models/modules.py

@@ -1,14 +1,7 @@
-from collections import OrderedDict
-import math
 import torch
 import torch.nn.functional as F
 import torch.nn as nn
-from torch.utils import model_zoo
-import copy
-import numpy as np
-import senet
-import resnet
-import densenet
+
 
 class _UpProjection(nn.Sequential):
 

models/net.py

@@ -1,17 +1,7 @@
-from collections import OrderedDict
-import math
 import torch
-import torch.nn.functional as F
 import torch.nn as nn
-from torch.utils import model_zoo
-import copy
-import numpy as np
-import modules
-from torchvision import utils
+from models import modules
 
-import senet
-import resnet
-import densenet
 
 class model(nn.Module):
     def __init__(self, Encoder, num_features, block_channel):

nyu_transform.py

@@ -10,8 +10,6 @@
 import random
 import scipy.ndimage as ndimage
 
-import pdb
-
 
 def _is_pil_image(img):
     if accimage is not None:

requirements.txt

@@ -0,0 +1,4 @@
+numpy
+matplotlib
+torch
+torchvision

test.py

@@ -1,6 +1,4 @@
-import argparse
 import torch
-import torch.nn as nn
 import torch.nn.parallel
 
 from models import modules, net, resnet, densenet, senet

train.py

@@ -7,8 +7,6 @@
 import torch.backends.cudnn as cudnn
 import torch.optim
 import loaddata
-import util
-import numpy as np
 import sobel
 from models import modules, net, resnet, densenet, senet