Spaces:
Running
Running
File size: 6,501 Bytes
7088d16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
# 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)
|