Llff & blender convention fix

Summary: Images were coming out in the wrong format.

Reviewed By: shapovalov

Differential Revision: D37291278

fbshipit-source-id: c10871c37dd186982e7abf2071ac66ed583df2e6

pytorch3d/implicitron/dataset/blender_dataset_map_provider.py

@@ -42,7 +42,7 @@ class BlenderDatasetMapProvider(SingleSceneDatasetMapProviderBase):
         )
         H, W, focal = hwf
         H, W = int(H), int(W)
-        images = torch.from_numpy(images)
+        images = torch.from_numpy(images).permute(0, 3, 1, 2)[:, :3]
 
         # pyre-ignore[16]
         self.poses = _interpret_blender_cameras(poses, H, W, focal)

pytorch3d/implicitron/dataset/llff_dataset_map_provider.py

@@ -50,7 +50,7 @@ class LlffDatasetMapProvider(SingleSceneDatasetMapProviderBase):
         i_split = (i_train, i_test, i_test)
         H, W, focal = hwf
         H, W = int(H), int(W)
-        images = torch.from_numpy(images)
+        images = torch.from_numpy(images).permute(0, 3, 1, 2)
         poses = torch.from_numpy(poses)
 
         # pyre-ignore[16]

tests/implicitron/test_data_llff.py

@@ -40,6 +40,7 @@ class TestDataLlff(TestCaseMixin, unittest.TestCase):
             self.assertEqual(len(dataset), length)
             # try getting a value
             value = dataset[0]
+            self.assertEqual(value.image_rgb.shape, (3, 800, 800))
             self.assertIsInstance(value, FrameData)
 
     def test_llff(self):
@@ -62,6 +63,7 @@ class TestDataLlff(TestCaseMixin, unittest.TestCase):
             value = dataset[0]
             self.assertIsInstance(value, FrameData)
             self.assertEqual(value.frame_type, frame_type)
+            self.assertEqual(value.image_rgb.shape, (3, 378, 504))
 
         self.assertEqual(len(dataset_map.test.get_eval_batches()), 3)
         for batch in dataset_map.test.get_eval_batches():