def __init__(
self,
path,
stage="train",
list_prefix="softras_",
image_size=None,
scale_focal=False,
max_imgs=100000,
z_near=0.0,
z_far=50.0,
skip_step=None,
):
"""
:param path dataset root path
:param image_size result image size (resizes if different); None to keep original size
:param scale_focal if true, assume focal length is specified for
image of side length 2 instead of actual image size. This is used
where image coordinates are placed in [-1, 1].
"""
super().__init__()
self.base_path = path
assert os.path.exists(self.base_path)
# transforms for images
self.image_to_tensor = get_image_to_tensor_balanced()
self.mask_to_tensor = get_mask_to_tensor()
self.image_size = image_size
self._coord_trans_world = torch.tensor(
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
dtype=torch.float32,
)
self._coord_trans_cam = torch.tensor(
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
dtype=torch.float32,
)
self.scale_focal = scale_focal
self.max_imgs = max_imgs
self.z_near = z_near
self.z_far = z_far
self.lindisp = False
self.rgb_paths = []
self.project_mats = []
files = os.listdir(self.base_path)
for f in files:
if f.endswith(".png"):
self.rgb_paths.append(os.path.join(self.base_path, f))
elif f.endswith(".txt"):
self.project_mats.append(os.path.join(self.base_path, f))
self.rgb_paths.sort()
self.project_mats.sort()
def __init__(self,
path,
stage="train",
image_size=(128, 128),
world_scale=1.0,
only_load_these_ids=None):
"""
:param stage train | val | test
:param image_size result image size (resizes if different)
:param world_scale amount to scale entire world by
"""
super().__init__()
self.base_path = path + "_" + stage
self.dataset_name = os.path.basename(path)
print("Loading SRN dataset", self.base_path, "name:",
self.dataset_name)
self.stage = stage
assert os.path.exists(self.base_path)
is_chair = "chair" in self.dataset_name
if is_chair and stage == "train":
# Ugly thing from SRN's public dataset
tmp = os.path.join(self.base_path, "chairs_2.0_train")
if os.path.exists(tmp):
self.base_path = tmp
self.intrins = sorted(
glob.glob(os.path.join(self.base_path, "*", "intrinsics.txt")))
if only_load_these_ids:
self.intrins = [
p for p in self.intrins
if any(id in p for id in only_load_these_ids)
]
self.image_to_tensor = get_image_to_tensor_balanced()
self.mask_to_tensor = get_mask_to_tensor()
self.image_size = image_size
self.world_scale = world_scale
self._coord_trans = torch.diag(
torch.tensor([1, -1, -1, 1], dtype=torch.float32))
if is_chair:
self.z_near = 1.25
self.z_far = 2.75
else:
self.z_near = 0.8
self.z_far = 1.8
if "renders" in self.dataset_name: # For my dataset
self.z_near = 0.3
self.z_far = 1.9
self.lindisp = False
def __init__(self, path, z_near=4, z_far=9, n_views=None):
super().__init__()
self.base_path = path
print("Loading NeRF synthetic dataset", self.base_path)
trans_files = []
TRANS_FILE = "transforms.json"
for root, directories, filenames in os.walk(path):
if TRANS_FILE in filenames:
trans_files.append(os.path.join(root, TRANS_FILE))
self.trans_files = trans_files
self.image_to_tensor = get_image_to_tensor_balanced()
self.mask_to_tensor = get_mask_to_tensor()
self.z_near = z_near
self.z_far = z_far
self.compose = compose
self.n_views = n_views
def __init__(self,
path,
stage="train",
z_near=4,
z_far=9,
n_views=None,
compose=False,
split_seed=1234,
val_frac=0.2,
test_frac=0.2):
"""
:param path data directory
:para stage train | val | test
:param z_near near bound for ray sampling
:param z_far far bound for ray sampling
:param n_views optional: expected number of views per object in dataset
for validity checking only
:param compose if true, adds background to images.
Dataset must have background '*_env.png' in addition to foreground
'*_obj.png' for each view.
"""
super().__init__()
self.base_path = path
self.stage = stage
print("Loading NeRF synthetic dataset", self.base_path, "stage",
self.stage)
trans_files = []
TRANS_FILE = "transforms.json"
for root, directories, filenames in os.walk(path):
if TRANS_FILE in filenames:
trans_files.append(osp.join(root, TRANS_FILE))
self.trans_files = trans_files
self.image_to_tensor = get_image_to_tensor_balanced()
self.mask_to_tensor = get_mask_to_tensor()
self.z_near = z_near
self.z_far = z_far
self.lindisp = False
self.compose = compose
self.n_views = n_views
# Load data split
self._load_split(val_frac, test_frac, split_seed)
def __init__(self, path, stage="train", z_near=4, z_far=9, n_views=None):
super().__init__()
path = os.path.join(path, stage)
self.base_path = path
print("Loading NeRF synthetic dataset", self.base_path)
trans_files = []
TRANS_FILE = "transforms.json"
for root, directories, filenames in os.walk(self.base_path):
if TRANS_FILE in filenames:
trans_files.append(os.path.join(root, TRANS_FILE))
self.trans_files = trans_files
self.image_to_tensor = get_image_to_tensor_balanced()
self.mask_to_tensor = get_mask_to_tensor()
self.z_near = z_near
self.z_far = z_far
self.lindisp = False
self.n_views = n_views
print("{} instances in split {}".format(len(self.trans_files), stage))
def __init__(
self,
path,
stage="train",
list_prefix="softras_",
image_size=None,
sub_format="shapenet",
scale_focal=True,
max_imgs=100000,
z_near=1.2,
z_far=4.0,
skip_step=None,
):
"""
:param path dataset root path, contains metadata.yml
:param stage train | val | test
:param list_prefix prefix for split lists: <list_prefix>[train, val, test].lst
:param image_size result image size (resizes if different); None to keep original size
:param sub_format shapenet | dtu dataset sub-type.
:param scale_focal if true, assume focal length is specified for
image of side length 2 instead of actual image size. This is used
where image coordinates are placed in [-1, 1].
"""
super().__init__()
self.base_path = path
assert os.path.exists(self.base_path)
cats = [
x for x in glob.glob(os.path.join(path, "*")) if os.path.isdir(x)
]
if stage == "train":
file_lists = [
os.path.join(x, list_prefix + "train.lst") for x in cats
]
elif stage == "val":
file_lists = [
os.path.join(x, list_prefix + "val.lst") for x in cats
]
elif stage == "test":
file_lists = [
os.path.join(x, list_prefix + "test.lst") for x in cats
]
all_objs = []
for file_list in file_lists:
if not os.path.exists(file_list):
continue
base_dir = os.path.dirname(file_list)
cat = os.path.basename(base_dir)
with open(file_list, "r") as f:
objs = [(cat, os.path.join(base_dir, x.strip()))
for x in f.readlines()]
all_objs.extend(objs)
self.all_objs = all_objs
self.stage = stage
self.image_to_tensor = get_image_to_tensor_balanced()
self.mask_to_tensor = get_mask_to_tensor()
print(
"Loading DVR dataset",
self.base_path,
"stage",
stage,
len(self.all_objs),
"objs",
"type:",
sub_format,
)
self.image_size = image_size
if sub_format == "dtu":
self._coord_trans_world = torch.tensor(
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
dtype=torch.float32,
)
self._coord_trans_cam = torch.tensor(
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
dtype=torch.float32,
)
else:
self._coord_trans_world = torch.tensor(
[[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]],
dtype=torch.float32,
)
self._coord_trans_cam = torch.tensor(
[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],
dtype=torch.float32,
)
self.sub_format = sub_format
self.scale_focal = scale_focal
self.max_imgs = max_imgs
self.z_near = z_near
self.z_far = z_far
self.lindisp = False