|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
import unittest |
|
|
|
|
|
import torch |
|
|
from pytorch3d.implicitron.dataset.frame_data import FrameData |
|
|
from pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider import ( |
|
|
RenderedMeshDatasetMapProvider, |
|
|
) |
|
|
from pytorch3d.implicitron.tools.config import expand_args_fields |
|
|
from pytorch3d.renderer import FoVPerspectiveCameras |
|
|
from tests.common_testing import TestCaseMixin |
|
|
|
|
|
|
|
|
inside_re_worker = os.environ.get("INSIDE_RE_WORKER", False) |
|
|
|
|
|
|
|
|
class TestDataCow(TestCaseMixin, unittest.TestCase): |
|
|
def test_simple(self): |
|
|
if inside_re_worker: |
|
|
return |
|
|
expand_args_fields(RenderedMeshDatasetMapProvider) |
|
|
self._runtest(use_point_light=True, num_views=4) |
|
|
self._runtest(use_point_light=False, num_views=4) |
|
|
|
|
|
def _runtest(self, **kwargs): |
|
|
provider = RenderedMeshDatasetMapProvider(**kwargs) |
|
|
dataset_map = provider.get_dataset_map() |
|
|
known_matrix = torch.zeros(1, 4, 4) |
|
|
known_matrix[0, 0, 0] = 1.7321 |
|
|
known_matrix[0, 1, 1] = 1.7321 |
|
|
known_matrix[0, 2, 2] = 1.0101 |
|
|
known_matrix[0, 3, 2] = -1.0101 |
|
|
known_matrix[0, 2, 3] = 1 |
|
|
|
|
|
self.assertIsNone(dataset_map.val) |
|
|
self.assertIsNone(dataset_map.test) |
|
|
self.assertEqual(len(dataset_map.train), provider.num_views) |
|
|
|
|
|
value = dataset_map.train[0] |
|
|
self.assertIsInstance(value, FrameData) |
|
|
|
|
|
self.assertEqual(value.image_rgb.shape, (3, 128, 128)) |
|
|
self.assertEqual(value.fg_probability.shape, (1, 128, 128)) |
|
|
|
|
|
self.assertEqual(value.fg_probability[0, 0, 0], 0) |
|
|
self.assertEqual(value.fg_probability.max(), 1.0) |
|
|
self.assertIsInstance(value.camera, FoVPerspectiveCameras) |
|
|
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) |
|
|
|
