diff --git a/.gitignore b/.gitignore index 0617278..5c1996b 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ colored_* test_labels_13 train_labels_13 *.zip -*.mat \ No newline at end of file +*.mat +**/__pycache__ diff --git a/download_and_extract.sh b/download_and_extract.sh index 4657652..505629a 100644 --- a/download_and_extract.sh +++ b/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 \ No newline at end of file +python extract_nyuv2.py --mat nyu_depth_v2_labeled.mat --normal_zip nyuv2_surfacenormal_metadata.zip --data_root NYUv2 --save_colored \ No newline at end of file diff --git a/extract_nyuv2.py b/extract_nyuv2.py index 7ff37eb..32515c4 100644 --- a/extract_nyuv2.py +++ b/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' )) - diff --git a/nyuv2.py b/nyuv2.py index e6ab190..7bb74d7 100644 --- a/nyuv2.py +++ b/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") \ No newline at end of file