[MISC] Move unit test rendering snapshots on huggingface.

.gitignore

@@ -69,6 +69,7 @@ coverage.xml
 *.py,cover
 .hypothesis/
 .pytest_cache/
+__snapshots__
 
 # Translations
 *.mo

tests/test_fem.py

@@ -6,7 +6,7 @@
 import genesis as gs
 from genesis.utils.misc import tensor_to_array
 
-from .utils import assert_allclose, get_hf_assets
+from .utils import assert_allclose, get_hf_dataset
 
 
 @pytest.fixture(scope="session")
@@ -375,7 +375,7 @@ def test_fem_sphere_box_self(fem_material_linear_corotated, fem_material_linear_
 
     # Add second FEM entity
     scale = 0.1
-    asset_path = get_hf_assets(pattern="meshes/cube8.obj")
+    asset_path = get_hf_dataset(pattern="meshes/cube8.obj")
     scene.add_entity(
         morph=gs.morphs.Mesh(
             file=f"{asset_path}/meshes/cube8.obj",

tests/test_mesh.py

@@ -9,7 +9,7 @@
 import genesis.utils.mesh as mu
 import genesis.utils.usda as usda_utils
 
-from .utils import assert_allclose, assert_array_equal, get_hf_assets
+from .utils import assert_allclose, assert_array_equal, get_hf_dataset
 
 VERTICES_TOL = 1e-05  # Transformation loses a little precision in vertices
 NORMALS_TOL = 1e-02  # Conversion from .usd to .glb loses a little precision in normals
@@ -248,9 +248,9 @@ def test_glb_parse_material(glb_file):
 @pytest.mark.required
 @pytest.mark.parametrize("usd_filename", ["usd/sneaker_airforce", "usd/RoughnessTest"])
 def test_usd_parse(usd_filename):
-    asset_path = get_hf_assets(pattern=f"{usd_filename}.glb")
+    asset_path = get_hf_dataset(pattern=f"{usd_filename}.glb")
     glb_file = os.path.join(asset_path, f"{usd_filename}.glb")
-    asset_path = get_hf_assets(pattern=f"{usd_filename}.usdz")
+    asset_path = get_hf_dataset(pattern=f"{usd_filename}.usdz")
     usd_file = os.path.join(asset_path, f"{usd_filename}.usdz")
 
     gs_glb_meshes = gltf_utils.parse_mesh_glb(

tests/test_render.py

@@ -1,19 +1,21 @@
 import queue
 import sys
 from io import BytesIO
+from pathlib import Path
 
 import numpy as np
 import pytest
 import torch
 from PIL import Image
 from syrupy.extensions.image import PNGImageSnapshotExtension
+from syrupy.location import PyTestLocation
 
 import genesis as gs
 import genesis.utils.geom as gu
 from genesis.utils import set_random_seed
 from genesis.utils.image_exporter import FrameImageExporter
 
-from .utils import assert_allclose, assert_array_equal
+from .utils import assert_allclose, assert_array_equal, get_hf_dataset
 
 
 IMG_STD_ERR_THR = 0.8
@@ -430,6 +432,9 @@ def test_madrona_batch_rendering(tmp_path, use_rasterizer, render_all_cameras, n
 
     pytest.importorskip("gs_madrona", reason="Python module 'gs-madrona' not installed.")
 
+    snapshot_dir = Path(PixelMatchSnapshotExtension.dirname(test_location=png_snapshot.test_location))
+    get_hf_dataset(pattern=f"{snapshot_dir.name}/*", repo_name="snapshots", local_dir=snapshot_dir.parent)
+
     scene = gs.Scene(
         sim_options=gs.options.SimOptions(
             dt=0.02,

tests/test_rigid_physics.py

@@ -21,7 +21,7 @@
     build_mujoco_sim,
     check_mujoco_data_consistency,
     check_mujoco_model_consistency,
-    get_hf_assets,
+    get_hf_dataset,
     init_simulators,
     simulate_and_check_mujoco_consistency,
 )
@@ -548,7 +548,7 @@ def test_urdf_rope(
     dof_damping,
     show_viewer,
 ):
-    asset_path = get_hf_assets(pattern="linear_deformable.urdf")
+    asset_path = get_hf_dataset(pattern="linear_deformable.urdf")
     xml_path = os.path.join(asset_path, "linear_deformable.urdf")
 
     mj_sim = build_mujoco_sim(
@@ -1705,7 +1705,7 @@ def test_mesh_repair(convexify, show_viewer, gjk_collision):
         show_viewer=show_viewer,
         show_FPS=False,
     )
-    asset_path = get_hf_assets(pattern="work_table.glb")
+    asset_path = get_hf_dataset(pattern="work_table.glb")
     table = scene.add_entity(
         gs.morphs.Mesh(
             file=f"{asset_path}/work_table.glb",
@@ -1714,7 +1714,7 @@ def test_mesh_repair(convexify, show_viewer, gjk_collision):
         ),
         vis_mode="collision",
     )
-    asset_path = get_hf_assets(pattern="spoon.glb")
+    asset_path = get_hf_dataset(pattern="spoon.glb")
     obj = scene.add_entity(
         gs.morphs.Mesh(
             file=f"{asset_path}/spoon.glb",
@@ -1787,7 +1787,7 @@ def test_convexify(euler, backend, show_viewer, gjk_collision):
     )
     objs = []
     for i, asset_name in enumerate(("mug_1", "donut_0", "cup_2", "apple_15")):
-        asset_path = get_hf_assets(pattern=f"{asset_name}/*")
+        asset_path = get_hf_dataset(pattern=f"{asset_name}/*")
         obj = scene.add_entity(
             gs.morphs.MJCF(
                 file=f"{asset_path}/{asset_name}/output.xml",
@@ -1896,7 +1896,7 @@ def test_collision_plane_convex(show_viewer, tol):
 
         scene.add_entity(morph)
 
-        asset_path = get_hf_assets(pattern="image_0000_segmented.glb")
+        asset_path = get_hf_dataset(pattern="image_0000_segmented.glb")
         asset = scene.add_entity(
             gs.morphs.Mesh(
                 file=f"{asset_path}/image_0000_segmented.glb",
@@ -2078,7 +2078,7 @@ def test_urdf_parsing(show_viewer, tol):
         show_viewer=show_viewer,
         show_FPS=False,
     )
-    asset_path = get_hf_assets(pattern="microwave/*")
+    asset_path = get_hf_dataset(pattern="microwave/*")
     entities = {}
     for i, (fixed, merge_fixed_links) in enumerate(
         ((False, False), (False, True), (True, False), (True, True)),
@@ -2340,7 +2340,7 @@ def test_drone_advanced(show_viewer):
         show_FPS=False,
     )
     plane = scene.add_entity(gs.morphs.Plane())
-    asset_path = get_hf_assets(pattern="drone_sus/*")
+    asset_path = get_hf_dataset(pattern="drone_sus/*")
     drones = []
     for offset, merge_fixed_links in ((-0.3, False), (0.3, True)):
         drone = scene.add_entity(

tests/utils.py

@@ -16,6 +16,7 @@
 import mujoco
 import torch
 from huggingface_hub import snapshot_download
+from PIL import Image, UnidentifiedImageError
 from requests.exceptions import HTTPError
 
 import genesis as gs
@@ -28,6 +29,9 @@
 REPOSITY_URL = "Genesis-Embodied-AI/Genesis"
 DEFAULT_BRANCH_NAME = "main"
 
+MESH_EXTENSIONS = (".mtl", ".glb", ".obj", ".stl", ".usb", ".usdz", ".mdl")
+IMAGE_EXTENSIONS = (".png", ".jpg")
+
 # Get repository "root" path (actually test dir is good enough)
 TEST_DIR = os.path.dirname(__file__)
 
@@ -164,7 +168,9 @@ def get_git_commit_info(ref="HEAD"):
     return revision, timestamp
 
 
-def get_hf_assets(pattern, num_retry: int = 4, retry_delay: float = 30.0, check: bool = True):
+def get_hf_dataset(
+    pattern, repo_name: str = "assets", local_dir: str | None = None, num_retry: int = 4, retry_delay: float = 30.0
+):
     assert num_retry >= 1
 
     for _ in range(num_retry):
@@ -173,16 +179,22 @@ def get_hf_assets(pattern, num_retry: int = 4, retry_delay: float = 30.0, check:
             # Try downloading the assets
             asset_path = snapshot_download(
                 repo_type="dataset",
-                repo_id="Genesis-Intelligence/assets",
+                repo_id=f"Genesis-Intelligence/{repo_name}",
                 allow_patterns=pattern,
                 max_workers=1,
+                local_dir=local_dir,
             )
 
             # Make sure that download was successful
             has_files = False
             for path in Path(asset_path).rglob(pattern):
                 if not path.is_file():
                     continue
+
+                ext = path.suffix.lower()
+                if not ext in (".xml", ".urdf", *IMAGE_EXTENSIONS, *MESH_EXTENSIONS):
+                    continue
+
                 has_files = True
 
                 if path.stat().st_size == 0:
@@ -193,9 +205,18 @@ def get_hf_assets(pattern, num_retry: int = 4, retry_delay: float = 30.0, check:
                         ET.parse(path)
                     except ET.ParseError as e:
                         raise HTTPError(f"Impossible to parse XML file.") from e
+                elif path.suffix.lower() in IMAGE_EXTENSIONS:
+                    try:
+                        Image.open(path)
+                    except UnidentifiedImageError as e:
+                        raise HTTPError(f"Impossible to parse Image file.") from e
+                elif path.suffix.lower() in MESH_EXTENSIONS:
+                    # TODO: Validating mesh files is more tricky. Ignoring them for now.
+                    pass
+
             if not has_files:
                 raise HTTPError("No file downloaded.")
-        except HTTPError:
+        except HTTPError as e:
             num_trials += 1
             if num_trials == num_retry:
                 raise