mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-03-19 11:36:03 +08:00
Compare commits
21 Commits
cbcae096a0
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b6a77ad7aa | ||
|
|
52164b8324 | ||
|
|
61cc79aa34 | ||
|
|
7a6157e38e | ||
|
|
d9839a95f2 | ||
|
|
7b5c78460a | ||
|
|
e3c80a4368 | ||
|
|
b9b5ea3428 | ||
|
|
0e435c297c | ||
|
|
d631b56fba | ||
|
|
3ba2030aa4 | ||
|
|
79a7fcf02b | ||
|
|
e43ed8c76e | ||
|
|
49f43402c6 | ||
|
|
90646d93ab | ||
|
|
eabb511410 | ||
|
|
e70188ebbc | ||
|
|
1bd911d534 | ||
|
|
3aadd19a2b | ||
|
|
42d66c1145 | ||
|
|
e9ed1cb178 |
@@ -3,11 +3,6 @@ pytorch3d.implicitron.dataset 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
|
||||
:members:
|
||||
:undoc-members:
|
||||
@@ -18,11 +13,6 @@ specific datasets
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. automodule:: pytorch3d.implicitron.dataset.llff_dataset_map_provider
|
||||
:members:
|
||||
:undoc-members:
|
||||
:show-inheritance:
|
||||
|
||||
.. automodule:: pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider
|
||||
: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:
|
||||
dataset_map_provider_class_type: ???
|
||||
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:
|
||||
category: ???
|
||||
task_str: singlesequence
|
||||
@@ -91,14 +84,6 @@ data_source_ImplicitronDataSource_args:
|
||||
sort_frames: false
|
||||
path_manager_factory_PathManagerFactory_args:
|
||||
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:
|
||||
num_views: 40
|
||||
data_file: null
|
||||
|
||||
@@ -82,10 +82,12 @@ class _SymEig3x3(nn.Module):
|
||||
q = inputs_trace / 3.0
|
||||
|
||||
# Calculate squared sum of elements outside the main diagonal / 2
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
p1 = ((inputs**2).sum(dim=(-1, -2)) - (inputs_diag**2).sum(-1)) / 2
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
p2 = ((inputs_diag - q[..., None]) ** 2).sum(dim=-1) + 2.0 * p1.clamp(self._eps)
|
||||
p1 = (
|
||||
torch.square(inputs).sum(dim=(-1, -2)) - torch.square(inputs_diag).sum(-1)
|
||||
) / 2
|
||||
p2 = torch.square(inputs_diag - q[..., None]).sum(dim=-1) + 2.0 * p1.clamp(
|
||||
self._eps
|
||||
)
|
||||
|
||||
p = torch.sqrt(p2 / 6.0)
|
||||
B = (inputs - q[..., None, None] * self._identity) / p[..., None, None]
|
||||
@@ -104,7 +106,9 @@ class _SymEig3x3(nn.Module):
|
||||
# Soft dispatch between the degenerate case (diagonal A) and general.
|
||||
# diag_soft_cond -> 1.0 when p1 < 6 * eps and diag_soft_cond -> 0.0 otherwise.
|
||||
# We use 6 * eps to take into account the error accumulated during the p1 summation
|
||||
diag_soft_cond = torch.exp(-((p1 / (6 * self._eps)) ** 2)).detach()[..., None]
|
||||
diag_soft_cond = torch.exp(-torch.square(p1 / (6 * self._eps))).detach()[
|
||||
..., None
|
||||
]
|
||||
|
||||
# Eigenvalues are the ordered elements of main diagonal in the degenerate case
|
||||
diag_eigenvals, _ = torch.sort(inputs_diag, dim=-1)
|
||||
@@ -199,8 +203,7 @@ class _SymEig3x3(nn.Module):
|
||||
cross_products[..., :1, :]
|
||||
)
|
||||
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
norms_sq = (cross_products**2).sum(dim=-1)
|
||||
norms_sq = torch.square(cross_products).sum(dim=-1)
|
||||
max_norms_index = norms_sq.argmax(dim=-1)
|
||||
|
||||
# Pick only the cross-product with highest squared norm for each input
|
||||
|
||||
@@ -18,68 +18,89 @@ namespace Renderer {
|
||||
|
||||
template <bool DEV>
|
||||
HOST void destruct(Renderer* self) {
|
||||
if (self->result_d != NULL)
|
||||
if (self->result_d != NULL) {
|
||||
FREE(self->result_d);
|
||||
}
|
||||
self->result_d = NULL;
|
||||
if (self->min_depth_d != NULL)
|
||||
if (self->min_depth_d != NULL) {
|
||||
FREE(self->min_depth_d);
|
||||
}
|
||||
self->min_depth_d = NULL;
|
||||
if (self->min_depth_sorted_d != NULL)
|
||||
if (self->min_depth_sorted_d != NULL) {
|
||||
FREE(self->min_depth_sorted_d);
|
||||
}
|
||||
self->min_depth_sorted_d = NULL;
|
||||
if (self->ii_d != NULL)
|
||||
if (self->ii_d != NULL) {
|
||||
FREE(self->ii_d);
|
||||
}
|
||||
self->ii_d = NULL;
|
||||
if (self->ii_sorted_d != NULL)
|
||||
if (self->ii_sorted_d != NULL) {
|
||||
FREE(self->ii_sorted_d);
|
||||
}
|
||||
self->ii_sorted_d = NULL;
|
||||
if (self->ids_d != NULL)
|
||||
if (self->ids_d != NULL) {
|
||||
FREE(self->ids_d);
|
||||
}
|
||||
self->ids_d = NULL;
|
||||
if (self->ids_sorted_d != NULL)
|
||||
if (self->ids_sorted_d != NULL) {
|
||||
FREE(self->ids_sorted_d);
|
||||
}
|
||||
self->ids_sorted_d = NULL;
|
||||
if (self->workspace_d != NULL)
|
||||
if (self->workspace_d != NULL) {
|
||||
FREE(self->workspace_d);
|
||||
}
|
||||
self->workspace_d = NULL;
|
||||
if (self->di_d != NULL)
|
||||
if (self->di_d != NULL) {
|
||||
FREE(self->di_d);
|
||||
}
|
||||
self->di_d = NULL;
|
||||
if (self->di_sorted_d != NULL)
|
||||
if (self->di_sorted_d != NULL) {
|
||||
FREE(self->di_sorted_d);
|
||||
}
|
||||
self->di_sorted_d = NULL;
|
||||
if (self->region_flags_d != NULL)
|
||||
if (self->region_flags_d != NULL) {
|
||||
FREE(self->region_flags_d);
|
||||
}
|
||||
self->region_flags_d = NULL;
|
||||
if (self->num_selected_d != NULL)
|
||||
if (self->num_selected_d != NULL) {
|
||||
FREE(self->num_selected_d);
|
||||
}
|
||||
self->num_selected_d = NULL;
|
||||
if (self->forw_info_d != NULL)
|
||||
if (self->forw_info_d != NULL) {
|
||||
FREE(self->forw_info_d);
|
||||
}
|
||||
self->forw_info_d = NULL;
|
||||
if (self->min_max_pixels_d != NULL)
|
||||
if (self->min_max_pixels_d != NULL) {
|
||||
FREE(self->min_max_pixels_d);
|
||||
}
|
||||
self->min_max_pixels_d = NULL;
|
||||
if (self->grad_pos_d != NULL)
|
||||
if (self->grad_pos_d != NULL) {
|
||||
FREE(self->grad_pos_d);
|
||||
}
|
||||
self->grad_pos_d = NULL;
|
||||
if (self->grad_col_d != NULL)
|
||||
if (self->grad_col_d != NULL) {
|
||||
FREE(self->grad_col_d);
|
||||
}
|
||||
self->grad_col_d = NULL;
|
||||
if (self->grad_rad_d != NULL)
|
||||
if (self->grad_rad_d != NULL) {
|
||||
FREE(self->grad_rad_d);
|
||||
}
|
||||
self->grad_rad_d = NULL;
|
||||
if (self->grad_cam_d != NULL)
|
||||
if (self->grad_cam_d != NULL) {
|
||||
FREE(self->grad_cam_d);
|
||||
}
|
||||
self->grad_cam_d = NULL;
|
||||
if (self->grad_cam_buf_d != NULL)
|
||||
if (self->grad_cam_buf_d != NULL) {
|
||||
FREE(self->grad_cam_buf_d);
|
||||
}
|
||||
self->grad_cam_buf_d = NULL;
|
||||
if (self->grad_opy_d != NULL)
|
||||
if (self->grad_opy_d != NULL) {
|
||||
FREE(self->grad_opy_d);
|
||||
}
|
||||
self->grad_opy_d = NULL;
|
||||
if (self->n_grad_contributions_d != NULL)
|
||||
if (self->n_grad_contributions_d != NULL) {
|
||||
FREE(self->n_grad_contributions_d);
|
||||
}
|
||||
self->n_grad_contributions_d = NULL;
|
||||
}
|
||||
|
||||
|
||||
@@ -64,8 +64,9 @@ GLOBAL void norm_sphere_gradients(Renderer renderer, const int num_balls) {
|
||||
// The sphere only contributes to the camera gradients if it is
|
||||
// large enough in screen space.
|
||||
if (renderer.ids_sorted_d[idx] > 0 && ii.max.x >= ii.min.x + 3 &&
|
||||
ii.max.y >= ii.min.y + 3)
|
||||
ii.max.y >= ii.min.y + 3) {
|
||||
renderer.ids_sorted_d[idx] = 1;
|
||||
}
|
||||
END_PARALLEL_NORET();
|
||||
};
|
||||
|
||||
|
||||
@@ -139,8 +139,9 @@ GLOBAL void render(
|
||||
coord_y < cam_norm.film_border_top + cam_norm.film_height) {
|
||||
// Initialize the result.
|
||||
if (mode == 0u) {
|
||||
for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id)
|
||||
for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id) {
|
||||
result[c_id] = bg_col[c_id];
|
||||
}
|
||||
} else {
|
||||
result[0] = 0.f;
|
||||
}
|
||||
@@ -190,20 +191,22 @@ GLOBAL void render(
|
||||
"render|found intersection with sphere %u.\n",
|
||||
sphere_id_l[write_idx]);
|
||||
}
|
||||
if (ii.min.x == MAX_USHORT)
|
||||
if (ii.min.x == MAX_USHORT) {
|
||||
// This is an invalid sphere (out of image). These spheres have
|
||||
// maximum depth. Since we ordered the spheres by earliest possible
|
||||
// intersection depth we re certain that there will no other sphere
|
||||
// that is relevant after this one.
|
||||
loading_done = true;
|
||||
}
|
||||
}
|
||||
// Reset n_pixels_done.
|
||||
n_pixels_done = 0;
|
||||
thread_block.sync(); // Make sure n_loaded is updated.
|
||||
if (n_loaded > RENDER_BUFFER_LOAD_THRESH) {
|
||||
// The load buffer is full enough. Draw.
|
||||
if (thread_block.thread_rank() == 0)
|
||||
if (thread_block.thread_rank() == 0) {
|
||||
n_balls_loaded += n_loaded;
|
||||
}
|
||||
max_closest_possible_intersection = 0.f;
|
||||
// This excludes threads outside of the image boundary. Also, it reduces
|
||||
// block artifacts.
|
||||
@@ -290,8 +293,9 @@ GLOBAL void render(
|
||||
uint warp_done = thread_warp.ballot(done);
|
||||
int warp_done_bit_cnt = POPC(warp_done);
|
||||
#endif //__CUDACC__ && __HIP_PLATFORM_AMD__
|
||||
if (thread_warp.thread_rank() == 0)
|
||||
if (thread_warp.thread_rank() == 0) {
|
||||
ATOMICADD_B(&n_pixels_done, warp_done_bit_cnt);
|
||||
}
|
||||
// This sync is necessary to keep n_loaded until all threads are done with
|
||||
// painting.
|
||||
thread_block.sync();
|
||||
@@ -299,8 +303,9 @@ GLOBAL void render(
|
||||
}
|
||||
thread_block.sync();
|
||||
}
|
||||
if (thread_block.thread_rank() == 0)
|
||||
if (thread_block.thread_rank() == 0) {
|
||||
n_balls_loaded += n_loaded;
|
||||
}
|
||||
PULSAR_LOG_DEV_PIX(
|
||||
PULSAR_LOG_RENDER_PIX,
|
||||
"render|loaded %d balls in total.\n",
|
||||
@@ -386,8 +391,9 @@ GLOBAL void render(
|
||||
static_cast<float>(tracker.get_n_hits());
|
||||
} else {
|
||||
float sm_d_normfac = FRCP(FMAX(sm_d, FEPS));
|
||||
for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id)
|
||||
for (uint c_id = 0; c_id < cam_norm.n_channels; ++c_id) {
|
||||
result[c_id] *= sm_d_normfac;
|
||||
}
|
||||
int write_loc = (coord_y - cam_norm.film_border_top) * cam_norm.film_width *
|
||||
(3 + 2 * n_track) +
|
||||
(coord_x - cam_norm.film_border_left) * (3 + 2 * n_track);
|
||||
|
||||
@@ -860,8 +860,9 @@ std::tuple<torch::Tensor, torch::Tensor> Renderer::forward(
|
||||
? (cudaStream_t) nullptr
|
||||
#endif
|
||||
: (cudaStream_t) nullptr);
|
||||
if (mode == 1)
|
||||
if (mode == 1) {
|
||||
results[batch_i] = results[batch_i].slice(2, 0, 1, 1);
|
||||
}
|
||||
forw_infos[batch_i] = from_blob(
|
||||
this->renderer_vec[batch_i].forw_info_d,
|
||||
{this->renderer_vec[0].cam.film_height,
|
||||
|
||||
@@ -128,8 +128,9 @@ struct Renderer {
|
||||
stream << "pulsar::Renderer[";
|
||||
// Device info.
|
||||
stream << self.device_type;
|
||||
if (self.device_index != -1)
|
||||
if (self.device_index != -1) {
|
||||
stream << ", ID " << self.device_index;
|
||||
}
|
||||
stream << "]";
|
||||
return stream;
|
||||
}
|
||||
|
||||
@@ -106,6 +106,8 @@ auto ComputeFaceAreas(const torch::Tensor& face_verts) {
|
||||
return face_areas;
|
||||
}
|
||||
|
||||
namespace {
|
||||
|
||||
// Helper function to use with std::find_if to find the index of any
|
||||
// values in the top k struct which match a given idx.
|
||||
struct IsNeighbor {
|
||||
@@ -118,7 +120,6 @@ struct IsNeighbor {
|
||||
int neighbor_idx;
|
||||
};
|
||||
|
||||
namespace {
|
||||
void RasterizeMeshesNaiveCpu_worker(
|
||||
const int start_yi,
|
||||
const int end_yi,
|
||||
|
||||
@@ -19,7 +19,7 @@ template <
|
||||
std::is_same<T, double>::value || std::is_same<T, float>::value>>
|
||||
struct vec2 {
|
||||
T x, y;
|
||||
typedef T scalar_t;
|
||||
using scalar_t = T;
|
||||
vec2(T x, T y) : x(x), y(y) {}
|
||||
};
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ template <
|
||||
std::is_same<T, double>::value || std::is_same<T, float>::value>>
|
||||
struct vec3 {
|
||||
T x, y, z;
|
||||
typedef T scalar_t;
|
||||
using scalar_t = T;
|
||||
vec3(T x, T y, T z) : x(x), y(y), z(z) {}
|
||||
};
|
||||
|
||||
|
||||
@@ -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:
|
||||
# use try/finally to bypass cinder's lazy imports
|
||||
try:
|
||||
from .blender_dataset_map_provider import ( # noqa: F401
|
||||
BlenderDatasetMapProvider,
|
||||
)
|
||||
from .json_index_dataset_map_provider import ( # noqa: F401
|
||||
JsonIndexDatasetMapProvider,
|
||||
)
|
||||
from .json_index_dataset_map_provider_v2 import ( # noqa: F401
|
||||
JsonIndexDatasetMapProviderV2,
|
||||
)
|
||||
from .llff_dataset_map_provider import LlffDatasetMapProvider # noqa: F401
|
||||
from .rendered_mesh_dataset_map_provider import ( # noqa: F401
|
||||
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):
|
||||
"""
|
||||
Base for provider of data for one scene from LLFF or blender datasets.
|
||||
Base for provider of data for one scene.
|
||||
|
||||
Members:
|
||||
base_dir: directory holding the data for the scene.
|
||||
@@ -171,40 +171,3 @@ class SingleSceneDatasetMapProviderBase(DatasetMapProviderBase):
|
||||
# pyre-ignore[16]
|
||||
cameras = [self.poses[i] for i in self.i_split[0]]
|
||||
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
|
||||
|
||||
@@ -114,9 +114,7 @@ def mesh_laplacian_smoothing(meshes, method: str = "uniform"):
|
||||
if method == "cot":
|
||||
norm_w = torch.sparse.sum(L, dim=1).to_dense().view(-1, 1)
|
||||
idx = norm_w > 0
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and
|
||||
# `Tensor`.
|
||||
norm_w[idx] = 1.0 / norm_w[idx]
|
||||
norm_w[idx] = torch.reciprocal(norm_w[idx])
|
||||
else:
|
||||
L_sum = torch.sparse.sum(L, dim=1).to_dense().view(-1, 1)
|
||||
norm_w = 0.25 * inv_areas
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
|
||||
# pyre-unsafe
|
||||
|
||||
import torch
|
||||
from pytorch3d import _C
|
||||
from pytorch3d.structures import Meshes, Pointclouds
|
||||
from torch.autograd import Function
|
||||
@@ -302,8 +303,7 @@ def point_mesh_edge_distance(meshes: Meshes, pcls: Pointclouds):
|
||||
point_to_cloud_idx = pcls.packed_to_cloud_idx() # (sum(P_i), )
|
||||
num_points_per_cloud = pcls.num_points_per_cloud() # (N,)
|
||||
weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
weights_p = 1.0 / weights_p.float()
|
||||
weights_p = torch.reciprocal(weights_p.float())
|
||||
point_to_edge = point_to_edge * weights_p
|
||||
point_dist = point_to_edge.sum() / N
|
||||
|
||||
@@ -377,8 +377,7 @@ def point_mesh_face_distance(
|
||||
point_to_cloud_idx = pcls.packed_to_cloud_idx() # (sum(P_i),)
|
||||
num_points_per_cloud = pcls.num_points_per_cloud() # (N,)
|
||||
weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx)
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
weights_p = 1.0 / weights_p.float()
|
||||
weights_p = torch.reciprocal(weights_p.float())
|
||||
point_to_face = point_to_face * weights_p
|
||||
point_dist = point_to_face.sum() / N
|
||||
|
||||
|
||||
@@ -223,8 +223,7 @@ def _align_camera_extrinsics(
|
||||
# of centered A and centered B
|
||||
Ac = A - Amu
|
||||
Bc = B - Bmu
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
align_t_s = (Ac * Bc).mean() / (Ac**2).mean().clamp(eps)
|
||||
align_t_s = (Ac * Bc).mean() / torch.square(Ac).mean().clamp(eps)
|
||||
else:
|
||||
# set the scale to identity
|
||||
align_t_s = 1.0
|
||||
|
||||
@@ -55,11 +55,9 @@ def laplacian(verts: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
|
||||
# We construct the Laplacian matrix by adding the non diagonal values
|
||||
# i.e. L[i, j] = 1 ./ deg(i) if (i, j) is an edge
|
||||
deg0 = deg[e0]
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
deg0 = torch.where(deg0 > 0.0, 1.0 / deg0, deg0)
|
||||
deg0 = torch.where(deg0 > 0.0, torch.reciprocal(deg0), deg0)
|
||||
deg1 = deg[e1]
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
deg1 = torch.where(deg1 > 0.0, 1.0 / deg1, deg1)
|
||||
deg1 = torch.where(deg1 > 0.0, torch.reciprocal(deg1), deg1)
|
||||
val = torch.cat([deg0, deg1])
|
||||
L = torch.sparse_coo_tensor(idx, val, (V, V), dtype=torch.float32)
|
||||
|
||||
@@ -137,8 +135,7 @@ def cot_laplacian(
|
||||
val = torch.stack([area] * 3, dim=1).view(-1)
|
||||
inv_areas.scatter_add_(0, idx, val)
|
||||
idx = inv_areas > 0
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
inv_areas[idx] = 1.0 / inv_areas[idx]
|
||||
inv_areas[idx] = torch.reciprocal(inv_areas[idx])
|
||||
inv_areas = inv_areas.view(-1, 1)
|
||||
|
||||
return L, inv_areas
|
||||
|
||||
@@ -182,8 +182,7 @@ def iterative_closest_point(
|
||||
t_history.append(SimilarityTransform(R, T, s))
|
||||
|
||||
# compute the root mean squared error
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
Xt_sq_diff = ((Xt - Xt_nn_points) ** 2).sum(2)
|
||||
Xt_sq_diff = torch.square(Xt - Xt_nn_points).sum(2)
|
||||
rmse = oputil.wmean(Xt_sq_diff[:, :, None], mask_X).sqrt()[:, 0, 0]
|
||||
|
||||
# compute the relative rmse
|
||||
|
||||
@@ -179,9 +179,7 @@ def sample_farthest_points_naive(
|
||||
# and all the other points. If a point has already been selected
|
||||
# it's distance will be 0.0 so it will not be selected again as the max.
|
||||
dist = points[n, selected_idx, :] - points[n, : lengths[n], :]
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
||||
# `int`.
|
||||
dist_to_last_selected = (dist**2).sum(-1) # (P - i)
|
||||
dist_to_last_selected = torch.square(dist).sum(-1) # (P - i)
|
||||
|
||||
# If closer than currently saved distance to one of the selected
|
||||
# points, then updated closest_dists
|
||||
|
||||
@@ -629,10 +629,8 @@ class FoVPerspectiveCameras(CamerasBase):
|
||||
# so the so the z sign is 1.0.
|
||||
z_sign = 1.0
|
||||
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
K[:, 0, 0] = 2.0 * znear / (max_x - min_x)
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
K[:, 1, 1] = 2.0 * znear / (max_y - min_y)
|
||||
K[:, 0, 0] = torch.div(2.0 * znear, max_x - min_x)
|
||||
K[:, 1, 1] = torch.div(2.0 * znear, max_y - min_y)
|
||||
K[:, 0, 2] = (max_x + min_x) / (max_x - min_x)
|
||||
K[:, 1, 2] = (max_y + min_y) / (max_y - min_y)
|
||||
K[:, 3, 2] = z_sign * ones
|
||||
@@ -1178,9 +1176,7 @@ class PerspectiveCameras(CamerasBase):
|
||||
xy_inv_depth = torch.cat(
|
||||
# pyre-fixme[6]: For 1st argument expected `Union[List[Tensor],
|
||||
# tuple[Tensor, ...]]` but got `Tuple[Tensor, float]`.
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and
|
||||
# `Tensor`.
|
||||
(xy_depth[..., :2], 1.0 / xy_depth[..., 2:3]),
|
||||
(xy_depth[..., :2], torch.reciprocal(xy_depth[..., 2:3])),
|
||||
dim=-1, # type: ignore
|
||||
)
|
||||
return unprojection_transform.transform_points(xy_inv_depth)
|
||||
|
||||
@@ -434,13 +434,7 @@ def clip_faces(
|
||||
# These will then be filled in for each case.
|
||||
###########################################
|
||||
F_clipped = (
|
||||
F
|
||||
# pyre-fixme[58]: `+` is not supported for operand types `int` and
|
||||
# `Union[bool, float, int]`.
|
||||
+ faces_delta_cum[-1].item()
|
||||
# pyre-fixme[58]: `+` is not supported for operand types `int` and
|
||||
# `Union[bool, float, int]`.
|
||||
+ faces_delta[-1].item()
|
||||
F + int(faces_delta_cum[-1].item()) + int(faces_delta[-1].item())
|
||||
) # Total number of faces in the new Meshes
|
||||
face_verts_clipped = torch.zeros(
|
||||
(F_clipped, 3, 3), dtype=face_verts_unclipped.dtype, device=device
|
||||
|
||||
@@ -71,9 +71,7 @@ def _list_to_padded_wrapper(
|
||||
# pyre-fixme[6]: For 2nd param expected `int` but got `Union[bool, float, int]`.
|
||||
x_reshaped.append(y.reshape(-1, D))
|
||||
x_padded = list_to_padded(x_reshaped, pad_size=pad_size, pad_value=pad_value)
|
||||
# pyre-fixme[58]: `+` is not supported for operand types `Tuple[int, int]` and
|
||||
# `Size`.
|
||||
return x_padded.reshape((N, -1) + reshape_dims)
|
||||
return x_padded.reshape((N, -1) + tuple(reshape_dims))
|
||||
|
||||
|
||||
def _padded_to_list_wrapper(
|
||||
@@ -104,9 +102,7 @@ def _padded_to_list_wrapper(
|
||||
# pyre-fixme[6]: For 3rd param expected `int` but got `Union[bool, float, int]`.
|
||||
x_reshaped = x.reshape(N, M, D)
|
||||
x_list = padded_to_list(x_reshaped, split_size=split_size)
|
||||
# pyre-fixme[58]: `+` is not supported for operand types `Tuple[typing.Any]` and
|
||||
# `Size`.
|
||||
x_list = [xl.reshape((xl.shape[0],) + reshape_dims) for xl in x_list]
|
||||
x_list = [xl.reshape((xl.shape[0],) + tuple(reshape_dims)) for xl in x_list]
|
||||
return x_list
|
||||
|
||||
|
||||
|
||||
@@ -132,15 +132,13 @@ def _get_splat_kernel_normalization(
|
||||
|
||||
epsilon = 0.05
|
||||
normalization_constant = torch.exp(
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
-(offsets**2).sum(dim=1) / (2 * sigma**2)
|
||||
-torch.square(offsets).sum(dim=1) / (2 * sigma**2)
|
||||
).sum()
|
||||
|
||||
# We add an epsilon to the normalization constant to ensure the gradient will travel
|
||||
# through non-boundary pixels' normalization factor, see Sec. 3.3.1 in "Differentia-
|
||||
# ble Surface Rendering via Non-Differentiable Sampling", Cole et al.
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
return (1 + epsilon) / normalization_constant
|
||||
return torch.div(1 + epsilon, normalization_constant)
|
||||
|
||||
|
||||
def _compute_occlusion_layers(
|
||||
@@ -264,8 +262,9 @@ def _compute_splatting_colors_and_weights(
|
||||
torch.floor(pixel_coords_screen[..., :2]) - pixel_coords_screen[..., :2] + 0.5
|
||||
).view((N, H, W, K, 1, 2))
|
||||
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
dist2_p_q = torch.sum((q_to_px_center + offsets) ** 2, dim=5) # (N, H, W, K, 9)
|
||||
dist2_p_q = torch.sum(
|
||||
torch.square(q_to_px_center + offsets), dim=5
|
||||
) # (N, H, W, K, 9)
|
||||
splat_weights = torch.exp(-dist2_p_q / (2 * sigma**2))
|
||||
alpha = colors[..., 3:4]
|
||||
splat_weights = (alpha * splat_kernel_normalization * splat_weights).unsqueeze(
|
||||
@@ -417,12 +416,12 @@ def _normalize_and_compose_all_layers(
|
||||
device = splatted_colors_per_occlusion_layer.device
|
||||
|
||||
# Normalize each of bg/surface/fg splat layers separately.
|
||||
normalization_scales = 1.0 / (
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
normalization_scales = torch.div(
|
||||
1.0,
|
||||
torch.maximum(
|
||||
splatted_weights_per_occlusion_layer,
|
||||
torch.tensor([1.0], device=device),
|
||||
)
|
||||
),
|
||||
) # (N, H, W, 1, 3)
|
||||
|
||||
normalized_splatted_colors = (
|
||||
|
||||
@@ -269,9 +269,7 @@ class TensorProperties(nn.Module):
|
||||
# to have the same shape as the input tensor.
|
||||
new_dims = len(tensor_dims) - len(idx_dims)
|
||||
new_shape = idx_dims + (1,) * new_dims
|
||||
# pyre-fixme[58]: `+` is not supported for operand types
|
||||
# `Tuple[int]` and `torch._C.Size`
|
||||
expand_dims = (-1,) + tensor_dims[1:]
|
||||
expand_dims = (-1,) + tuple(tensor_dims[1:])
|
||||
_batch_idx = _batch_idx.view(*new_shape)
|
||||
_batch_idx = _batch_idx.expand(*expand_dims)
|
||||
|
||||
|
||||
@@ -52,8 +52,7 @@ def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
|
||||
Rotation matrices as tensor of shape (..., 3, 3).
|
||||
"""
|
||||
r, i, j, k = torch.unbind(quaternions, -1)
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
two_s = 2.0 / (quaternions * quaternions).sum(-1)
|
||||
two_s = torch.div(2.0, (quaternions * quaternions).sum(-1))
|
||||
|
||||
o = torch.stack(
|
||||
(
|
||||
@@ -95,13 +94,9 @@ def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
|
||||
Returns torch.sqrt(torch.max(0, x))
|
||||
but with a zero subgradient where x is 0.
|
||||
"""
|
||||
ret = torch.zeros_like(x)
|
||||
positive_mask = x > 0
|
||||
if torch.is_grad_enabled():
|
||||
ret[positive_mask] = torch.sqrt(x[positive_mask])
|
||||
else:
|
||||
ret = torch.where(positive_mask, torch.sqrt(x), ret)
|
||||
return ret
|
||||
safe_x = torch.where(positive_mask, x, 1.0)
|
||||
return torch.where(positive_mask, torch.sqrt(safe_x), 0.0)
|
||||
|
||||
|
||||
def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
|
||||
@@ -137,18 +132,18 @@ def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
|
||||
# we produce the desired quaternion multiplied by each of r, i, j, k
|
||||
quat_by_rijk = torch.stack(
|
||||
[
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
||||
# `int`.
|
||||
torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
||||
# `int`.
|
||||
torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
||||
# `int`.
|
||||
torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
||||
# `int`.
|
||||
torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
|
||||
torch.stack(
|
||||
[torch.square(q_abs[..., 0]), m21 - m12, m02 - m20, m10 - m01], dim=-1
|
||||
),
|
||||
torch.stack(
|
||||
[m21 - m12, torch.square(q_abs[..., 1]), m10 + m01, m02 + m20], dim=-1
|
||||
),
|
||||
torch.stack(
|
||||
[m02 - m20, m10 + m01, torch.square(q_abs[..., 2]), m12 + m21], dim=-1
|
||||
),
|
||||
torch.stack(
|
||||
[m10 - m01, m20 + m02, m21 + m12, torch.square(q_abs[..., 3])], dim=-1
|
||||
),
|
||||
],
|
||||
dim=-2,
|
||||
)
|
||||
|
||||
@@ -195,15 +195,15 @@ def _se3_V_matrix(
|
||||
V = (
|
||||
torch.eye(3, dtype=log_rotation.dtype, device=log_rotation.device)[None]
|
||||
+ log_rotation_hat
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
* ((1 - torch.cos(rotation_angles)) / (rotation_angles**2))[:, None, None]
|
||||
* ((1 - torch.cos(rotation_angles)) / torch.square(rotation_angles))[
|
||||
:, None, None
|
||||
]
|
||||
+ (
|
||||
log_rotation_hat_square
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
||||
# `int`.
|
||||
* ((rotation_angles - torch.sin(rotation_angles)) / (rotation_angles**3))[
|
||||
:, None, None
|
||||
]
|
||||
* (
|
||||
(rotation_angles - torch.sin(rotation_angles))
|
||||
/ torch.pow(rotation_angles, 3)
|
||||
)[:, None, None]
|
||||
)
|
||||
)
|
||||
|
||||
@@ -215,8 +215,7 @@ def _get_se3_V_input(log_rotation: torch.Tensor, eps: float = 1e-4):
|
||||
A helper function that computes the input variables to the `_se3_V_matrix`
|
||||
function.
|
||||
"""
|
||||
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and `int`.
|
||||
nrms = (log_rotation**2).sum(-1)
|
||||
nrms = torch.square(log_rotation).sum(-1)
|
||||
rotation_angles = torch.clamp(nrms, eps).sqrt()
|
||||
log_rotation_hat = hat(log_rotation)
|
||||
log_rotation_hat_square = torch.bmm(log_rotation_hat, log_rotation_hat)
|
||||
|
||||
@@ -623,9 +623,7 @@ class Scale(Transform3d):
|
||||
Return the inverse of self._matrix.
|
||||
"""
|
||||
xyz = torch.stack([self._matrix[:, i, i] for i in range(4)], dim=1)
|
||||
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
||||
ixyz = 1.0 / xyz
|
||||
# pyre-fixme[6]: For 1st param expected `Tensor` but got `float`.
|
||||
ixyz = torch.reciprocal(xyz)
|
||||
imat = torch.diag_embed(ixyz, dim1=1, dim2=2)
|
||||
return imat
|
||||
|
||||
|
||||
@@ -1,12 +1,5 @@
|
||||
dataset_map_provider_class_type: ???
|
||||
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:
|
||||
category: ???
|
||||
task_str: singlesequence
|
||||
@@ -78,14 +71,6 @@ dataset_map_provider_JsonIndexDatasetMapProviderV2_args:
|
||||
sort_frames: false
|
||||
path_manager_factory_PathManagerFactory_args:
|
||||
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:
|
||||
num_views: 40
|
||||
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)
|
||||
@@ -74,20 +74,6 @@ class HomeSplash extends React.Component {
|
||||
}
|
||||
}
|
||||
|
||||
function SocialBanner() {
|
||||
return (
|
||||
<div className="socialBanner">
|
||||
<div>
|
||||
Support Ukraine 🇺🇦{' '}
|
||||
<a href="https://opensource.fb.com/support-ukraine">
|
||||
Help Provide Humanitarian Aid to Ukraine
|
||||
</a>
|
||||
.
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
class Index extends React.Component {
|
||||
render() {
|
||||
const {config: siteConfig, language = ''} = this.props;
|
||||
@@ -226,7 +212,6 @@ loss_chamfer, _ = chamfer_distance(sample_sphere, sample_test)
|
||||
|
||||
return (
|
||||
<div>
|
||||
<SocialBanner />
|
||||
<HomeSplash siteConfig={siteConfig} language={language} />
|
||||
<div className="landingPage mainContainer">
|
||||
<Features />
|
||||
|
||||
Reference in New Issue
Block a user