Successfully configured the dataloader and trained for one epoch. Results are not so good, but it's something. Still the Fastaiv1 looked better qualitatively

dvc.lock

@@ -3,21 +3,21 @@ process_data:
     src/data/processed
   deps:
   - path: src/code/make_dataset.py
-    md5: b15a09b30657f303d62215ceb73c6f4f
+    md5: 726bf2bed948f73c5c342a96d017539e
   - path: src/data/raw/nyu_depth_v2_labeled.mat
     md5: 520609c519fba3ba5ac58c8fefcc3530
   - path: src/data/raw/splits.mat
     md5: 08e3c3aea27130ac7c01ffd739a4535f
   outs:
   - path: src/data/processed/
-    md5: 78d932b016048d527f9e73e1781127cd.dir
+    md5: 77adb8603dbf31f3b272e0f51b6c2c29.dir
 train:
   cmd: python3 src/code/training.py src/data/processed
   deps:
   - path: src/code/training.py
-    md5: d36fea3f2f8524ea040fe7330ce8ba40
+    md5: 1d5f2b07b208bf062526e5ebfddca043
   - path: src/data/processed/
-    md5: 78d932b016048d527f9e73e1781127cd.dir
+    md5: 77adb8603dbf31f3b272e0f51b6c2c29.dir
   outs:
   - path: src/models/
-    md5: 48f953dd81c67ea0f6958fe4e4226e01.dir
+    md5: e6f3667c5e3ff28faaf9172adab28107.dir

src/code/make_dataset.py

@@ -44,7 +44,9 @@ import cv2
 def convert_image(i, scene, depth, image, folder):
     img_depth = depth * 1000.0
     img_depth_uint16 = img_depth.astype(np.uint16)
-    cv2.imwrite("%s/%05d_depth.png" % (folder, i), img_depth_uint16)
+    normalized_depth = np.zeros(img_depth_uint16.shape)
+    normalized_depth = cv2.normalize(img_depth_uint16,  normalized_depth, 0, 255, cv2.NORM_MINMAX)
+    cv2.imwrite("%s/%05d_depth.png" % (folder, i), normalized_depth)
 
     image = image[:, :, ::-1]
     image_black_boundary = np.zeros((480, 640, 3), dtype=np.uint8)

src/code/training.py

@@ -3,6 +3,20 @@ import sys
 from fastai2.vision.all import *
 from torchvision.utils import save_image
 
+class ImageImageDataLoaders(DataLoaders):
+    "Basic wrapper around several `DataLoader`s with factory methods for Image to Image problems"
+    @classmethod
+    @delegates(DataLoaders.from_dblock)
+    def from_label_func(cls, path, fnames, label_func, valid_pct=0.2, seed=None, item_tfms=None, batch_tfms=None, **kwargs):
+        "Create from list of `fnames` in `path`s with `label_func`."
+        dblock = DataBlock(blocks=(ImageBlock(cls=PILImage), ImageBlock(cls=PILImageBW)),
+                           splitter=RandomSplitter(valid_pct, seed=seed),
+                           get_y=label_func,
+                           item_tfms=item_tfms,
+                           batch_tfms=batch_tfms)
+        res = cls.from_dblock(dblock, fnames, path=path, **kwargs)
+        return res
+
 
 def get_y_fn(x):
     y = str(x.absolute()).replace('.jpg', '_depth.png')
@@ -13,7 +27,7 @@ def get_y_fn(x):
 
 def create_data(data_path):
     fnames = get_files(data_path/'train', extensions='.jpg')
-    data = SegmentationDataLoaders.from_label_func(data_path/'train',  bs=4, num_workers=0, fnames=fnames, label_func=get_y_fn)
+    data = ImageImageDataLoaders.from_label_func(data_path/'train', seed=42, bs=4, num_workers=0, fnames=fnames, label_func=get_y_fn)
     return data
 
 
@@ -23,7 +37,7 @@ if __name__ == "__main__":
         sys.exit(0)
 
     data = create_data(Path(sys.argv[1]))
-    learner = unet_learner(data, resnet34, metrics=rmse, wd=1e-2, n_out=1, loss_func=MSELossFlat(), path='src/')
+    learner = unet_learner(data, resnet34, metrics=rmse, wd=1e-2, n_out=3, loss_func=MSELossFlat(), path='src/')
     learner.fine_tune(1)
 
     learner.save('model')