test_data_llff.py 5.35 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
# 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

from pytorch3d.implicitron.dataset.blender_dataset_map_provider import (
    BlenderDatasetMapProvider,
)
13
from pytorch3d.implicitron.dataset.data_source import ImplicitronDataSource
14
15
16
17
from pytorch3d.implicitron.dataset.dataset_base import FrameData
from pytorch3d.implicitron.dataset.llff_dataset_map_provider import (
    LlffDatasetMapProvider,
)
18
from pytorch3d.implicitron.tools.config import expand_args_fields, get_default_args
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
19
from pytorch3d.renderer import PerspectiveCameras
20
21
22
23
24
25
26
27
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)


Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
28
@unittest.skipUnless(internal, "no data")
29
30
class TestDataLlff(TestCaseMixin, unittest.TestCase):
    def test_synthetic(self):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
31
32
        if inside_re_worker:
            return
33
34
35
36
37
38
39
40
41
42
43
44
45
        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()

        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]
46
            self.assertEqual(value.image_rgb.shape, (3, 800, 800))
47
48
49
50
            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)
51
52
53
            self.assertIsInstance(value, FrameData)

    def test_llff(self):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
54
55
        if inside_re_worker:
            return
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
        expand_args_fields(LlffDatasetMapProvider)

        provider = LlffDatasetMapProvider(
            base_dir="manifold://co3d/tree/nerf_data/nerf_llff_data/fern",
            object_name="fern",
        )
        dataset_map = provider.get_dataset_map()

        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)
75
            self.assertEqual(value.image_rgb.shape, (3, 378, 504))
76
77
78
79
80
81
82

        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):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
83
84
        if inside_re_worker:
            return
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        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)
104
                self.assertIsNone(value.fg_probability)
105
106
107
108
109
110
111

        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")
112
113

    def test_loaders(self):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
114
115
        if inside_re_worker:
            return
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
        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))
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
133
134
135
136

        cameras = data_source.get_all_train_cameras()
        self.assertIsInstance(cameras, PerspectiveCameras)
        self.assertEqual(len(cameras), 100)