mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-03-20 13:05:59 +08:00
[pytorch3d[ Remove LlffDatasetMapProvider and BlenderDatasetMapProvider
Summary: No one is using these. (The minify part has been broken for a couple of years, too) Reviewed By: patricklabatut Differential Revision: D96977684 fbshipit-source-id: 4708dfd37b14d1930f1370677eb126a61a0d9d3c
This commit is contained in:
committed by
meta-codesync[bot]
parent
52164b8324
commit
b6a77ad7aa
@@ -3,11 +3,6 @@ pytorch3d.implicitron.dataset specific datasets
|
|||||||
|
|
||||||
specific datasets
|
specific datasets
|
||||||
|
|
||||||
.. automodule:: pytorch3d.implicitron.dataset.blender_dataset_map_provider
|
|
||||||
:members:
|
|
||||||
:undoc-members:
|
|
||||||
:show-inheritance:
|
|
||||||
|
|
||||||
.. automodule:: pytorch3d.implicitron.dataset.json_index_dataset_map_provider
|
.. automodule:: pytorch3d.implicitron.dataset.json_index_dataset_map_provider
|
||||||
:members:
|
:members:
|
||||||
:undoc-members:
|
:undoc-members:
|
||||||
@@ -18,11 +13,6 @@ specific datasets
|
|||||||
:undoc-members:
|
:undoc-members:
|
||||||
:show-inheritance:
|
:show-inheritance:
|
||||||
|
|
||||||
.. automodule:: pytorch3d.implicitron.dataset.llff_dataset_map_provider
|
|
||||||
:members:
|
|
||||||
:undoc-members:
|
|
||||||
:show-inheritance:
|
|
||||||
|
|
||||||
.. automodule:: pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider
|
.. automodule:: pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider
|
||||||
:members:
|
:members:
|
||||||
:undoc-members:
|
:undoc-members:
|
||||||
|
|||||||
@@ -1,56 +0,0 @@
|
|||||||
defaults:
|
|
||||||
- overfit_singleseq_base
|
|
||||||
- _self_
|
|
||||||
exp_dir: "./data/overfit_nerf_blender_repro/${oc.env:BLENDER_SINGLESEQ_CLASS}"
|
|
||||||
data_source_ImplicitronDataSource_args:
|
|
||||||
data_loader_map_provider_SequenceDataLoaderMapProvider_args:
|
|
||||||
dataset_length_train: 100
|
|
||||||
dataset_map_provider_class_type: BlenderDatasetMapProvider
|
|
||||||
dataset_map_provider_BlenderDatasetMapProvider_args:
|
|
||||||
base_dir: ${oc.env:BLENDER_DATASET_ROOT}/${oc.env:BLENDER_SINGLESEQ_CLASS}
|
|
||||||
n_known_frames_for_test: null
|
|
||||||
object_name: ${oc.env:BLENDER_SINGLESEQ_CLASS}
|
|
||||||
path_manager_factory_class_type: PathManagerFactory
|
|
||||||
path_manager_factory_PathManagerFactory_args:
|
|
||||||
silence_logs: true
|
|
||||||
|
|
||||||
model_factory_ImplicitronModelFactory_args:
|
|
||||||
model_class_type: "OverfitModel"
|
|
||||||
model_OverfitModel_args:
|
|
||||||
mask_images: false
|
|
||||||
raysampler_class_type: AdaptiveRaySampler
|
|
||||||
raysampler_AdaptiveRaySampler_args:
|
|
||||||
n_pts_per_ray_training: 64
|
|
||||||
n_pts_per_ray_evaluation: 64
|
|
||||||
n_rays_per_image_sampled_from_mask: 4096
|
|
||||||
stratified_point_sampling_training: true
|
|
||||||
stratified_point_sampling_evaluation: false
|
|
||||||
scene_extent: 2.0
|
|
||||||
scene_center:
|
|
||||||
- 0.0
|
|
||||||
- 0.0
|
|
||||||
- 0.0
|
|
||||||
renderer_MultiPassEmissionAbsorptionRenderer_args:
|
|
||||||
density_noise_std_train: 0.0
|
|
||||||
n_pts_per_ray_fine_training: 128
|
|
||||||
n_pts_per_ray_fine_evaluation: 128
|
|
||||||
raymarcher_EmissionAbsorptionRaymarcher_args:
|
|
||||||
blend_output: false
|
|
||||||
loss_weights:
|
|
||||||
loss_rgb_mse: 1.0
|
|
||||||
loss_prev_stage_rgb_mse: 1.0
|
|
||||||
loss_mask_bce: 0.0
|
|
||||||
loss_prev_stage_mask_bce: 0.0
|
|
||||||
loss_autodecoder_norm: 0.00
|
|
||||||
|
|
||||||
optimizer_factory_ImplicitronOptimizerFactory_args:
|
|
||||||
exponential_lr_step_size: 3001
|
|
||||||
lr_policy: LinearExponential
|
|
||||||
linear_exponential_lr_milestone: 200
|
|
||||||
|
|
||||||
training_loop_ImplicitronTrainingLoop_args:
|
|
||||||
max_epochs: 6000
|
|
||||||
metric_print_interval: 10
|
|
||||||
store_checkpoints_purge: 3
|
|
||||||
test_when_finished: true
|
|
||||||
validation_interval: 100
|
|
||||||
@@ -1,55 +0,0 @@
|
|||||||
defaults:
|
|
||||||
- repro_singleseq_base
|
|
||||||
- _self_
|
|
||||||
exp_dir: "./data/nerf_blender_repro/${oc.env:BLENDER_SINGLESEQ_CLASS}"
|
|
||||||
data_source_ImplicitronDataSource_args:
|
|
||||||
data_loader_map_provider_SequenceDataLoaderMapProvider_args:
|
|
||||||
dataset_length_train: 100
|
|
||||||
dataset_map_provider_class_type: BlenderDatasetMapProvider
|
|
||||||
dataset_map_provider_BlenderDatasetMapProvider_args:
|
|
||||||
base_dir: ${oc.env:BLENDER_DATASET_ROOT}/${oc.env:BLENDER_SINGLESEQ_CLASS}
|
|
||||||
n_known_frames_for_test: null
|
|
||||||
object_name: ${oc.env:BLENDER_SINGLESEQ_CLASS}
|
|
||||||
path_manager_factory_class_type: PathManagerFactory
|
|
||||||
path_manager_factory_PathManagerFactory_args:
|
|
||||||
silence_logs: true
|
|
||||||
|
|
||||||
model_factory_ImplicitronModelFactory_args:
|
|
||||||
model_GenericModel_args:
|
|
||||||
mask_images: false
|
|
||||||
raysampler_class_type: AdaptiveRaySampler
|
|
||||||
raysampler_AdaptiveRaySampler_args:
|
|
||||||
n_pts_per_ray_training: 64
|
|
||||||
n_pts_per_ray_evaluation: 64
|
|
||||||
n_rays_per_image_sampled_from_mask: 4096
|
|
||||||
stratified_point_sampling_training: true
|
|
||||||
stratified_point_sampling_evaluation: false
|
|
||||||
scene_extent: 2.0
|
|
||||||
scene_center:
|
|
||||||
- 0.0
|
|
||||||
- 0.0
|
|
||||||
- 0.0
|
|
||||||
renderer_MultiPassEmissionAbsorptionRenderer_args:
|
|
||||||
density_noise_std_train: 0.0
|
|
||||||
n_pts_per_ray_fine_training: 128
|
|
||||||
n_pts_per_ray_fine_evaluation: 128
|
|
||||||
raymarcher_EmissionAbsorptionRaymarcher_args:
|
|
||||||
blend_output: false
|
|
||||||
loss_weights:
|
|
||||||
loss_rgb_mse: 1.0
|
|
||||||
loss_prev_stage_rgb_mse: 1.0
|
|
||||||
loss_mask_bce: 0.0
|
|
||||||
loss_prev_stage_mask_bce: 0.0
|
|
||||||
loss_autodecoder_norm: 0.00
|
|
||||||
|
|
||||||
optimizer_factory_ImplicitronOptimizerFactory_args:
|
|
||||||
exponential_lr_step_size: 3001
|
|
||||||
lr_policy: LinearExponential
|
|
||||||
linear_exponential_lr_milestone: 200
|
|
||||||
|
|
||||||
training_loop_ImplicitronTrainingLoop_args:
|
|
||||||
max_epochs: 6000
|
|
||||||
metric_print_interval: 10
|
|
||||||
store_checkpoints_purge: 3
|
|
||||||
test_when_finished: true
|
|
||||||
validation_interval: 100
|
|
||||||
@@ -13,13 +13,6 @@ hydra:
|
|||||||
data_source_ImplicitronDataSource_args:
|
data_source_ImplicitronDataSource_args:
|
||||||
dataset_map_provider_class_type: ???
|
dataset_map_provider_class_type: ???
|
||||||
data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
|
data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
|
||||||
dataset_map_provider_BlenderDatasetMapProvider_args:
|
|
||||||
base_dir: ???
|
|
||||||
object_name: ???
|
|
||||||
path_manager_factory_class_type: PathManagerFactory
|
|
||||||
n_known_frames_for_test: null
|
|
||||||
path_manager_factory_PathManagerFactory_args:
|
|
||||||
silence_logs: true
|
|
||||||
dataset_map_provider_JsonIndexDatasetMapProvider_args:
|
dataset_map_provider_JsonIndexDatasetMapProvider_args:
|
||||||
category: ???
|
category: ???
|
||||||
task_str: singlesequence
|
task_str: singlesequence
|
||||||
@@ -91,14 +84,6 @@ data_source_ImplicitronDataSource_args:
|
|||||||
sort_frames: false
|
sort_frames: false
|
||||||
path_manager_factory_PathManagerFactory_args:
|
path_manager_factory_PathManagerFactory_args:
|
||||||
silence_logs: true
|
silence_logs: true
|
||||||
dataset_map_provider_LlffDatasetMapProvider_args:
|
|
||||||
base_dir: ???
|
|
||||||
object_name: ???
|
|
||||||
path_manager_factory_class_type: PathManagerFactory
|
|
||||||
n_known_frames_for_test: null
|
|
||||||
path_manager_factory_PathManagerFactory_args:
|
|
||||||
silence_logs: true
|
|
||||||
downscale_factor: 4
|
|
||||||
dataset_map_provider_RenderedMeshDatasetMapProvider_args:
|
dataset_map_provider_RenderedMeshDatasetMapProvider_args:
|
||||||
num_views: 40
|
num_views: 40
|
||||||
data_file: null
|
data_file: null
|
||||||
|
|||||||
@@ -1,55 +0,0 @@
|
|||||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
||||||
# All rights reserved.
|
|
||||||
#
|
|
||||||
# This source code is licensed under the BSD-style license found in the
|
|
||||||
# LICENSE file in the root directory of this source tree.
|
|
||||||
|
|
||||||
# pyre-unsafe
|
|
||||||
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from pytorch3d.implicitron.tools.config import registry
|
|
||||||
|
|
||||||
from .load_blender import load_blender_data
|
|
||||||
from .single_sequence_dataset import (
|
|
||||||
_interpret_blender_cameras,
|
|
||||||
SingleSceneDatasetMapProviderBase,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@registry.register
|
|
||||||
class BlenderDatasetMapProvider(SingleSceneDatasetMapProviderBase):
|
|
||||||
"""
|
|
||||||
Provides data for one scene from Blender synthetic dataset.
|
|
||||||
Uses the code in load_blender.py
|
|
||||||
|
|
||||||
Members:
|
|
||||||
base_dir: directory holding the data for the scene.
|
|
||||||
object_name: The name of the scene (e.g. "lego"). This is just used as a label.
|
|
||||||
It will typically be equal to the name of the directory self.base_dir.
|
|
||||||
path_manager_factory: Creates path manager which may be used for
|
|
||||||
interpreting paths.
|
|
||||||
n_known_frames_for_test: If set, training frames are included in the val
|
|
||||||
and test datasets, and this many random training frames are added to
|
|
||||||
each test batch. If not set, test batches each contain just a single
|
|
||||||
testing frame.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def _load_data(self) -> None:
|
|
||||||
path_manager = self.path_manager_factory.get()
|
|
||||||
images, poses, _, hwf, i_split = load_blender_data(
|
|
||||||
self.base_dir,
|
|
||||||
testskip=1,
|
|
||||||
path_manager=path_manager,
|
|
||||||
)
|
|
||||||
H, W, focal = hwf
|
|
||||||
images_masks = torch.from_numpy(images).permute(0, 3, 1, 2)
|
|
||||||
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.poses = _interpret_blender_cameras(poses, focal)
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.images = images_masks[:, :3]
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.fg_probabilities = images_masks[:, 3:4]
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.i_split = i_split
|
|
||||||
@@ -64,16 +64,12 @@ class ImplicitronDataSource(DataSourceBase):
|
|||||||
def pre_expand(cls) -> None:
|
def pre_expand(cls) -> None:
|
||||||
# use try/finally to bypass cinder's lazy imports
|
# use try/finally to bypass cinder's lazy imports
|
||||||
try:
|
try:
|
||||||
from .blender_dataset_map_provider import ( # noqa: F401
|
|
||||||
BlenderDatasetMapProvider,
|
|
||||||
)
|
|
||||||
from .json_index_dataset_map_provider import ( # noqa: F401
|
from .json_index_dataset_map_provider import ( # noqa: F401
|
||||||
JsonIndexDatasetMapProvider,
|
JsonIndexDatasetMapProvider,
|
||||||
)
|
)
|
||||||
from .json_index_dataset_map_provider_v2 import ( # noqa: F401
|
from .json_index_dataset_map_provider_v2 import ( # noqa: F401
|
||||||
JsonIndexDatasetMapProviderV2,
|
JsonIndexDatasetMapProviderV2,
|
||||||
)
|
)
|
||||||
from .llff_dataset_map_provider import LlffDatasetMapProvider # noqa: F401
|
|
||||||
from .rendered_mesh_dataset_map_provider import ( # noqa: F401
|
from .rendered_mesh_dataset_map_provider import ( # noqa: F401
|
||||||
RenderedMeshDatasetMapProvider,
|
RenderedMeshDatasetMapProvider,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1,68 +0,0 @@
|
|||||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
||||||
# All rights reserved.
|
|
||||||
#
|
|
||||||
# This source code is licensed under the BSD-style license found in the
|
|
||||||
# LICENSE file in the root directory of this source tree.
|
|
||||||
|
|
||||||
# pyre-unsafe
|
|
||||||
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from pytorch3d.implicitron.tools.config import registry
|
|
||||||
|
|
||||||
from .load_llff import load_llff_data
|
|
||||||
from .single_sequence_dataset import (
|
|
||||||
_interpret_blender_cameras,
|
|
||||||
SingleSceneDatasetMapProviderBase,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@registry.register
|
|
||||||
class LlffDatasetMapProvider(SingleSceneDatasetMapProviderBase):
|
|
||||||
"""
|
|
||||||
Provides data for one scene from the LLFF dataset.
|
|
||||||
|
|
||||||
Members:
|
|
||||||
base_dir: directory holding the data for the scene.
|
|
||||||
object_name: The name of the scene (e.g. "fern"). This is just used as a label.
|
|
||||||
It will typically be equal to the name of the directory self.base_dir.
|
|
||||||
path_manager_factory: Creates path manager which may be used for
|
|
||||||
interpreting paths.
|
|
||||||
n_known_frames_for_test: If set, training frames are included in the val
|
|
||||||
and test datasets, and this many random training frames are added to
|
|
||||||
each test batch. If not set, test batches each contain just a single
|
|
||||||
testing frame.
|
|
||||||
downscale_factor: determines image sizes.
|
|
||||||
"""
|
|
||||||
|
|
||||||
downscale_factor: int = 4
|
|
||||||
|
|
||||||
def _load_data(self) -> None:
|
|
||||||
path_manager = self.path_manager_factory.get()
|
|
||||||
images, poses, _ = load_llff_data(
|
|
||||||
self.base_dir, factor=self.downscale_factor, path_manager=path_manager
|
|
||||||
)
|
|
||||||
hwf = poses[0, :3, -1]
|
|
||||||
poses = poses[:, :3, :4]
|
|
||||||
|
|
||||||
llffhold = 8
|
|
||||||
i_test = np.arange(images.shape[0])[::llffhold]
|
|
||||||
i_test_index = set(i_test.tolist())
|
|
||||||
i_train = np.array(
|
|
||||||
[i for i in np.arange(images.shape[0]) if i not in i_test_index]
|
|
||||||
)
|
|
||||||
i_split = (i_train, i_test, i_test)
|
|
||||||
H, W, focal = hwf
|
|
||||||
focal_ndc = 2 * focal / min(H, W)
|
|
||||||
images = torch.from_numpy(images).permute(0, 3, 1, 2)
|
|
||||||
poses = torch.from_numpy(poses)
|
|
||||||
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.poses = _interpret_blender_cameras(poses, focal_ndc)
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.images = images
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.fg_probabilities = None
|
|
||||||
# pyre-ignore[16]
|
|
||||||
self.i_split = i_split
|
|
||||||
@@ -1,143 +0,0 @@
|
|||||||
# @lint-ignore-every LICENSELINT
|
|
||||||
# Adapted from https://github.com/bmild/nerf/blob/master/load_blender.py
|
|
||||||
# Copyright (c) 2020 bmild
|
|
||||||
|
|
||||||
# pyre-unsafe
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
|
|
||||||
def translate_by_t_along_z(t):
|
|
||||||
tform = np.eye(4).astype(np.float32)
|
|
||||||
tform[2][3] = t
|
|
||||||
return tform
|
|
||||||
|
|
||||||
|
|
||||||
def rotate_by_phi_along_x(phi):
|
|
||||||
tform = np.eye(4).astype(np.float32)
|
|
||||||
tform[1, 1] = tform[2, 2] = np.cos(phi)
|
|
||||||
tform[1, 2] = -np.sin(phi)
|
|
||||||
tform[2, 1] = -tform[1, 2]
|
|
||||||
return tform
|
|
||||||
|
|
||||||
|
|
||||||
def rotate_by_theta_along_y(theta):
|
|
||||||
tform = np.eye(4).astype(np.float32)
|
|
||||||
tform[0, 0] = tform[2, 2] = np.cos(theta)
|
|
||||||
tform[0, 2] = -np.sin(theta)
|
|
||||||
tform[2, 0] = -tform[0, 2]
|
|
||||||
return tform
|
|
||||||
|
|
||||||
|
|
||||||
def pose_spherical(theta, phi, radius):
|
|
||||||
c2w = translate_by_t_along_z(radius)
|
|
||||||
c2w = rotate_by_phi_along_x(phi / 180.0 * np.pi) @ c2w
|
|
||||||
c2w = rotate_by_theta_along_y(theta / 180 * np.pi) @ c2w
|
|
||||||
c2w = np.array([[-1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]) @ c2w
|
|
||||||
return c2w
|
|
||||||
|
|
||||||
|
|
||||||
def _local_path(path_manager, path):
|
|
||||||
if path_manager is None:
|
|
||||||
return path
|
|
||||||
return path_manager.get_local_path(path)
|
|
||||||
|
|
||||||
|
|
||||||
def load_blender_data(
|
|
||||||
basedir,
|
|
||||||
half_res=False,
|
|
||||||
testskip=1,
|
|
||||||
debug=False,
|
|
||||||
path_manager=None,
|
|
||||||
focal_length_in_screen_space=False,
|
|
||||||
):
|
|
||||||
splits = ["train", "val", "test"]
|
|
||||||
metas = {}
|
|
||||||
for s in splits:
|
|
||||||
path = os.path.join(basedir, f"transforms_{s}.json")
|
|
||||||
with open(_local_path(path_manager, path)) as fp:
|
|
||||||
metas[s] = json.load(fp)
|
|
||||||
|
|
||||||
all_imgs = []
|
|
||||||
all_poses = []
|
|
||||||
counts = [0]
|
|
||||||
for s in splits:
|
|
||||||
meta = metas[s]
|
|
||||||
imgs = []
|
|
||||||
poses = []
|
|
||||||
if s == "train" or testskip == 0:
|
|
||||||
skip = 1
|
|
||||||
else:
|
|
||||||
skip = testskip
|
|
||||||
|
|
||||||
for frame in meta["frames"][::skip]:
|
|
||||||
fname = os.path.join(basedir, frame["file_path"] + ".png")
|
|
||||||
imgs.append(np.array(Image.open(_local_path(path_manager, fname))))
|
|
||||||
poses.append(np.array(frame["transform_matrix"]))
|
|
||||||
imgs = (np.array(imgs) / 255.0).astype(np.float32)
|
|
||||||
poses = np.array(poses).astype(np.float32)
|
|
||||||
counts.append(counts[-1] + imgs.shape[0])
|
|
||||||
all_imgs.append(imgs)
|
|
||||||
all_poses.append(poses)
|
|
||||||
|
|
||||||
i_split = [np.arange(counts[i], counts[i + 1]) for i in range(3)]
|
|
||||||
|
|
||||||
imgs = np.concatenate(all_imgs, 0)
|
|
||||||
poses = np.concatenate(all_poses, 0)
|
|
||||||
|
|
||||||
H, W = imgs[0].shape[:2]
|
|
||||||
camera_angle_x = float(meta["camera_angle_x"])
|
|
||||||
if focal_length_in_screen_space:
|
|
||||||
focal = 0.5 * W / np.tan(0.5 * camera_angle_x)
|
|
||||||
else:
|
|
||||||
focal = 1 / np.tan(0.5 * camera_angle_x)
|
|
||||||
|
|
||||||
render_poses = torch.stack(
|
|
||||||
[
|
|
||||||
torch.from_numpy(pose_spherical(angle, -30.0, 4.0))
|
|
||||||
for angle in np.linspace(-180, 180, 40 + 1)[:-1]
|
|
||||||
],
|
|
||||||
0,
|
|
||||||
)
|
|
||||||
|
|
||||||
# In debug mode, return extremely tiny images
|
|
||||||
if debug:
|
|
||||||
import cv2
|
|
||||||
|
|
||||||
H = H // 32
|
|
||||||
W = W // 32
|
|
||||||
if focal_length_in_screen_space:
|
|
||||||
focal = focal / 32.0
|
|
||||||
imgs = [
|
|
||||||
torch.from_numpy(
|
|
||||||
cv2.resize(imgs[i], dsize=(25, 25), interpolation=cv2.INTER_AREA)
|
|
||||||
)
|
|
||||||
for i in range(imgs.shape[0])
|
|
||||||
]
|
|
||||||
imgs = torch.stack(imgs, 0)
|
|
||||||
poses = torch.from_numpy(poses)
|
|
||||||
return imgs, poses, render_poses, [H, W, focal], i_split
|
|
||||||
|
|
||||||
if half_res:
|
|
||||||
import cv2
|
|
||||||
|
|
||||||
# TODO: resize images using INTER_AREA (cv2)
|
|
||||||
H = H // 2
|
|
||||||
W = W // 2
|
|
||||||
if focal_length_in_screen_space:
|
|
||||||
focal = focal / 2.0
|
|
||||||
imgs = [
|
|
||||||
torch.from_numpy(
|
|
||||||
cv2.resize(imgs[i], dsize=(400, 400), interpolation=cv2.INTER_AREA)
|
|
||||||
)
|
|
||||||
for i in range(imgs.shape[0])
|
|
||||||
]
|
|
||||||
imgs = torch.stack(imgs, 0)
|
|
||||||
|
|
||||||
poses = torch.from_numpy(poses)
|
|
||||||
|
|
||||||
return imgs, poses, render_poses, [H, W, focal], i_split
|
|
||||||
@@ -1,335 +0,0 @@
|
|||||||
# @lint-ignore-every LICENSELINT
|
|
||||||
# Adapted from https://github.com/bmild/nerf/blob/master/load_llff.py
|
|
||||||
# Copyright (c) 2020 bmild
|
|
||||||
|
|
||||||
# pyre-unsafe
|
|
||||||
import logging
|
|
||||||
import os
|
|
||||||
import warnings
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
|
|
||||||
# Slightly modified version of LLFF data loading code
|
|
||||||
# see https://github.com/Fyusion/LLFF for original
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
def _minify(basedir, path_manager, factors=(), resolutions=()):
|
|
||||||
needtoload = False
|
|
||||||
for r in factors:
|
|
||||||
imgdir = os.path.join(basedir, "images_{}".format(r))
|
|
||||||
if not _exists(path_manager, imgdir):
|
|
||||||
needtoload = True
|
|
||||||
for r in resolutions:
|
|
||||||
imgdir = os.path.join(basedir, "images_{}x{}".format(r[1], r[0]))
|
|
||||||
if not _exists(path_manager, imgdir):
|
|
||||||
needtoload = True
|
|
||||||
if not needtoload:
|
|
||||||
return
|
|
||||||
assert path_manager is None
|
|
||||||
|
|
||||||
from subprocess import check_output
|
|
||||||
|
|
||||||
imgdir = os.path.join(basedir, "images")
|
|
||||||
imgs = [os.path.join(imgdir, f) for f in sorted(_ls(path_manager, imgdir))]
|
|
||||||
imgs = [f for f in imgs if f.endswith("JPG", "jpg", "png", "jpeg", "PNG")]
|
|
||||||
imgdir_orig = imgdir
|
|
||||||
|
|
||||||
wd = os.getcwd()
|
|
||||||
|
|
||||||
for r in factors + resolutions:
|
|
||||||
if isinstance(r, int):
|
|
||||||
name = "images_{}".format(r)
|
|
||||||
resizearg = "{}%".format(100.0 / r)
|
|
||||||
else:
|
|
||||||
name = "images_{}x{}".format(r[1], r[0])
|
|
||||||
resizearg = "{}x{}".format(r[1], r[0])
|
|
||||||
imgdir = os.path.join(basedir, name)
|
|
||||||
if os.path.exists(imgdir):
|
|
||||||
continue
|
|
||||||
|
|
||||||
logger.info(f"Minifying {r}, {basedir}")
|
|
||||||
|
|
||||||
os.makedirs(imgdir)
|
|
||||||
check_output("cp {}/* {}".format(imgdir_orig, imgdir), shell=True)
|
|
||||||
|
|
||||||
ext = imgs[0].split(".")[-1]
|
|
||||||
args = " ".join(
|
|
||||||
["mogrify", "-resize", resizearg, "-format", "png", "*.{}".format(ext)]
|
|
||||||
)
|
|
||||||
logger.info(args)
|
|
||||||
os.chdir(imgdir)
|
|
||||||
check_output(args, shell=True)
|
|
||||||
os.chdir(wd)
|
|
||||||
|
|
||||||
if ext != "png":
|
|
||||||
check_output("rm {}/*.{}".format(imgdir, ext), shell=True)
|
|
||||||
logger.info("Removed duplicates")
|
|
||||||
logger.info("Done")
|
|
||||||
|
|
||||||
|
|
||||||
def _load_data(
|
|
||||||
basedir, factor=None, width=None, height=None, load_imgs=True, path_manager=None
|
|
||||||
):
|
|
||||||
poses_arr = np.load(
|
|
||||||
_local_path(path_manager, os.path.join(basedir, "poses_bounds.npy"))
|
|
||||||
)
|
|
||||||
poses = poses_arr[:, :-2].reshape([-1, 3, 5]).transpose([1, 2, 0])
|
|
||||||
bds = poses_arr[:, -2:].transpose([1, 0])
|
|
||||||
|
|
||||||
img0 = [
|
|
||||||
os.path.join(basedir, "images", f)
|
|
||||||
for f in sorted(_ls(path_manager, os.path.join(basedir, "images")))
|
|
||||||
if f.endswith("JPG") or f.endswith("jpg") or f.endswith("png")
|
|
||||||
][0]
|
|
||||||
|
|
||||||
def imread(f):
|
|
||||||
return np.array(Image.open(f))
|
|
||||||
|
|
||||||
sh = imread(_local_path(path_manager, img0)).shape
|
|
||||||
|
|
||||||
sfx = ""
|
|
||||||
|
|
||||||
if factor is not None:
|
|
||||||
sfx = "_{}".format(factor)
|
|
||||||
_minify(basedir, path_manager, factors=[factor])
|
|
||||||
factor = factor
|
|
||||||
elif height is not None:
|
|
||||||
factor = sh[0] / float(height)
|
|
||||||
width = int(sh[1] / factor)
|
|
||||||
_minify(basedir, path_manager, resolutions=[[height, width]])
|
|
||||||
sfx = "_{}x{}".format(width, height)
|
|
||||||
elif width is not None:
|
|
||||||
factor = sh[1] / float(width)
|
|
||||||
height = int(sh[0] / factor)
|
|
||||||
_minify(basedir, path_manager, resolutions=[[height, width]])
|
|
||||||
sfx = "_{}x{}".format(width, height)
|
|
||||||
else:
|
|
||||||
factor = 1
|
|
||||||
|
|
||||||
imgdir = os.path.join(basedir, "images" + sfx)
|
|
||||||
if not _exists(path_manager, imgdir):
|
|
||||||
raise ValueError(f"{imgdir} does not exist, returning")
|
|
||||||
|
|
||||||
imgfiles = [
|
|
||||||
_local_path(path_manager, os.path.join(imgdir, f))
|
|
||||||
for f in sorted(_ls(path_manager, imgdir))
|
|
||||||
if f.endswith("JPG") or f.endswith("jpg") or f.endswith("png")
|
|
||||||
]
|
|
||||||
if poses.shape[-1] != len(imgfiles):
|
|
||||||
raise ValueError(
|
|
||||||
"Mismatch between imgs {} and poses {} !!!!".format(
|
|
||||||
len(imgfiles), poses.shape[-1]
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
sh = imread(imgfiles[0]).shape
|
|
||||||
poses[:2, 4, :] = np.array(sh[:2]).reshape([2, 1])
|
|
||||||
poses[2, 4, :] = poses[2, 4, :] * 1.0 / factor
|
|
||||||
|
|
||||||
if not load_imgs:
|
|
||||||
return poses, bds
|
|
||||||
|
|
||||||
imgs = imgs = [imread(f)[..., :3] / 255.0 for f in imgfiles]
|
|
||||||
imgs = np.stack(imgs, -1)
|
|
||||||
|
|
||||||
logger.info(f"Loaded image data, shape {imgs.shape}")
|
|
||||||
return poses, bds, imgs
|
|
||||||
|
|
||||||
|
|
||||||
def normalize(x):
|
|
||||||
denom = np.linalg.norm(x)
|
|
||||||
if denom < 0.001:
|
|
||||||
warnings.warn("unsafe normalize()")
|
|
||||||
return x / denom
|
|
||||||
|
|
||||||
|
|
||||||
def viewmatrix(z, up, pos):
|
|
||||||
vec2 = normalize(z)
|
|
||||||
vec1_avg = up
|
|
||||||
vec0 = normalize(np.cross(vec1_avg, vec2))
|
|
||||||
vec1 = normalize(np.cross(vec2, vec0))
|
|
||||||
m = np.stack([vec0, vec1, vec2, pos], 1)
|
|
||||||
return m
|
|
||||||
|
|
||||||
|
|
||||||
def ptstocam(pts, c2w):
|
|
||||||
tt = np.matmul(c2w[:3, :3].T, (pts - c2w[:3, 3])[..., np.newaxis])[..., 0]
|
|
||||||
return tt
|
|
||||||
|
|
||||||
|
|
||||||
def poses_avg(poses):
|
|
||||||
hwf = poses[0, :3, -1:]
|
|
||||||
|
|
||||||
center = poses[:, :3, 3].mean(0)
|
|
||||||
vec2 = normalize(poses[:, :3, 2].sum(0))
|
|
||||||
up = poses[:, :3, 1].sum(0)
|
|
||||||
c2w = np.concatenate([viewmatrix(vec2, up, center), hwf], 1)
|
|
||||||
|
|
||||||
return c2w
|
|
||||||
|
|
||||||
|
|
||||||
def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, rots, N):
|
|
||||||
render_poses = []
|
|
||||||
rads = np.array(list(rads) + [1.0])
|
|
||||||
hwf = c2w[:, 4:5]
|
|
||||||
|
|
||||||
for theta in np.linspace(0.0, 2.0 * np.pi * rots, N + 1)[:-1]:
|
|
||||||
c = np.dot(
|
|
||||||
c2w[:3, :4],
|
|
||||||
np.array([np.cos(theta), -np.sin(theta), -np.sin(theta * zrate), 1.0])
|
|
||||||
* rads,
|
|
||||||
)
|
|
||||||
z = normalize(c - np.dot(c2w[:3, :4], np.array([0, 0, -focal, 1.0])))
|
|
||||||
render_poses.append(np.concatenate([viewmatrix(z, up, c), hwf], 1))
|
|
||||||
return render_poses
|
|
||||||
|
|
||||||
|
|
||||||
def recenter_poses(poses):
|
|
||||||
poses_ = poses + 0
|
|
||||||
bottom = np.reshape([0, 0, 0, 1.0], [1, 4])
|
|
||||||
c2w = poses_avg(poses)
|
|
||||||
c2w = np.concatenate([c2w[:3, :4], bottom], -2)
|
|
||||||
bottom = np.tile(np.reshape(bottom, [1, 1, 4]), [poses.shape[0], 1, 1])
|
|
||||||
poses = np.concatenate([poses[:, :3, :4], bottom], -2)
|
|
||||||
|
|
||||||
poses = np.linalg.inv(c2w) @ poses
|
|
||||||
poses_[:, :3, :4] = poses[:, :3, :4]
|
|
||||||
poses = poses_
|
|
||||||
return poses
|
|
||||||
|
|
||||||
|
|
||||||
def spherify_poses(poses, bds):
|
|
||||||
def add_row_to_homogenize_transform(p):
|
|
||||||
r"""Add the last row to homogenize 3 x 4 transformation matrices."""
|
|
||||||
return np.concatenate(
|
|
||||||
[p, np.tile(np.reshape(np.eye(4)[-1, :], [1, 1, 4]), [p.shape[0], 1, 1])], 1
|
|
||||||
)
|
|
||||||
|
|
||||||
# p34_to_44 = lambda p: np.concatenate(
|
|
||||||
# [p, np.tile(np.reshape(np.eye(4)[-1, :], [1, 1, 4]), [p.shape[0], 1, 1])], 1
|
|
||||||
# )
|
|
||||||
|
|
||||||
p34_to_44 = add_row_to_homogenize_transform
|
|
||||||
|
|
||||||
rays_d = poses[:, :3, 2:3]
|
|
||||||
rays_o = poses[:, :3, 3:4]
|
|
||||||
|
|
||||||
def min_line_dist(rays_o, rays_d):
|
|
||||||
A_i = np.eye(3) - rays_d * np.transpose(rays_d, [0, 2, 1])
|
|
||||||
b_i = -A_i @ rays_o
|
|
||||||
pt_mindist = np.squeeze(
|
|
||||||
-np.linalg.inv((np.transpose(A_i, [0, 2, 1]) @ A_i).mean(0)) @ (b_i).mean(0)
|
|
||||||
)
|
|
||||||
return pt_mindist
|
|
||||||
|
|
||||||
pt_mindist = min_line_dist(rays_o, rays_d)
|
|
||||||
|
|
||||||
center = pt_mindist
|
|
||||||
up = (poses[:, :3, 3] - center).mean(0)
|
|
||||||
|
|
||||||
vec0 = normalize(up)
|
|
||||||
vec1 = normalize(np.cross([0.1, 0.2, 0.3], vec0))
|
|
||||||
vec2 = normalize(np.cross(vec0, vec1))
|
|
||||||
pos = center
|
|
||||||
c2w = np.stack([vec1, vec2, vec0, pos], 1)
|
|
||||||
|
|
||||||
poses_reset = np.linalg.inv(p34_to_44(c2w[None])) @ p34_to_44(poses[:, :3, :4])
|
|
||||||
|
|
||||||
rad = np.sqrt(np.mean(np.sum(np.square(poses_reset[:, :3, 3]), -1)))
|
|
||||||
|
|
||||||
sc = 1.0 / rad
|
|
||||||
poses_reset[:, :3, 3] *= sc
|
|
||||||
bds *= sc
|
|
||||||
rad *= sc
|
|
||||||
|
|
||||||
centroid = np.mean(poses_reset[:, :3, 3], 0)
|
|
||||||
zh = centroid[2]
|
|
||||||
radcircle = np.sqrt(rad**2 - zh**2)
|
|
||||||
new_poses = []
|
|
||||||
|
|
||||||
for th in np.linspace(0.0, 2.0 * np.pi, 120):
|
|
||||||
camorigin = np.array([radcircle * np.cos(th), radcircle * np.sin(th), zh])
|
|
||||||
up = np.array([0, 0, -1.0])
|
|
||||||
|
|
||||||
vec2 = normalize(camorigin)
|
|
||||||
vec0 = normalize(np.cross(vec2, up))
|
|
||||||
vec1 = normalize(np.cross(vec2, vec0))
|
|
||||||
pos = camorigin
|
|
||||||
p = np.stack([vec0, vec1, vec2, pos], 1)
|
|
||||||
|
|
||||||
new_poses.append(p)
|
|
||||||
|
|
||||||
new_poses = np.stack(new_poses, 0)
|
|
||||||
|
|
||||||
new_poses = np.concatenate(
|
|
||||||
[new_poses, np.broadcast_to(poses[0, :3, -1:], new_poses[:, :3, -1:].shape)], -1
|
|
||||||
)
|
|
||||||
poses_reset = np.concatenate(
|
|
||||||
[
|
|
||||||
poses_reset[:, :3, :4],
|
|
||||||
np.broadcast_to(poses[0, :3, -1:], poses_reset[:, :3, -1:].shape),
|
|
||||||
],
|
|
||||||
-1,
|
|
||||||
)
|
|
||||||
|
|
||||||
return poses_reset, new_poses, bds
|
|
||||||
|
|
||||||
|
|
||||||
def _local_path(path_manager, path):
|
|
||||||
if path_manager is None:
|
|
||||||
return path
|
|
||||||
return path_manager.get_local_path(path)
|
|
||||||
|
|
||||||
|
|
||||||
def _ls(path_manager, path):
|
|
||||||
if path_manager is None:
|
|
||||||
return os.listdir(path)
|
|
||||||
return path_manager.ls(path)
|
|
||||||
|
|
||||||
|
|
||||||
def _exists(path_manager, path):
|
|
||||||
if path_manager is None:
|
|
||||||
return os.path.exists(path)
|
|
||||||
return path_manager.exists(path)
|
|
||||||
|
|
||||||
|
|
||||||
def load_llff_data(
|
|
||||||
basedir,
|
|
||||||
factor=8,
|
|
||||||
recenter=True,
|
|
||||||
bd_factor=0.75,
|
|
||||||
spherify=False,
|
|
||||||
path_zflat=False,
|
|
||||||
path_manager=None,
|
|
||||||
):
|
|
||||||
poses, bds, imgs = _load_data(
|
|
||||||
basedir, factor=factor, path_manager=path_manager
|
|
||||||
) # factor=8 downsamples original imgs by 8x
|
|
||||||
logger.info(f"Loaded {basedir}, {bds.min()}, {bds.max()}")
|
|
||||||
|
|
||||||
# Correct rotation matrix ordering and move variable dim to axis 0
|
|
||||||
poses = np.concatenate([poses[:, 1:2, :], -poses[:, 0:1, :], poses[:, 2:, :]], 1)
|
|
||||||
poses = np.moveaxis(poses, -1, 0).astype(np.float32)
|
|
||||||
imgs = np.moveaxis(imgs, -1, 0).astype(np.float32)
|
|
||||||
images = imgs
|
|
||||||
bds = np.moveaxis(bds, -1, 0).astype(np.float32)
|
|
||||||
|
|
||||||
# Rescale if bd_factor is provided
|
|
||||||
sc = 1.0 if bd_factor is None else 1.0 / (bds.min() * bd_factor)
|
|
||||||
poses[:, :3, 3] *= sc
|
|
||||||
bds *= sc
|
|
||||||
|
|
||||||
if recenter:
|
|
||||||
poses = recenter_poses(poses)
|
|
||||||
|
|
||||||
if spherify:
|
|
||||||
poses, render_poses, bds = spherify_poses(poses, bds)
|
|
||||||
|
|
||||||
images = images.astype(np.float32)
|
|
||||||
poses = poses.astype(np.float32)
|
|
||||||
|
|
||||||
return images, poses, bds
|
|
||||||
@@ -85,7 +85,7 @@ class SingleSceneDataset(DatasetBase, Configurable):
|
|||||||
|
|
||||||
class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
|
class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
|
||||||
"""
|
"""
|
||||||
Base for provider of data for one scene from LLFF or blender datasets.
|
Base for provider of data for one scene.
|
||||||
|
|
||||||
Members:
|
Members:
|
||||||
base_dir: directory holding the data for the scene.
|
base_dir: directory holding the data for the scene.
|
||||||
@@ -171,40 +171,3 @@ class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
|
|||||||
# pyre-ignore[16]
|
# pyre-ignore[16]
|
||||||
cameras = [self.poses[i] for i in self.i_split[0]]
|
cameras = [self.poses[i] for i in self.i_split[0]]
|
||||||
return join_cameras_as_batch(cameras)
|
return join_cameras_as_batch(cameras)
|
||||||
|
|
||||||
|
|
||||||
def _interpret_blender_cameras(
|
|
||||||
poses: torch.Tensor, focal: float
|
|
||||||
) -> List[PerspectiveCameras]:
|
|
||||||
"""
|
|
||||||
Convert 4x4 matrices representing cameras in blender format
|
|
||||||
to PyTorch3D format.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
poses: N x 3 x 4 camera matrices
|
|
||||||
focal: ndc space focal length
|
|
||||||
"""
|
|
||||||
pose_target_cameras = []
|
|
||||||
for pose_target in poses:
|
|
||||||
pose_target = pose_target[:3, :4]
|
|
||||||
mtx = torch.eye(4, dtype=pose_target.dtype)
|
|
||||||
mtx[:3, :3] = pose_target[:3, :3].t()
|
|
||||||
mtx[3, :3] = pose_target[:, 3]
|
|
||||||
mtx = mtx.inverse()
|
|
||||||
|
|
||||||
# flip the XZ coordinates.
|
|
||||||
mtx[:, [0, 2]] *= -1.0
|
|
||||||
|
|
||||||
Rpt3, Tpt3 = mtx[:, :3].split([3, 1], dim=0)
|
|
||||||
|
|
||||||
focal_length_pt3 = torch.FloatTensor([[focal, focal]])
|
|
||||||
principal_point_pt3 = torch.FloatTensor([[0.0, 0.0]])
|
|
||||||
|
|
||||||
cameras = PerspectiveCameras(
|
|
||||||
focal_length=focal_length_pt3,
|
|
||||||
principal_point=principal_point_pt3,
|
|
||||||
R=Rpt3[None],
|
|
||||||
T=Tpt3,
|
|
||||||
)
|
|
||||||
pose_target_cameras.append(cameras)
|
|
||||||
return pose_target_cameras
|
|
||||||
|
|||||||
@@ -1,12 +1,5 @@
|
|||||||
dataset_map_provider_class_type: ???
|
dataset_map_provider_class_type: ???
|
||||||
data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
|
data_loader_map_provider_class_type: SequenceDataLoaderMapProvider
|
||||||
dataset_map_provider_BlenderDatasetMapProvider_args:
|
|
||||||
base_dir: ???
|
|
||||||
object_name: ???
|
|
||||||
path_manager_factory_class_type: PathManagerFactory
|
|
||||||
n_known_frames_for_test: null
|
|
||||||
path_manager_factory_PathManagerFactory_args:
|
|
||||||
silence_logs: true
|
|
||||||
dataset_map_provider_JsonIndexDatasetMapProvider_args:
|
dataset_map_provider_JsonIndexDatasetMapProvider_args:
|
||||||
category: ???
|
category: ???
|
||||||
task_str: singlesequence
|
task_str: singlesequence
|
||||||
@@ -78,14 +71,6 @@ dataset_map_provider_JsonIndexDatasetMapProviderV2_args:
|
|||||||
sort_frames: false
|
sort_frames: false
|
||||||
path_manager_factory_PathManagerFactory_args:
|
path_manager_factory_PathManagerFactory_args:
|
||||||
silence_logs: true
|
silence_logs: true
|
||||||
dataset_map_provider_LlffDatasetMapProvider_args:
|
|
||||||
base_dir: ???
|
|
||||||
object_name: ???
|
|
||||||
path_manager_factory_class_type: PathManagerFactory
|
|
||||||
n_known_frames_for_test: null
|
|
||||||
path_manager_factory_PathManagerFactory_args:
|
|
||||||
silence_logs: true
|
|
||||||
downscale_factor: 4
|
|
||||||
dataset_map_provider_RenderedMeshDatasetMapProvider_args:
|
dataset_map_provider_RenderedMeshDatasetMapProvider_args:
|
||||||
num_views: 40
|
num_views: 40
|
||||||
data_file: null
|
data_file: null
|
||||||
|
|||||||
@@ -1,158 +0,0 @@
|
|||||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
||||||
# All rights reserved.
|
|
||||||
#
|
|
||||||
# This source code is licensed under the BSD-style license found in the
|
|
||||||
# LICENSE file in the root directory of this source tree.
|
|
||||||
|
|
||||||
import os
|
|
||||||
import unittest
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from pytorch3d.implicitron.dataset.blender_dataset_map_provider import (
|
|
||||||
BlenderDatasetMapProvider,
|
|
||||||
)
|
|
||||||
from pytorch3d.implicitron.dataset.data_source import ImplicitronDataSource
|
|
||||||
from pytorch3d.implicitron.dataset.dataset_base import FrameData
|
|
||||||
from pytorch3d.implicitron.dataset.llff_dataset_map_provider import (
|
|
||||||
LlffDatasetMapProvider,
|
|
||||||
)
|
|
||||||
from pytorch3d.implicitron.tools.config import expand_args_fields, get_default_args
|
|
||||||
from pytorch3d.renderer import PerspectiveCameras
|
|
||||||
from tests.common_testing import TestCaseMixin
|
|
||||||
|
|
||||||
|
|
||||||
# These tests are only run internally, where the data is available.
|
|
||||||
internal = os.environ.get("FB_TEST", False)
|
|
||||||
inside_re_worker = os.environ.get("INSIDE_RE_WORKER", False)
|
|
||||||
|
|
||||||
|
|
||||||
@unittest.skipUnless(internal, "no data")
|
|
||||||
class TestDataLlff(TestCaseMixin, unittest.TestCase):
|
|
||||||
def test_synthetic(self):
|
|
||||||
if inside_re_worker:
|
|
||||||
return
|
|
||||||
expand_args_fields(BlenderDatasetMapProvider)
|
|
||||||
|
|
||||||
provider = BlenderDatasetMapProvider(
|
|
||||||
base_dir="manifold://co3d/tree/nerf_data/nerf_synthetic/lego",
|
|
||||||
object_name="lego",
|
|
||||||
)
|
|
||||||
dataset_map = provider.get_dataset_map()
|
|
||||||
known_matrix = torch.zeros(1, 4, 4)
|
|
||||||
known_matrix[0, 0, 0] = 2.7778
|
|
||||||
known_matrix[0, 1, 1] = 2.7778
|
|
||||||
known_matrix[0, 2, 3] = 1
|
|
||||||
known_matrix[0, 3, 2] = 1
|
|
||||||
|
|
||||||
for name, length in [("train", 100), ("val", 100), ("test", 200)]:
|
|
||||||
dataset = getattr(dataset_map, name)
|
|
||||||
self.assertEqual(len(dataset), length)
|
|
||||||
# try getting a value
|
|
||||||
value = dataset[0]
|
|
||||||
self.assertEqual(value.image_rgb.shape, (3, 800, 800))
|
|
||||||
self.assertEqual(value.fg_probability.shape, (1, 800, 800))
|
|
||||||
# corner of image is background
|
|
||||||
self.assertEqual(value.fg_probability[0, 0, 0], 0)
|
|
||||||
self.assertEqual(value.fg_probability.max(), 1.0)
|
|
||||||
self.assertIsInstance(value.camera, PerspectiveCameras)
|
|
||||||
self.assertEqual(len(value.camera), 1)
|
|
||||||
self.assertIsNone(value.camera.K)
|
|
||||||
matrix = value.camera.get_projection_transform().get_matrix()
|
|
||||||
self.assertClose(matrix, known_matrix, atol=1e-4)
|
|
||||||
self.assertIsInstance(value, FrameData)
|
|
||||||
|
|
||||||
def test_llff(self):
|
|
||||||
if inside_re_worker:
|
|
||||||
return
|
|
||||||
expand_args_fields(LlffDatasetMapProvider)
|
|
||||||
|
|
||||||
provider = LlffDatasetMapProvider(
|
|
||||||
base_dir="manifold://co3d/tree/nerf_data/nerf_llff_data/fern",
|
|
||||||
object_name="fern",
|
|
||||||
downscale_factor=8,
|
|
||||||
)
|
|
||||||
dataset_map = provider.get_dataset_map()
|
|
||||||
known_matrix = torch.zeros(1, 4, 4)
|
|
||||||
known_matrix[0, 0, 0] = 2.1564
|
|
||||||
known_matrix[0, 1, 1] = 2.1564
|
|
||||||
known_matrix[0, 2, 3] = 1
|
|
||||||
known_matrix[0, 3, 2] = 1
|
|
||||||
|
|
||||||
for name, length, frame_type in [
|
|
||||||
("train", 17, "known"),
|
|
||||||
("test", 3, "unseen"),
|
|
||||||
("val", 3, "unseen"),
|
|
||||||
]:
|
|
||||||
dataset = getattr(dataset_map, name)
|
|
||||||
self.assertEqual(len(dataset), length)
|
|
||||||
# try getting a value
|
|
||||||
value = dataset[0]
|
|
||||||
self.assertIsInstance(value, FrameData)
|
|
||||||
self.assertEqual(value.frame_type, frame_type)
|
|
||||||
self.assertEqual(value.image_rgb.shape, (3, 378, 504))
|
|
||||||
self.assertIsInstance(value.camera, PerspectiveCameras)
|
|
||||||
self.assertEqual(len(value.camera), 1)
|
|
||||||
self.assertIsNone(value.camera.K)
|
|
||||||
matrix = value.camera.get_projection_transform().get_matrix()
|
|
||||||
self.assertClose(matrix, known_matrix, atol=1e-4)
|
|
||||||
|
|
||||||
self.assertEqual(len(dataset_map.test.get_eval_batches()), 3)
|
|
||||||
for batch in dataset_map.test.get_eval_batches():
|
|
||||||
self.assertEqual(len(batch), 1)
|
|
||||||
self.assertEqual(dataset_map.test[batch[0]].frame_type, "unseen")
|
|
||||||
|
|
||||||
def test_include_known_frames(self):
|
|
||||||
if inside_re_worker:
|
|
||||||
return
|
|
||||||
expand_args_fields(LlffDatasetMapProvider)
|
|
||||||
|
|
||||||
provider = LlffDatasetMapProvider(
|
|
||||||
base_dir="manifold://co3d/tree/nerf_data/nerf_llff_data/fern",
|
|
||||||
object_name="fern",
|
|
||||||
n_known_frames_for_test=2,
|
|
||||||
)
|
|
||||||
dataset_map = provider.get_dataset_map()
|
|
||||||
|
|
||||||
for name, types in [
|
|
||||||
("train", ["known"] * 17),
|
|
||||||
("val", ["unseen"] * 3 + ["known"] * 17),
|
|
||||||
("test", ["unseen"] * 3 + ["known"] * 17),
|
|
||||||
]:
|
|
||||||
dataset = getattr(dataset_map, name)
|
|
||||||
self.assertEqual(len(dataset), len(types))
|
|
||||||
for i, frame_type in enumerate(types):
|
|
||||||
value = dataset[i]
|
|
||||||
self.assertEqual(value.frame_type, frame_type)
|
|
||||||
self.assertIsNone(value.fg_probability)
|
|
||||||
|
|
||||||
self.assertEqual(len(dataset_map.test.get_eval_batches()), 3)
|
|
||||||
for batch in dataset_map.test.get_eval_batches():
|
|
||||||
self.assertEqual(len(batch), 3)
|
|
||||||
self.assertEqual(dataset_map.test[batch[0]].frame_type, "unseen")
|
|
||||||
for i in batch[1:]:
|
|
||||||
self.assertEqual(dataset_map.test[i].frame_type, "known")
|
|
||||||
|
|
||||||
def test_loaders(self):
|
|
||||||
if inside_re_worker:
|
|
||||||
return
|
|
||||||
args = get_default_args(ImplicitronDataSource)
|
|
||||||
args.dataset_map_provider_class_type = "BlenderDatasetMapProvider"
|
|
||||||
dataset_args = args.dataset_map_provider_BlenderDatasetMapProvider_args
|
|
||||||
dataset_args.object_name = "lego"
|
|
||||||
dataset_args.base_dir = "manifold://co3d/tree/nerf_data/nerf_synthetic/lego"
|
|
||||||
|
|
||||||
data_source = ImplicitronDataSource(**args)
|
|
||||||
_, data_loaders = data_source.get_datasets_and_dataloaders()
|
|
||||||
for i in data_loaders.train:
|
|
||||||
self.assertEqual(i.frame_type, ["known"])
|
|
||||||
self.assertEqual(i.image_rgb.shape, (1, 3, 800, 800))
|
|
||||||
for i in data_loaders.val:
|
|
||||||
self.assertEqual(i.frame_type, ["unseen"])
|
|
||||||
self.assertEqual(i.image_rgb.shape, (1, 3, 800, 800))
|
|
||||||
for i in data_loaders.test:
|
|
||||||
self.assertEqual(i.frame_type, ["unseen"])
|
|
||||||
self.assertEqual(i.image_rgb.shape, (1, 3, 800, 800))
|
|
||||||
|
|
||||||
cameras = data_source.all_train_cameras
|
|
||||||
self.assertIsInstance(cameras, PerspectiveCameras)
|
|
||||||
self.assertEqual(len(cameras), 100)
|
|
||||||
Reference in New Issue
Block a user