text: # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
  """
  Sanity checks for loading ShapeNetCore.
  """
  import os
  import unittest
  
  import numpy as np
  import torch
  from common_testing import TestCaseMixin, load_rgb_image, get_tests_dir
  from PIL import Image
  from pytorch3d.datasets import ShapeNetCore, collate_batched_meshes
  from pytorch3d.renderer import (
      FoVPerspectiveCameras,
      PointLights,
      RasterizationSettings,
      look_at_view_transform,
  )
  from torch.utils.data import DataLoader
  
  
  # Set the SHAPENET_PATH to the local path to the dataset
  SHAPENET_PATH = None
  VERSION = 1
  # If DEBUG=True, save out images generated in the tests for debugging.
  # All saved images have prefix DEBUG_
  DEBUG = False
  DATA_DIR = get_tests_dir() / "data"
  
  
  class TestShapenetCore(TestCaseMixin, unittest.TestCase):
      def setUp(self):
          """
          Check if the ShapeNet dataset is provided in the repo.
          If not, download this separately and update the shapenet_path`
          with the location of the dataset in order to run the tests.
          """
          if SHAPENET_PATH is None or not os.path.exists(SHAPENET_PATH):
              url = "https://www.shapenet.org/"
              msg = (
                  "ShapeNet data not found, download from %s, update "
                  "SHAPENET_PATH at the top of the file, and rerun."
              )
  
              self.skipTest(msg % url)
  
      def test_load_shapenet_core(self):
          """
          Test loading both the entire ShapeNetCore dataset and a subset of the ShapeNetCore
          dataset. Check the loaded datasets return items of the correct shapes and types.
          """
          # Try loading ShapeNetCore with an invalid version number and catch error.
          with self.assertRaises(ValueError) as err:
              ShapeNetCore(SHAPENET_PATH, version=3)
          self.assertTrue("Version number must be either 1 or 2." in str(err.exception))
  
          # Load ShapeNetCore without specifying any particular categories.
          shapenet_dataset = ShapeNetCore(SHAPENET_PATH, version=VERSION)
  
          # Count the number of grandchildren directories (which should be equal to
          # the total number of objects in the dataset) by walking through the given
          # directory.
          wnsynset_list = [
              wnsynset
              for wnsynset in os.listdir(SHAPENET_PATH)
              if os.path.isdir(os.path.join(SHAPENET_PATH, wnsynset))
          ]
          model_num_list = [
              (len(next(os.walk(os.path.join(SHAPENET_PATH, wnsynset)))[1]))
              for wnsynset in wnsynset_list
          ]
          # Check total number of objects in the dataset is correct.
          self.assertEqual(len(shapenet_dataset), sum(model_num_list))
  
          # Randomly retrieve an object from the dataset.
          rand_obj = shapenet_dataset[torch.randint(len(shapenet_dataset), (1,))]
          # Check that data types and shapes of items returned by __getitem__ are correct.
          verts, faces = rand_obj["verts"], rand_obj["faces"]
          self.assertTrue(verts.dtype == torch.float32)
          self.assertTrue(faces.dtype == torch.int64)
          self.assertEqual(verts.ndim, 2)
          self.assertEqual(verts.shape[-1], 3)
          self.assertEqual(faces.ndim, 2)
          self.assertEqual(faces.shape[-1], 3)
  
          # Load six categories from ShapeNetCore.
          # Specify categories with a combination of offsets and labels.
          shapenet_subset = ShapeNetCore(
              SHAPENET_PATH,
              synsets=[
                  "04330267",
                  "guitar",
                  "02801938",
                  "birdhouse",
                  "03991062",
                  "tower",
              ],
              version=1,
          )
          subset_offsets = [
              "04330267",
              "03467517",
              "02801938",
              "02843684",
              "03991062",
              "04460130",
          ]
          subset_model_nums = [
              (len(next(os.walk(os.path.join(SHAPENET_PATH, offset)))[1]))
              for offset in subset_offsets
          ]
          self.assertEqual(len(shapenet_subset), sum(subset_model_nums))
  
      def test_collate_models(self):
          """
          Test collate_batched_meshes returns items of the correct shapes and types.
          Check that when collate_batched_meshes is passed to Dataloader, batches of
          the correct shapes and types are returned.
          """
          # Load ShapeNetCore without specifying any particular categories.
          shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
          # Randomly retrieve several objects from the dataset.
          rand_idxs = torch.randint(len(shapenet_dataset), (6,))
          rand_objs = [shapenet_dataset[idx] for idx in rand_idxs]
  
          # Collate the randomly selected objects
          collated_meshes = collate_batched_meshes(rand_objs)
          verts, faces = (collated_meshes["verts"], collated_meshes["faces"])
          self.assertEqual(len(verts), 6)
          self.assertEqual(len(faces), 6)
  
          # Pass the custom collate_fn function to DataLoader and check elements
          # in batch have the correct shape.
          batch_size = 12
          shapenet_core_loader = DataLoader(
              shapenet_dataset, batch_size=batch_size, collate_fn=collate_batched_meshes
          )
          it = iter(shapenet_core_loader)
          object_batch = next(it)
          self.assertEqual(len(object_batch["synset_id"]), batch_size)
          self.assertEqual(len(object_batch["model_id"]), batch_size)
          self.assertEqual(len(object_batch["label"]), batch_size)
          self.assertEqual(object_batch["mesh"].verts_padded().shape[0], batch_size)
          self.assertEqual(object_batch["mesh"].faces_padded().shape[0], batch_size)
  
      def test_catch_render_arg_errors(self):
          """
          Test rendering ShapeNetCore with invalid model_ids, categories or indices,
          and catch corresponding errors.
          """
          # Load ShapeNetCore.
          shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
  
          # Try loading with an invalid model_id and catch error.
          with self.assertRaises(ValueError) as err:
              shapenet_dataset.render(model_ids=["piano0"])
          self.assertTrue("not found in the loaded dataset" in str(err.exception))
  
          # Try loading with an index out of bounds and catch error.
          with self.assertRaises(IndexError) as err:
              shapenet_dataset.render(idxs=[100000])
          self.assertTrue("are out of bounds" in str(err.exception))
  
      def test_render_shapenet_core(self):
          """
          Test rendering objects from ShapeNetCore.
          """
          # Setup device and seed for random selections.
          device = torch.device("cuda:0")
          torch.manual_seed(39)
  
          # Load category piano from ShapeNetCore.
          piano_dataset = ShapeNetCore(SHAPENET_PATH, synsets=["piano"])
  
          # Rendering settings.
          R, T = look_at_view_transform(1.0, 1.0, 90)
          cameras = FoVPerspectiveCameras(R=R, T=T, device=device)
          raster_settings = RasterizationSettings(image_size=512)
          lights = PointLights(
              location=torch.tensor([0.0, 1.0, -2.0], device=device)[None],
              # TODO: debug the source of the discrepancy in two images when rendering on GPU.
              diffuse_color=((0, 0, 0),),
              specular_color=((0, 0, 0),),
              device=device,
          )
  
          # Render first three models in the piano category.
          pianos = piano_dataset.render(
              idxs=list(range(3)),
              device=device,
              cameras=cameras,
              raster_settings=raster_settings,
              lights=lights,
          )
          # Check that there are three images in the batch.
          self.assertEqual(pianos.shape[0], 3)
  
          # Compare the rendered models to the reference images.
          for idx in range(3):
              piano_rgb = pianos[idx, ..., :3].squeeze().cpu()
              if DEBUG:
                  Image.fromarray((piano_rgb.numpy() * 255).astype(np.uint8)).save(
                      DATA_DIR / ("DEBUG_shapenet_core_render_piano_by_idxs_%s.png" % idx)
                  )
              image_ref = load_rgb_image(
                  "test_shapenet_core_render_piano_%s.png" % idx, DATA_DIR
              )
              self.assertClose(piano_rgb, image_ref, atol=0.05)
  
          # Render the same piano models but by model_ids this time.
          pianos_2 = piano_dataset.render(
              model_ids=[
                  "13394ca47c89f91525a3aaf903a41c90",
                  "14755c2ee8e693aba508f621166382b0",
                  "156c4207af6d2c8f1fdc97905708b8ea",
              ],
              device=device,
              cameras=cameras,
              raster_settings=raster_settings,
              lights=lights,
          )
  
          # Compare the rendered models to the reference images.
          for idx in range(3):
              piano_rgb_2 = pianos_2[idx, ..., :3].squeeze().cpu()
              if DEBUG:
                  Image.fromarray((piano_rgb_2.numpy() * 255).astype(np.uint8)).save(
                      DATA_DIR / ("DEBUG_shapenet_core_render_piano_by_ids_%s.png" % idx)
                  )
              image_ref = load_rgb_image(
                  "test_shapenet_core_render_piano_%s.png" % idx, DATA_DIR
              )
              self.assertClose(piano_rgb_2, image_ref, atol=0.05)
  
          #######################
          # Render by categories
          #######################
  
          # Load ShapeNetCore.
          shapenet_dataset = ShapeNetCore(SHAPENET_PATH)
  
          # Render a mixture of categories and specify the number of models to be
          # randomly sampled from each category.
          mixed_objs = shapenet_dataset.render(
              categories=["faucet", "chair"],
              sample_nums=[2, 1],
              device=device,
              cameras=cameras,
              raster_settings=raster_settings,
              lights=lights,
          )
          # Compare the rendered models to the reference images.
          for idx in range(3):
              mixed_rgb = mixed_objs[idx, ..., :3].squeeze().cpu()
              if DEBUG:
                  Image.fromarray((mixed_rgb.numpy() * 255).astype(np.uint8)).save(
                      DATA_DIR
                      / ("DEBUG_shapenet_core_render_mixed_by_categories_%s.png" % idx)
                  )
              image_ref = load_rgb_image(
                  "test_shapenet_core_render_mixed_by_categories_%s.png" % idx, DATA_DIR
              )
              self.assertClose(mixed_rgb, image_ref, atol=0.05)
  
          # Render a mixture of categories without specifying sample_nums.
          mixed_objs_2 = shapenet_dataset.render(
              categories=["faucet", "chair"],
              device=device,
              cameras=cameras,
              raster_settings=raster_settings,
              lights=lights,
          )
          # Compare the rendered models to the reference images.
          for idx in range(2):
              mixed_rgb_2 = mixed_objs_2[idx, ..., :3].squeeze().cpu()
              if DEBUG:
                  Image.fromarray((mixed_rgb_2.numpy() * 255).astype(np.uint8)).save(
                      DATA_DIR
                      / ("DEBUG_shapenet_core_render_without_sample_nums_%s.png" % idx)
                  )
              image_ref = load_rgb_image(
                  "test_shapenet_core_render_without_sample_nums_%s.png" % idx, DATA_DIR
              )
              self.assertClose(mixed_rgb_2, image_ref, atol=0.05)
  