def load_data(channels_first, dataset_type, root_dir, train_files, test_files, crop, min_depth, max_depth):
"""
DORN:
Input: Resized version of bgr_cropped
Output size: Same as unresized bgr_cropped
Comapre to: rawdepth_cropped, mask
:param root_dir: The root directory from which to load the dataset
:param use_dorn_normalization: Whether or not to normalize the rgb images according to DORN statistics.
:return: test: a NYUDepthv2TestDataset object.
"""
if dataset_type == "train":
dataset = NYUDepthv2LabeledDataset(root_dir, train_files["rgb"],
train_files["depth"],
train_files["rawdepth"],
crop,
transform=None, bgr_mode=False)
elif dataset_type == "test":
dataset = NYUDepthv2LabeledDataset(root_dir, test_files["rgb"],
test_files["depth"],
test_files["rawdepth"],
crop,
transform=None, bgr_mode=False)
transform_list = [
AddDepthMask(min_depth, max_depth, "rawdepth_cropped", "mask_cropped"),
AddDepthMask(min_depth, max_depth, "rawdepth", "mask"),
ToTensorAll(keys=["rgb", "rgb_cropped", "depth_cropped"],
channels_first=channels_first)
]
dataset.transform = transforms.Compose(transform_list)
return dataset
def load_data(root_dir, min_depth, max_depth, dorn_mode, sid_bins, alpha, beta,
offset, spad_config):
"""
:param root_dir: The root directory from which to load the dataset
:param use_dorn_normalization: Whether or not to normalize the rgb images according to DORN statistics.
:return: test: a NYUDepthv2TestDataset object.
"""
test = NYUDepthv2TestDataset(root_dir, transform=None, dorn_mode=dorn_mode)
if dorn_mode:
transform_mean = np.array([[[103.0626, 115.9029,
123.1516]]]).astype(np.float32)
else:
transform_mean = np.zeros((1, 1, 3)).astype(np.float32)
transform_var = np.ones((1, 1, 3)).astype(np.float32)
transform_list = [
AddDepthMask(min_depth, max_depth, "depth_cropped"),
Save(["rgb_cropped", "mask", "depth_cropped"], "_orig"),
]
if dorn_mode:
transform_list.append(
ResizeAll((353, 257), keys=["rgb_cropped", "depth_cropped"]))
transform_list += [
Normalize(transform_mean, transform_var, key="rgb_cropped"),
AddDepthMask(min_depth, max_depth, "depth_cropped"),
SimulateSpadIntensity("depth_cropped_orig",
"rgb_cropped_orig",
"mask_orig",
"spad",
min_depth,
max_depth,
spad_config["spad_bins"],
spad_config["photon_count"],
spad_config["dc_count"],
spad_config["fwhm_ps"],
spad_config["use_intensity"],
spad_config["use_squared_falloff"],
sid_obj=SID(sid_bins, alpha, beta, offset)),
]
# TODO: Determine if BGR or RGB
# Answer: RGB - need to flip for DORN.
if dorn_mode:
print("Using dataset in DORN mode.")
transform_list.append(
ToTensorAll(keys=[
"rgb_cropped", "rgb_cropped_orig", "depth_cropped",
"depth_cropped_orig", "mask", "mask_orig", "spad"
]))
else:
print("Using dataset in Wonka mode.")
# Don't flip channels of RGB input
# Flip channels of rgb_cropped used for intensity
transform_list.append(
ToTensorAll(keys=["rgb_cropped", "depth_cropped", "mask", "spad"]))
test.transform = transforms.Compose(transform_list)
return test
def load_data(channels_first, rootdir, scene_dict, min_depth, max_depth, crop):
"""
Load the dataset from files and initialize transforms accordingly.
:return: The dataset.
"""
dataset = PNGDataset(rootdir, scene_dict, min_depth, max_depth, crop)
transform_list = [
AddDepthMask(min_depth, max_depth, "rawDepth_cropped", "mask_cropped"),
AddDepthMask(min_depth, max_depth, "rawDepth", "mask"),
ToTensorAll(keys=["rgb", "rgb_cropped", "depth_cropped"],
channels_first=channels_first)
]
dataset.transform = transforms.Compose(transform_list)
return dataset
def load_data(root_dir, train_files, test_files, crop_file, min_depth, max_depth, dorn_mode):
"""
Wonka:
Input: rgb
Output size: Same as *_cropped
Compare to: rawdepth_cropped, mask
DORN:
Input: Resized version of bgr_cropped
Output size: Same as unresized bgr_cropped
Comapre to: rawdepth_cropped, mask
:param root_dir: The root directory from which to load the dataset
:param use_dorn_normalization: Whether or not to normalize the rgb images according to DORN statistics.
:return: test: a NYUDepthv2TestDataset object.
"""
train = NYUDepthv2LabeledDataset(root_dir, train_files["rgb"],
train_files["depth"],
train_files["rawdepth"],
crop_file,
transform=None, bgr_mode=dorn_mode)
test = NYUDepthv2LabeledDataset(root_dir, test_files["rgb"],
test_files["depth"],
test_files["rawdepth"],
crop_file,
transform=None, bgr_mode=dorn_mode)
transform_list = [
AddDepthMask(min_depth, max_depth, "rawdepth_cropped", "mask_cropped"),
AddDepthMask(min_depth, max_depth, "rawdepth", "mask")
]
if dorn_mode:
print("Using dataset in DORN mode.")
transform_mean = np.array([[[103.0626, 115.9029, 123.1516]]]).astype(np.float32)
transform_var = np.ones((1, 1, 3)).astype(np.float32)
transform_list += [
Save(["bgr"], "_orig"),
ResizeAll((353, 257), keys=["bgr"]),
Normalize(transform_mean, transform_var, key="bgr"),
ToTensorAll(keys=["bgr", "bgr_orig", "depth_cropped"],
channels_first=dorn_mode)
]
else:
print("Using dataset in Wonka mode.")
transform_list.append(ToTensorAll(keys=["rgb", "depth_cropped"], channels_first=dorn_mode))
train.transform = transforms.Compose(transform_list)
test.transform = transforms.Compose(transform_list)
return train, test
def load_data(root_dir, min_depth, max_depth, dorn_mode):
"""
:param root_dir: The root directory from which to load the dataset
:param use_dorn_normalization: Whether or not to normalize the rgb images according to DORN statistics.
:return: test: a NYUDepthv2TestDataset object.
"""
test = NYUDepthv2TestDataset(root_dir, transform=None, dorn_mode=dorn_mode)
if dorn_mode:
print("Using dataset in DORN mode.")
# Give data entries as DORN expects
transform_mean = np.array([[[103.0626, 115.9029, 123.1516]]]).astype(np.float32)
transform_var = np.ones((1, 1, 3))
test_transform = transforms.Compose([
AddDepthMask(min_depth, max_depth, "depth_cropped"),
Save(["rgb_cropped", "mask", "depth_cropped"], "_orig"),
ResizeAll((353, 257), keys=["rgb_cropped", "depth_cropped"]),
Normalize(transform_mean, transform_var, key="rgb_cropped"),
ToTensorAll(keys=["rgb", "rgb_cropped", "rgb_cropped_orig","depth", "depth_cropped",
"depth_cropped_orig", "mask", "mask_orig"])
])
# TODO: Determine if BGR or RGB
# Answer: RGB - need to flip for DORN.
else:
print("Using dataset in Wonka mode.")
# Give data entries as Wonka et. al. expects.
test_transform = None # Leaves numpy data as-is
test.transform = test_transform
return test
def load_data(root_dir, train_files, test_files, crop_file, min_depth,
max_depth, dorn_mode, sid_bins, alpha, beta, offset,
spad_config):
"""
Wonka:
Input: rgb
Output size: Same as *_cropped
Compare to: rawdepth_cropped, mask
DORN:
Input: Resized version of bgr_cropped
Output size: Same as unresized bgr_cropped
Comapre to: rawdepth_cropped, mask
:return: test: a NYUDepthv2TestDataset object.
"""
train = NYUDepthv2LabeledDataset(root_dir,
train_files["rgb"],
train_files["depth"],
train_files["rawdepth"],
crop_file,
transform=None,
bgr_mode=dorn_mode)
test = NYUDepthv2LabeledDataset(root_dir,
test_files["rgb"],
test_files["depth"],
test_files["rawdepth"],
crop_file,
transform=None,
bgr_mode=dorn_mode)
transform_list = [
AddDepthMask(min_depth, max_depth, "rawdepth_cropped"),
]
if dorn_mode:
print("Using dataset in DORN mode.")
transform_mean = np.array([[[103.0626, 115.9029,
123.1516]]]).astype(np.float32)
transform_var = np.ones((1, 1, 3)).astype(np.float32)
transform_list += [
Save(["bgr_cropped"], "_orig"),
ResizeAll((353, 257), keys=["bgr_cropped"]),
Normalize(transform_mean, transform_var, key="bgr_cropped"),
ToTensorAll(keys=[
"bgr", "bgr_cropped", "bgr_cropped_orig", "rawdepth_cropped",
"mask"
],
channels_first=dorn_mode)
]
else:
print("Using dataset in Wonka mode.")
transform_list.append(
ToTensorAll(keys=["rgb_cropped", "rawdepth_cropped", "mask"],
channels_first=dorn_mode))
train.transform = transforms.Compose(transform_list)
test.transform = transforms.Compose(transform_list)
return train, test
def load_data(dataset_type, root_dir, train_files, test_files, crop, min_depth,
max_depth):
"""
:param root_dir: The root directory from which to load the dataset
:param use_dorn_normalization: Whether or not to normalize the rgb images according to DORN statistics.
:return: test: a NYUDepthv2TestDataset object.
"""
print("Using dataset in DORN mode.")
if dataset_type == "train":
dataset = NYUDepthv2LabeledDataset(root_dir,
train_files["rgb"],
train_files["depth"],
train_files["rawdepth"],
crop,
transform=None,
bgr_mode=True)
elif dataset_type == "test":
dataset = NYUDepthv2LabeledDataset(root_dir,
test_files["rgb"],
test_files["depth"],
test_files["rawdepth"],
crop,
transform=None,
bgr_mode=True)
transform_list = [
AddDepthMask(min_depth, max_depth, "rawdepth_cropped", "mask_cropped"),
AddDepthMask(min_depth, max_depth, "rawdepth", "mask")
]
transform_mean = np.array([[[103.0626, 115.9029,
123.1516]]]).astype(np.float32)
transform_var = np.ones((1, 1, 3)).astype(np.float32)
transform_list += [
Save(["bgr"], "_orig"),
ResizeAll((353, 257), keys=["bgr"]),
Normalize(transform_mean, transform_var, key="bgr"),
ToTensorAll(keys=["bgr", "bgr_orig", "depth_cropped"],
channels_first=True)
]
dataset.transform = transforms.Compose(transform_list)
return dataset
def load_data(train_file, train_dir, val_file, val_dir, test_file, test_dir,
min_depth, max_depth, use_dorn_normalization, blacklist_file):
"""Generates training and validation datasets from
text files and directories. Sets up datasets with transforms.py.
*_file - string - a text file containing info for DepthDataset to load the images
*_dir - string - the folder containing the images to load
min_depth - the minimum depth for this dataset
max_depth - the maximum depth for this dataset
blacklist_file - string - a text file listing, on each line, an image_id of an image to exclude
from the dataset.
test_loader - bool - whether or not to test the loader and not set the dataset-wide mean and
variance.
Returns
-------
train, val, test - torch.data_utils.data.Dataset objects containing the relevant splits
"""
train = NYUDepthv2Dataset(train_file,
train_dir,
transform=None,
file_types=["rgb", "rawdepth"],
min_depth=min_depth,
max_depth=max_depth,
blacklist_file=blacklist_file)
train.rgb_mean, train.rgb_var = train.get_mean_and_var()
# Transform:
# Size is set to (353, 257) to conform to DORN conventions
# If use_dorn_normalization is true:
# Mean is set to np.array([[[103.0626, 115.9029, 123.1516]]]).astype(np.float32) to conform to DORN conventions
# Var is set to np.ones((1,1,3)) to conform to DORN conventions
if use_dorn_normalization:
transform_mean = np.array([[[103.0626, 115.9029,
123.1516]]]).astype(np.float32)
transform_var = np.ones((1, 1, 3))
else:
transform_mean = train.rgb_mean
transform_var = train.rgb_var
train_transform = transforms.Compose([
ResizeAll((353, 257), keys=["rgb", "depth", "rawdepth"]),
AddDepthMask(min_depth, max_depth, "rawdepth"), # introduces "mask"
RandomHorizontalFlipAll(flip_prob=0.5,
keys=["rgb", "depth", "rawdepth", "mask"]),
Normalize(transform_mean, transform_var,
key="rgb"), # introduces "rgb_orig"
ToTensorAll(keys=["rgb", "rgb_orig", "rawdepth", "mask"])
])
val_transform = transforms.Compose([
ResizeAll((353, 257), keys=["rgb", "depth", "rawdepth"]),
AddDepthMask(min_depth, max_depth, "rawdepth"),
Normalize(transform_mean, transform_var, key="rgb"),
ToTensorAll(keys=["rgb", "rgb_orig", "depth", "rawdepth", "mask"])
])
test_transform = transforms.Compose([
ResizeAll((353, 257), keys=["rgb", "rawdepth"]),
AddDepthMask(min_depth, max_depth, "rawdepth"),
Normalize(transform_mean, transform_var, key="rgb"),
ToTensorAll(keys=["rgb", "rgb_orig", "rawdepth", "mask"])
])
train.transform = train_transform
print("Loaded training dataset from {} with size {}.".format(
train_file, len(train)))
val = None
if val_file is not None:
val = NYUDepthv2Dataset(val_file,
val_dir,
transform=val_transform,
file_types=["rgb", "rawdepth"],
min_depth=min_depth,
max_depth=max_depth)
val.rgb_mean, val.rgb_var = train.rgb_mean, train.rgb_var
print("Loaded val dataset from {} with size {}.".format(
val_file, len(val)))
test = None
if test_file is not None:
test = NYUDepthv2Dataset(test_file,
test_dir,
transform=test_transform,
file_types=["rgb", "rawdepth"],
min_depth=min_depth,
max_depth=max_depth)
test.rgb_mean, test.rgb_var = train.rgb_mean, train.rgb_var
print("Loaded test dataset from {} with size {}.".format(
test_file, len(test)))
return train, val, test
def load_data(train_file, train_dir,
val_file, val_dir,
test_file, test_dir,
min_depth, max_depth, normalization,
sid_bins, alpha, beta, offset,
blacklist_file,
spad_config):
"""Generates training and validation datasets from
text files and directories. Sets up datasets with transforms.py.
*_file - string - a text file containing info for DepthDataset to load the images
*_dir - string - the folder containing the images to load
min_depth - the minimum depth for this dataset
max_depth - the maximum depth for this dataset
normalization - The type of normalization to use.
sid_bins - the number of Spacing Increasing Discretization bins to add.
blacklist_file - string - a text file listing, on each line, an image_id of an image to exclude
from the dataset.
test_loader - bool - whether or not to test the loader and not set the dataset-wide mean and
variance.
Returns
-------
train, val, test - torch.data_utils.data.Dataset objects containing the relevant splits
"""
# print(spad_config)
train = NYUDepthv2Dataset(train_file, train_dir, transform=None,
file_types=["rgb", "rawdepth"],
min_depth=min_depth, max_depth=max_depth,
blacklist_file=blacklist_file)
train.rgb_mean, train.rgb_var = train.get_mean_and_var()
# Transform:
# Size is set to (353, 257) to conform to DORN conventions
# If normalization == "dorn":
# Mean is set to np.array([[[103.0626, 115.9029, 123.1516]]]).astype(np.float32) to conform to DORN conventions
# Var is set to np.ones((1,1,3)) to conform to DORN conventions
if normalization == "dorn":
# Use normalization as in the github code for DORN.
print("Using DORN normalization.")
transform_mean = np.array([[[103.0626, 115.9029, 123.1516]]]).astype(np.float32)
transform_var = np.ones((1, 1, 3))
elif normalization == "none":
print("No normalization.")
transform_mean = np.zeros((1, 1, 3))
transform_var = np.ones((1, 1, 3))
else:
transform_mean = train.rgb_mean
transform_var = train.rgb_var
train_transform = transforms.Compose([
AddDepthMask(min_depth, max_depth, "rawdepth"),
Save(["rgb", "mask", "rawdepth"], "_orig"),
Normalize(transform_mean, transform_var, key="rgb"),
ResizeAll((353, 257), keys=["rgb", "rawdepth"]),
RandomHorizontalFlipAll(flip_prob=0.5, keys=["rgb", "rawdepth"]),
AddDepthMask(min_depth, max_depth, "rawdepth"), # "mask"
AddSIDDepth(sid_bins, alpha, beta, offset, "rawdepth"), # "rawdepth_sid" "rawdepth_sid_index"
SimulateSpadIntensity("rawdepth", "rgb", "mask", "spad", min_depth, max_depth,
spad_config["spad_bins"],
spad_config["photon_count"],
spad_config["dc_count"],
spad_config["fwhm_ps"],
spad_config["use_intensity"],
spad_config["use_squared_falloff"],
sid_obj=SID(sid_bins, alpha, beta, offset)),
ToTensorAll(keys=["rgb", "rgb_orig", "rawdepth", "rawdepth_orig",
"rawdepth_sid", "rawdepth_sid_index", "mask", "mask_orig", "spad"])
]
)
val_transform = transforms.Compose([
AddDepthMask(min_depth, max_depth, "rawdepth"),
Save(["rgb", "depth", "mask", "rawdepth"], "_orig"),
Normalize(transform_mean, transform_var, key="rgb"),
ResizeAll((353, 257), keys=["rgb", "depth", "rawdepth"]),
AddDepthMask(min_depth, max_depth, "rawdepth"),
AddSIDDepth(sid_bins, alpha, beta, offset, "rawdepth"),
SimulateSpadIntensity("depth", "rgb", "mask", "spad", min_depth, max_depth,
spad_config["spad_bins"],
spad_config["photon_count"],
spad_config["dc_count"],
spad_config["fwhm_ps"],
spad_config["use_intensity"],
spad_config["use_squared_falloff"],
sid_obj=SID(sid_bins, alpha, beta, offset)),
ToTensorAll(keys=["rgb", "rgb_orig", "depth", "depth_orig", "rawdepth", "rawdepth_orig",
"rawdepth_sid", "rawdepth_sid_index", "mask", "mask_orig", "spad"])
]
)
test_transform = transforms.Compose([
AddDepthMask(min_depth, max_depth, "rawdepth"),
Save(["rgb", "depth", "mask", "rawdepth"], "_orig"),
Normalize(transform_mean, transform_var, key="rgb"),
ResizeAll((353, 257), keys=["rgb", "depth", "rawdepth"]),
AddDepthMask(min_depth, max_depth, "rawdepth"),
AddSIDDepth(sid_bins, alpha, beta, offset, "rawdepth"),
SimulateSpadIntensity("depth", "rgb", "mask", "spad", min_depth, max_depth,
spad_config["spad_bins"],
spad_config["photon_count"],
spad_config["dc_count"],
spad_config["fwhm_ps"],
spad_config["use_intensity"],
spad_config["use_squared_falloff"],
sid_obj=SID(sid_bins, alpha, beta, offset)),
ToTensorAll(keys=["rgb", "rgb_orig", "depth", "depth_orig", "rawdepth", "rawdepth_orig",
"rawdepth_sid", "rawdepth_sid_index", "mask", "mask_orig", "spad"])
]
)
train.transform = train_transform
print("Loaded training dataset from {} with size {}.".format(train_file, len(train)))
val = None
if val_file is not None:
val = NYUDepthv2Dataset(val_file, val_dir, transform=val_transform,
file_types = ["rgb", "depth", "rawdepth"],
min_depth=min_depth, max_depth=max_depth)
val.rgb_mean, val.rgb_var = train.rgb_mean, train.rgb_var
print("Loaded val dataset from {} with size {}.".format(val_file, len(val)))
test = None
if test_file is not None:
test = NYUDepthv2Dataset(test_file, test_dir, transform=test_transform,
file_types = ["rgb", "depth", "rawdepth"],
min_depth=min_depth, max_depth=max_depth)
test.rgb_mean, test.rgb_var = train.rgb_mean, train.rgb_var
print("Loaded test dataset from {} with size {}.".format(test_file, len(test)))
return train, val, test