changed object id logic

Author: albertaillet

builder.py

@@ -1,5 +1,5 @@
 from jax import numpy as np, tree_map
-from raymarch import Objects, Camera, OBJECT_IDX
+from raymarch import Objects, Camera, OBJECTS
 
 # typing
 from typeguard import check_type
@@ -18,7 +18,7 @@ class ObjectDict(TypedDict):
     color: Tuple[float, float, float]
     position: Tuple[float, float, float]
     rotation: Tuple[float, float, float]
-    mirror: Tuple[float, float, float]
+    mirror: Tuple[int, int, int]
     rounding: float
 
 
@@ -57,7 +57,7 @@ def check_scene_dict(scene_dict: Dict[str, Any]) -> SceneDict:
     # checking the obj_names and adding default values where needed
     for i in range(len(scene_dict['Objects'])):
         obj_name, obj_dict = next(iter(scene_dict['Objects'][i].items()))  # first item of dict
-        assert obj_name in OBJECT_IDX, f'Unknown object name {obj_name}'
+        assert obj_name in OBJECTS, f'Unknown object name {obj_name}'
         scene_dict['Objects'][i] = (obj_name, add_obj_dict_defaults(**obj_dict))
 
     return check_type(scene_dict, SceneDict)
@@ -72,7 +72,7 @@ def build_scene(scene_dict: SceneDict) -> Dict[str, Any]:
     obj_args['mirrorings'] = obj_args.pop('mirrors').astype(np.bool_)  # convert mirrorings
     return {
         'objects': Objects(
-            object_ids=np.uint8([OBJECT_IDX[o] for o in obj_names]),  # type: ignore
+            ids=np.uint8([OBJECTS.index(o) for o in obj_names]),  # type: ignore
             **obj_args,
             smoothing=np.float32(scene_dict['smoothing']),
         ),

raymarch.py

@@ -42,13 +42,7 @@ def sdf_ellipsoid(p: Vec3, r: Vec3) -> Scalar:
     return k0 * (k0 - 1) / k1
 
 
-OBJECT_IDX = {
-    'Box': 0,
-    'Sphere': 1,
-    'Plane': 2,
-    'Torus': 3,
-    'Ellipsoid': 4,
-}
+OBJECTS = ('Box', 'Sphere', 'Plane', 'Torus', 'Ellipsoid')
 BRANCHES = (sdf_box, sdf_sphere, sdf_plane, sdf_torus, sdf_ellipsoid)
 
 
@@ -74,7 +68,7 @@ def __call__(self, view_size: Tuple[int, int]) -> Vec3s:
 
 
 class Objects(NamedTuple):
-    object_ids: UInts
+    ids: UInts
     positions: Vec3s
     attributes: Vec3s
     rotations: Vec3s
@@ -84,15 +78,13 @@ class Objects(NamedTuple):
     smoothing: Scalar
 
     def sdfs(self, p: Vec3) -> Scalars:
-        def switch(
-            p: Vec3, obj_idx: UInt, pos: Vec3, attr: Vec3, rot: Vec3, mirror: Bool3
-        ) -> Scalar:
+        def switch(p: Vec3, id: UInt, pos: Vec3, attr: Vec3, rot: Vec3, mirror: Bool3) -> Scalar:
             p = np.where(mirror, smoothabs(p, 1e-3), p)
             p = (p - pos) @ Rxyz(rot)
-            return lax.switch(obj_idx, BRANCHES, p, attr)
+            return lax.switch(id, BRANCHES, p, attr)
 
         dists = vmap(partial(switch, p))(
-            self.object_ids,
+            self.ids,
             self.positions,
             self.attributes,
             self.rotations,

tests.py

@@ -5,7 +5,7 @@
 from matplotlib import pyplot as plt
 from raymarch import render_scene, RenderedImages
 from builder import build_scene, check_scene_dict, SceneDict
-from setup import SCENES_PATH
+from app import SCENES_PATH
 
 
 def test_all_scenes():