Bug: dead link to normals zip

.gitignore

@@ -3,4 +3,5 @@ colored_*
 test_labels_13
 train_labels_13
 *.zip
-*.mat
+*.mat
+**/__pycache__

download_and_extract.sh

@@ -1,4 +1,4 @@
 wget http://horatio.cs.nyu.edu/mit/silberman/nyu_depth_v2/nyu_depth_v2_labeled.mat
-wget https://inf.ethz.ch/personal/ladickyl/nyu_normals_gt.zip
+wget https://dl.fbaipublicfiles.com/fair_self_supervision_benchmark/nyuv2_surfacenormal_metadata.zip
 
-python extract_nyuv2.py --mat nyu_depth_v2_labeled.mat --normal_zip nyu_normals_gt.zip  --data_root NYUv2 --save_colored
+python extract_nyuv2.py --mat nyu_depth_v2_labeled.mat --normal_zip nyuv2_surfacenormal_metadata.zip  --data_root NYUv2 --save_colored

extract_nyuv2.py

@@ -1,15 +1,21 @@
+import argparse
+import gzip
 import os
-import sys
+import pickle
+import shutil
+import zipfile
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
 import h5py
-import argparse
+import matplotlib.pyplot as plt
 import numpy as np
-from skimage import io
+
 from scipy.io import loadmat
+from skimage import io
 from tqdm import tqdm
-import shutil
-import matplotlib
-import matplotlib.pyplot as plt
-import zipfile
+
 
 def colormap(N=256, normalized=False):
     def bitget(byteval, idx):
@@ -110,6 +116,39 @@ def extract_depths(depths, splits, DEPTH_DIR, save_colored=False):
                 colored = plt.cm.jet(norm(depth))
                 plt.imsave('colored_depth/%05d.png' % (idx), colored)
 
+
+def extract_normals(normal_zip: str, data_root: str, splits: dict) -> None:
+    # prepare directories
+    normal_test = Path(data_root) / "normal/test"
+    normal_test.mkdir(exist_ok=True, parents=True)
+    normal_train = Path(data_root) / "normal/train"
+    normal_train.mkdir(exist_ok=True, parents=True)
+
+    # extract raw normals in the form of array
+    with TemporaryDirectory() as tmpdir:
+        with zipfile.ZipFile(normal_zip, 'r') as normal_zip:
+            normal_zip.extractall(path=tmpdir)
+        with gzip.open(
+            f"{tmpdir}/surfacenormal_metadata/all_normals.pklz", "rb"
+        ) as f:
+            data = pickle.load(f)
+
+    # obtain indices of images in each of splits
+    reference_test = splits["testNdxs"].reshape(-1)
+    reference_train = splits["trainNdxs"].reshape(-1)
+
+    # save raw data in proper directory
+    for idx, f_name in enumerate(data["all_filenames"]):
+        if int(f_name) in reference_test:
+            out_path = normal_test / f"{f_name[1:]}.npy"
+        elif int(f_name) in reference_train:
+            out_path = normal_train / f"{f_name[1:]}.npy"
+        else:
+            raise ValueError(f"{f_name} not found in train and test splits!")
+        normal_map = data["all_normals"][idx, :]
+        np.save(out_path, normal_map)
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='RYU DATA Extraction')
     parser.add_argument('--mat', type=str, required=True,
@@ -119,7 +158,7 @@ def extract_depths(depths, splits, DEPTH_DIR, save_colored=False):
     parser.add_argument('--save_colored', action='store_true', default=False,
                         help="save colored labels and depth maps for visualization")
     parser.add_argument('--normal_zip', type=str, default=None,
-                        help='path to nyu_normals_gt.zip. https: // inf.ethz.ch/personal/ladickyl/nyu_normals_gt.zip')
+                        help='path to nyuv2_surfacenormal_metadata.zip. https://dl.fbaipublicfiles.com/fair_self_supervision_benchmark/nyuv2_surfacenormal_metadata.zip')
 
     args = parser.parse_args()
 
@@ -149,11 +188,7 @@ def extract_depths(depths, splits, DEPTH_DIR, save_colored=False):
         extract_images(np.array(images), splits, IMAGE_DIR)
 
         if args.normal_zip is not None and os.path.exists(args.normal_zip):
-            NORMAL_DIR = os.path.join(DATA_ROOT, 'normal')
-            os.makedirs(NORMAL_DIR, exist_ok=True)
-            with zipfile.ZipFile(args.normal_zip, 'r') as normal_zip:
-                normal_zip.extractall(path=NORMAL_DIR)
+            extract_normals(args.normal_zip, DATA_ROOT, splits)
 
         if not os.path.exists(os.path.join( DATA_ROOT, 'splits.mat' )):
             shutil.copy2( 'splits.mat', os.path.join( DATA_ROOT, 'splits.mat' ))
-

nyuv2.py

@@ -1,14 +1,13 @@
 #coding:utf-8
 import os
+
+import numpy as np
 import torch
-import torch.utils.data as data
-from PIL import Image
+
+from PIL import Image, UnidentifiedImageError
 from scipy.io import loadmat
-import numpy as np
-import glob
 from torchvision import transforms
 from torchvision.datasets import VisionDataset
-import random
 
 def colormap(N=256, normalized=False):
     def bitget(byteval, idx):
@@ -79,12 +78,15 @@ def __init__(self,
 
         if self.target_type=='normal':
             normal_dir = os.path.join(self.root, 'normal', self.split)
-            self.normals = [os.path.join(normal_dir, name) for name in img_names]
+            self.normals = [os.path.join(normal_dir, f"{name[:-3]}npy") for name in img_names]
             self.targets = self.normals
 
     def __getitem__(self, idx):
         image = Image.open(self.images[idx])
-        target = Image.open(self.targets[idx])
+        try:
+            target = Image.open(self.targets[idx])
+        except UnidentifiedImageError:
+            target = np.load(self.targets[idx])
         if self.transforms is not None:
             image, target = self.transforms( image, target )
         return image, target
@@ -128,37 +130,36 @@ def __len__(self):
                             ]),  
                         )
 
-    nyu_normal = NYUv2( root='NYUv2', split='train', target_type='normal', 
-                            transform=transforms.Compose([
-                                transforms.Resize(512),
-                                transforms.ToTensor()
-                            ]),
-                            target_transform=transforms.Compose([
-                                transforms.ToTensor(),
-                                transforms.Lambda(lambda normal: normal * 2 - 1)
-                            ]),  
-                        )
-        
-    os.makedirs('test', exist_ok=True)
+    nyu_normal = NYUv2(
+        root="NYUv2",
+        split="train",
+        target_type="normal", 
+        transform=transforms.Compose(
+            [transforms.Resize(512), transforms.ToTensor()]
+        ),
+        target_transform=transforms.ToTensor(),
+    )
+
+    target_dir = "test" 
+    os.makedirs(target_dir, exist_ok=True)
     # Semantic
-    img_id = 0
+    img_id = 20
 
     img, lbl13 = nyu_semantic13[img_id]
-    Image.fromarray((img*255).numpy().transpose( 1,2,0 ).astype('uint8')).save('test/image.png')
-    Image.fromarray( nyu_semantic13.cmap[ (lbl13.numpy().astype('uint8')+1) ] ).save('test/semantic13.png')
+    Image.fromarray((img*255).numpy().transpose( 1,2,0 ).astype('uint8')).save(f"{target_dir}/image.png")
+    Image.fromarray( nyu_semantic13.cmap[ (lbl13.numpy().astype('uint8')+1) ] ).save(f"{target_dir}/semantic13.png")
 
     img, lbl40 = nyu_semantic40[img_id]
-    Image.fromarray( nyu_semantic40.cmap[ (lbl40.numpy().astype('uint8')+1) ] ).save('test/semantic40.png')
+    Image.fromarray( nyu_semantic40.cmap[ (lbl40.numpy().astype('uint8')+1) ] ).save(f"{target_dir}/semantic40.png")
 
     # Depth
     img, depth = nyu_depth[img_id]
     norm = plt.Normalize()
     depth = plt.cm.jet(norm(depth))
-    plt.imsave('test/depth.png', depth)
+    plt.imsave(f"{target_dir}/depth.png", depth)
 
     # Normal
     img, normal = nyu_normal[img_id]
     normal = (normal+1)/2
-    Image.fromarray((normal*255).numpy().transpose( 1,2,0 ).astype('uint8')).save('test/normal.png')
-
+    Image.fromarray((normal*255).numpy().transpose( 1,2,0 ).astype('uint8')).save(f"{target_dir}/normal.png")