nyuv2 extraction tools

Author: VainF

.gitignore

@@ -0,0 +1,6 @@
+NYUv2
+colored_*
+test_labels_13
+train_labels_13
+*.zip
+*.mat

README.md

@@ -1,3 +1,58 @@
+# NYUv2 Python Toolkits
+
+<div>
+<img src="test/image.png"        width="18%">
+<img src="test/semantic13.png"   width="18%">
+<img src="test/semantic40.png"   width="18%">
+<img src="test/depth.png"        width="18%">
+<img src="test/normal.png"       width="18%">
+</div>
+
+This repo provides extraction tools and a pytorch dataloader written in python for NYUv2 dataset. All meta data comes from [ankurhanda/nyuv2-meta-data](https://github.com/ankurhanda/nyuv2-meta-data) 
+
+Supported Tasks:
+
+* Semantic Segmentation (13 classes and 40 classes)
+* Depth Estimation
+* Normal
+
+## Extraction
+
+```bash
+bash download_and_extract.sh
+```
+or 
+
+```bash
+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
+
+python extract_nyuv2.py --mat nyu_depth_v2_labeled.mat --normal_zip nyu_normals_gt.zip  --data_root NYUv2 --save_colored
+```
+All images and labels will be extracted to ./NYUv2 as the following:
+
+```
+NYUv2/
+    /image
+        /train
+            00003.png
+            00004.png
+            ...
+        /test
+    /seg13
+        /train
+            00003.png
+            00004.png
+            ...
+        /test
+    /seg40
+    /depth
+    /normal
+```
+
+
+# nyuv2-meta-data
+
 ## What does this repository contain?
 
 This repository contains 13 class labels for both train and test dataset in NYUv2. This is to avoid any hassle involved in parsing the data from the .mat files. If you are looking to train a network to do 13 class segmentation from RGB data, then this repository can provide you both the training/test dataset as well the corresponding ground truth labels. However, if your networks needs additionally depth data (either depth image or DHA features) then you will need to download the dataset from the [NYUv2 website](http://horatio.cs.nyu.edu/mit/silberman/nyu_depth_v2/nyu_depth_v2_labeled.mat) (~2.8GB) as well as the corresponding [toolbox](http://cs.nyu.edu/~silberman/code/toolbox_nyu_depth_v2.zip). To summarise, this repository contains the following

download_and_extract.sh

@@ -0,0 +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
+
+python extract_nyuv2.py --mat nyu_depth_v2_labeled.mat --normal_zip nyu_normals_gt.zip  --data_root NYUv2 --save_colored

extract_nyuv2.py

@@ -0,0 +1,159 @@
+import os
+import sys
+import cv2
+import h5py
+import argparse
+import multiprocessing
+import numpy as np
+from skimage import io
+from scipy.io import loadmat
+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):
+        return ((byteval & (1 << idx)) != 0)
+
+    dtype = 'float32' if normalized else 'uint8'
+    cmap = np.zeros((N, 3), dtype=dtype)
+    for i in range(N):
+        r = g = b = 0
+        c = i
+        for j in range(8):
+            r = r | (bitget(c, 0) << 7-j)
+            g = g | (bitget(c, 1) << 7-j)
+            b = b | (bitget(c, 2) << 7-j)
+            c = c >> 3
+
+        cmap[i] = np.array([r, g, b])
+
+    cmap = cmap/255 if normalized else cmap
+    return cmap
+
+
+def extract_images(imgs, splits, IMAGE_DIR):
+    print("Extracting images...")
+    imgs = imgs.transpose(0, 3, 2, 1)
+    for s in ['train', 'test']:
+        os.makedirs(os.path.join(IMAGE_DIR, s), exist_ok=True)
+        idxs = splits[s+'Ndxs'].reshape(-1)
+        for idx in tqdm(idxs):
+            img = imgs[idx-1]
+            path = os.path.join(IMAGE_DIR, s, '%05d.png' % (idx))
+            io.imsave(path, img)
+
+
+def extract_labels(labels, splits, SEG40_DIR, SEG13_DIR, save_colored=True):
+    mapping40 = loadmat('classMapping40.mat')['mapClass'][0]
+    mapping13 = loadmat('class13Mapping.mat')['classMapping13'][0][0][0][0]
+    mapping40 = np.insert(mapping40, 0, 0)
+    mapping13 = np.insert(mapping13, 0, 0)
+    labels = labels.transpose([0, 2, 1])
+
+    labels_40 = mapping40[labels]
+    labels_13 = mapping13[labels_40]
+
+    if save_colored:
+        cmap = colormap()
+        os.makedirs('colored_40', exist_ok=True)
+        os.makedirs('colored_13', exist_ok=True)
+
+    print("Extracting labels (40 classes)...")
+
+    for s in ['train', 'test']:
+        os.makedirs(os.path.join(SEG40_DIR, s), exist_ok=True)
+        idxs = splits[s+'Ndxs'].reshape(-1)
+
+        for idx in tqdm(idxs):
+            lbl = labels_40[idx-1]
+            path = os.path.join(SEG40_DIR, s, '%05d.png' % (idx))
+            io.imsave(path, lbl, check_contrast=False)
+            if save_colored:
+                colored_lbl = cmap[lbl]
+                io.imsave('colored_40/%05d.png' % idx, colored_lbl)
+
+    print("Extracting labels (13 classes)...")
+    for s in ['train', 'test']:
+        os.makedirs(os.path.join(SEG13_DIR, s), exist_ok=True)
+        idxs = splits[s+'Ndxs'].reshape(-1)
+
+        for idx in tqdm(idxs):
+            lbl = labels_13[idx-1]
+            path = os.path.join(SEG13_DIR, s, '%05d.png' % (idx))
+            io.imsave(path, lbl, check_contrast=False)
+            if save_colored:
+                colored_lbl = cmap[lbl]
+                io.imsave('colored_13/%05d.png' % idx, colored_lbl)
+
+
+def extract_depths(depths, splits, DEPTH_DIR, save_colored=False):
+    depths = depths.transpose(0, 2, 1)
+    if save_colored:
+        os.makedirs('colored_depth', exist_ok=True)
+    print("Extracting depths...")
+    depths = (depths*1e3).astype(np.uint16)
+
+    for s in ['train', 'test']:
+        os.makedirs(os.path.join(DEPTH_DIR, s), exist_ok=True)
+        idxs = splits[s+'Ndxs'].reshape(-1)
+        for idx in tqdm(idxs):
+            depth = depths[idx-1]
+            path = os.path.join(DEPTH_DIR, s, '%05d.png' % (idx))
+            io.imsave(path, depth, check_contrast=False)
+
+            if save_colored:
+                norm = plt.Normalize()
+                colored = plt.cm.jet(norm(depth))
+                plt.imsave('colored_depth/%05d.png' % (idx), colored)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='RYU DATA Extraction')
+    parser.add_argument('--mat', type=str, required=True,
+                        help='downloaded NYUv2 mat files. http://horatio.cs.nyu.edu/mit/silberman/nyu_depth_v2/nyu_depth_v2_labeled.mat')
+    parser.add_argument('--data_root', type=str,
+                        required=True, help='the output dir')
+    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')
+
+    args = parser.parse_args()
+
+    MAT_FILE = os.path.expanduser(args.mat)
+    DATA_ROOT = os.path.expanduser(args.data_root)
+    assert os.path.exists(MAT_FILE), "file does not exists: %s" % MAT_FILE
+
+    os.makedirs(DATA_ROOT, exist_ok=True)
+    IMAGE_DIR = os.path.join(DATA_ROOT, 'image')
+    SEG40_DIR = os.path.join(DATA_ROOT, 'seg40')
+    SEG13_DIR = os.path.join(DATA_ROOT, 'seg13')
+    DEPTH_DIR = os.path.join(DATA_ROOT, 'depth')
+    splits = loadmat('splits.mat')
+
+    os.makedirs(IMAGE_DIR, exist_ok=True)
+    os.makedirs(SEG40_DIR, exist_ok=True)
+    os.makedirs(SEG13_DIR, exist_ok=True)
+    os.makedirs(DEPTH_DIR, exist_ok=True)
+    import time
+    with h5py.File(MAT_FILE, 'r') as fr:
+        images = fr["images"]
+        labels = fr["labels"]
+        depths = fr["depths"]
+
+        #extract_images(np.array(images), splits, IMAGE_DIR)
+        #extract_labels(np.array(labels), splits, SEG40_DIR, SEG13_DIR, save_colored=args.save_colored )
+        extract_depths(np.array(depths), splits, DEPTH_DIR, save_colored=args.save_colored)
+
+        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)
+        
+        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

@@ -0,0 +1,164 @@
+#coding:utf-8
+import os
+import torch
+import torch.utils.data as data
+from PIL import Image
+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):
+        return ((byteval & (1 << idx)) != 0)
+
+    dtype = 'float32' if normalized else 'uint8'
+    cmap = np.zeros((N, 3), dtype=dtype)
+    for i in range(N):
+        r = g = b = 0
+        c = i
+        for j in range(8):
+            r = r | (bitget(c, 0) << 7-j)
+            g = g | (bitget(c, 1) << 7-j)
+            b = b | (bitget(c, 2) << 7-j)
+            c = c >> 3
+
+        cmap[i] = np.array([r, g, b])
+
+    cmap = cmap/255 if normalized else cmap
+    return cmap
+
+class NYUv2(VisionDataset):
+    """NYUv2 dataset
+    
+    Args:
+        root (string): Root directory path.
+        split (string, optional): 'train' for training set, and 'test' for test set. Default: 'train'.
+        target_type (string, optional): Type of target to use, ``semantic``, ``depth`` or ``normal``. 
+        num_classes (int, optional): The number of classes, must be 40 or 13. Default:13.
+        transform (callable, optional): A function/transform that takes in an PIL image and returns a transformed version.
+        target_transform (callable, optional): A function/transform that takes in the target and transforms it.
+        transforms (callable, optional): A function/transform that takes input sample and its target as entry and returns a transformed version.
+    """
+    cmap = colormap()
+    def __init__(self,
+                 root,
+                 split='train',
+                 target_type='semantic',
+                 num_classes=13,
+                 transforms=None,
+                 transform=None,
+                 target_transform=None):
+        super( NYUv2, self ).__init__(root, transforms=transforms, transform=transform, target_transform=target_transform)
+        assert(split in ('train', 'test'))
+        self.root = root
+        self.split = split
+        self.target_type = target_type
+        self.num_classes = num_classes
+        self.train_idx = np.array([255, ] + list(range(num_classes)))
+        
+        split_mat = loadmat(os.path.join(self.root, 'splits.mat'))
+        idxs = split_mat[self.split+'Ndxs'].reshape(-1) - 1
+
+        img_names = os.listdir( os.path.join(self.root, 'image', self.split) )
+        img_names.sort()
+        images_dir = os.path.join(self.root, 'image', self.split)
+        self.images = [os.path.join(images_dir, name) for name in img_names]
+
+        if self.target_type=='semantic':
+            semantic_dir = os.path.join(self.root, 'seg%d'%self.num_classes, self.split)
+            self.labels = [os.path.join(semantic_dir, name) for name in img_names]
+            self.targets = self.labels
+
+        if self.target_type=='depth':
+            depth_dir = os.path.join(self.root, 'depth', self.split)
+            self.depths = [os.path.join(depth_dir, name) for name in img_names]
+            self.targets = self.depths
+        
+        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.targets = self.normals
+        
+    def __getitem__(self, idx):
+        image = Image.open(self.images[idx])
+        target = Image.open(self.targets[idx])
+        if self.transforms is not None:
+            image, target = self.transforms( image, target )
+        return image, target
+
+    def __len__(self):
+        return len(self.images)
+
+if __name__=='__main__':
+    from torchvision import transforms
+    import matplotlib.pyplot as plt
+    nyu_semantic13 = NYUv2( root='NYUv2', split='train', target_type='semantic', num_classes=13, 
+                            transform=transforms.Compose([
+                                transforms.Resize(512),
+                                transforms.ToTensor()
+                            ]),
+                            target_transform=transforms.Compose([
+                                transforms.Resize(512, interpolation=Image.NEAREST),
+                                transforms.Lambda(lambda lbl: torch.from_numpy( np.array(lbl, dtype='uint8')-1 ) ) # 0->255, 1->0, 2->1
+                            ]),  
+                        )
+
+    nyu_semantic40 = NYUv2( root='NYUv2', split='train', target_type='semantic', num_classes=40, 
+                            transform=transforms.Compose([
+                                transforms.Resize(512),
+                                transforms.ToTensor()
+                            ]),
+                            target_transform=transforms.Compose([
+                                transforms.Resize(512, interpolation=Image.NEAREST),
+                                transforms.Lambda(lambda lbl: torch.from_numpy( np.array(lbl, dtype='uint8')-1 ) ) # 0->255, 1->0, 2->1
+                            ]),  
+                        )
+
+    nyu_depth = NYUv2( root='NYUv2', split='train', target_type='depth', 
+                            transform=transforms.Compose([
+                                transforms.Resize(512),
+                                transforms.ToTensor()
+                            ]),
+                            target_transform=transforms.Compose([
+                                transforms.Resize(512),
+                                transforms.Lambda(lambda lbl: torch.from_numpy( np.array(lbl, dtype='float') )/1e3 ) # uint16 to depth
+                            ]),  
+                        )
+
+    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)
+    # Semantic
+    img_id = 0
+    
+    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')
+
+    img, lbl40 = nyu_semantic40[img_id]
+    Image.fromarray( nyu_semantic40.cmap[ (lbl40.numpy().astype('uint8')+1) ] ).save('test/semantic40.png')
+
+    # Depth
+    img, depth = nyu_depth[img_id]
+    norm = plt.Normalize()
+    depth = plt.cm.jet(norm(depth))
+    plt.imsave('test/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')
+
+