Merge pull request #47 from probcomp/gm/task_scripts

Progress on Tracks->Segmentation and Video->[Tracks+Segmentation] Tasks, and script to visualize all tasks+solvers

Author: georgematheos

.gitignore

@@ -1,5 +1,6 @@
 *.egg-info
 *.pyc
 *.png
-assets/*
+assets/shared_data_bucket/*
+assets/test_results/*
 **/.ipynb_checkpoints

b3d/chisight/dense/model.py

@@ -6,6 +6,7 @@
 import b3d.utils as utils
 import b3d.chisight.dense.differentiable_renderer as rendering
 import rerun as rr
+import numpy as np
 
 def uniformpose_meshes_to_image_model__factory(likelihood):
     """
@@ -77,15 +78,15 @@ def rr_log_uniformpose_meshes_to_image_model_trace(trace, renderer, **kwargs):
     """
     return rr_log_meshes_to_image_model_trace(trace, renderer, **kwargs,
                                               model_args_to_densemodel_args=(
-        lambda args: (trace.get_choices()["camera_pose"], trace.get_choices()["poses"], *args)
+        lambda args: (trace.get_choices()["camera_pose"], trace.get_choices()("poses").c.v, *args)
     ))
 
 def rr_log_meshes_to_image_model_trace(
         trace, renderer,
         prefix="trace",
         timeless=False,
         model_args_to_densemodel_args=(lambda x: x),
-        transform=Pose.identity()
+        transform_Viz_Trace=Pose.identity()
     ):
     """
     Log to rerun a visualization of a trace from `meshes_to_image_model`.
@@ -94,11 +95,16 @@ def rr_log_meshes_to_image_model_trace(
     to visualize traces from other models that have the same return value as `meshes_to_image_model`.
     This function will call `model_args_to_densemodel_args` on the arguments of the given trace,
     and should produce arguments of the form accepted by `meshes_to_image_model`.
+
+    The argument `transform_Viz_Trace` can be used to visualize the trace at a transformed
+    coordinate frame.  `transform_Viz_Trace` is a Pose object so that for a 3D point
+    `point_Trace` in the trace, `transform_Viz_Trace.apply(point_Trace)` is the corresponding
+    3D point in the visualizer.
     """
     # 2D:
     (observed_rgbd, metadata) = trace.get_retval()
-    rr.log(f"/{prefix}/rgb/observed", rr.Image(observed_rgbd[:, :, :3]), timeless=timeless)
-    rr.log(f"/{prefix}/depth/observed", rr.DepthImage(observed_rgbd[:, :, 3]), timeless=timeless)
+    rr.log(f"/{prefix}/rgb/observed", rr.Image(np.array(observed_rgbd[:, :, :3])), timeless=timeless)
+    rr.log(f"/{prefix}/depth/observed", rr.DepthImage(np.array(observed_rgbd[:, :, 3])), timeless=timeless)
 
     # Visualization path for the average render,
     # if the likelihood metadata contains the output of the differentiable renderer.
@@ -107,46 +113,45 @@ def rr_log_meshes_to_image_model_trace(
         avg_obs = rendering.dist_params_to_average(weights, attributes, jnp.zeros(4))
         avg_obs_rgb_clipped = jnp.clip(avg_obs[:, :, :3], 0, 1)
         avg_obs_depth_clipped = jnp.clip(avg_obs[:, :, 3], 0, 1)
-        rr.log(f"/{prefix}/rgb/average_render", rr.Image(avg_obs_rgb_clipped), timeless=timeless)
-        rr.log(f"/{prefix}/depth/average_render", rr.DepthImage(avg_obs_depth_clipped), timeless=timeless)
+        rr.log(f"/{prefix}/rgb/average_render", rr.Image(np.array(avg_obs_rgb_clipped)), timeless=timeless)
+        rr.log(f"/{prefix}/depth/average_render", rr.DepthImage(np.array(avg_obs_depth_clipped)), timeless=timeless)
 
     # 3D:
-    rr.log(f"/{prefix}", rr.Transform3D(translation=transform.pos, mat3x3=transform.rot.as_matrix()), timeless=timeless)
+    rr.log(f"/{prefix}/3D/", rr.Transform3D(translation=transform_Viz_Trace.pos, mat3x3=transform_Viz_Trace.rot.as_matrix()), timeless=timeless)
 
     (X_WC, Xs_WO, vertices_O, faces, vertex_colors) = model_args_to_densemodel_args(trace.get_args())
-    Xs_WO = trace.strip()["poses"].inner.value # TODO: do this better
     vertices_W = jax.vmap(lambda X_WO, v_O: X_WO.apply(v_O), in_axes=(0, 0))(Xs_WO, vertices_O)
     N = vertices_O.shape[0]
     f = jax.vmap(lambda i, f: f + i*vertices_O.shape[1], in_axes=(0, 0))(jnp.arange(N), faces)
     f = f.reshape(-1, 3)
 
-    rr.log(f"/{prefix}/mesh", rr.Mesh3D(
-        vertex_positions=vertices_W.reshape(-1, 3),
-        triangle_indices=f,
-        vertex_colors=vertex_colors.reshape(-1, 3)
+    rr.log(f"/{prefix}/3D/mesh", rr.Mesh3D(
+        vertex_positions=np.array(vertices_W.reshape(-1, 3)),
+        triangle_indices=np.array(f),
+        vertex_colors=np.array(vertex_colors.reshape(-1, 3))
     ), timeless=timeless)
 
-    rr.log(f"/{prefix}/camera",
+    rr.log(f"/{prefix}/3D/camera",
         rr.Pinhole(
             focal_length=[float(renderer.fx), float(renderer.fy)],
             width=renderer.width,
             height=renderer.height,
             principal_point=jnp.array([renderer.cx, renderer.cy]),
             ), timeless=timeless
         )
-    rr.log(f"/{prefix}/camera", rr.Transform3D(translation=X_WC.pos, mat3x3=X_WC.rot.as_matrix()), timeless=timeless)
+    rr.log(f"/{prefix}/3D/camera", rr.Transform3D(translation=X_WC.pos, mat3x3=X_WC.rot.as_matrix()), timeless=timeless)
     xyzs_C = utils.xyz_from_depth(observed_rgbd[:, :, 3], renderer.fx, renderer.fy, renderer.cx, renderer.cy)
     xyzs_W = X_WC.apply(xyzs_C)
-    rr.log(f"/{prefix}/gt_pointcloud", rr.Points3D(
-        positions=xyzs_W.reshape(-1,3),
-        colors=observed_rgbd[:, :, :3].reshape(-1,3),
-        radii = 0.001*jnp.ones(xyzs_W.reshape(-1,3).shape[0])),
+    rr.log(f"/{prefix}/3D/gt_pointcloud", rr.Points3D(
+        positions=np.array(xyzs_W.reshape(-1,3)),
+        colors=np.array(observed_rgbd[:, :, :3].reshape(-1,3)),
+        radii = 0.001*np.ones(xyzs_W.reshape(-1,3).shape[0])),
         timeless=timeless
     )
 
     patch_centers_W = jax.vmap(lambda X_WO: X_WO.pos)(Xs_WO)
     rr.log(
-        f"/{prefix}/patch_centers_W",
-        rr.Points3D(positions=patch_centers_W, colors=jnp.array([0., 0., 1.]), radii=0.003),
+        f"/{prefix}/3D/patch_centers",
+        rr.Points3D(positions=np.array(patch_centers_W), colors=np.array([0., 0., 1.]), radii=0.003),
         timeless=timeless
     )

b3d/chisight/dense/patch_tracking.py

@@ -3,6 +3,7 @@
 from b3d import Pose
 import b3d
 import genjax
+from genjax import ChoiceMapBuilder as C
 import b3d.chisight.dense.differentiable_renderer as r
 import b3d.chisight.dense.model as m
 import b3d.chisight.dense.likelihoods as likelihoods
@@ -78,15 +79,10 @@ def get_adam_optimization_patch_tracker(model, patch_vertices_P, patch_faces, pa
     def importance_from_pos_quat(positions, quaternions, observed_rgbd):
         key = jax.random.PRNGKey(0) # This value shouldn't matter, in the current model version.
         poses = jax.vmap(lambda pos, quat: Pose.from_vec(jnp.concatenate([pos, quat])), in_axes=(0, 0))(positions, quaternions)
-        trace, weight = model.importance(
-            key,
-            genjax.ChoiceMap.d({
-                "poses": genjax.ChoiceMap.idx(jnp.arange(positions.shape[0]), poses),
-                "camera_pose": X_WC,
-                "observed_image": {"observed_image": {"obs": observed_rgbd}}
-            }),
-            (patch_vertices_P, patch_faces, patch_vertex_colors)
-        )
+        cm = jax.vmap(lambda i: C["poses", i].set(poses[i]))(jnp.arange(poses.shape[0]))
+        cm = cm.merge(C["camera_pose"].set(b3d.Pose.identity()))
+        cm = cm.merge(C["observed_image", "observed_image", "obs"].set(observed_rgbd))
+        trace, weight = model.importance(key, cm, (patch_vertices_P, patch_faces, patch_vertex_colors))
         return trace, weight
 
     def weight_from_pos_quat(pos, quat, observed_rgbd):

tests/common/solver.py

@@ -22,4 +22,11 @@ def visualize_solver_state(self):
         Visualize any information recorded by the solver during the last call to `solve`.
         This may log data to rerun, produce pyplots, etc.
         """
-        pass
+        pass
+
+    @property
+    def name(self) -> str:
+        """
+        Returns the name of the solver.
+        """
+        return self.__class__.__name__

tests/common/task.py

@@ -7,11 +7,16 @@ class Task:
     - `score`
 
     Optionally, task developers may also implement the following methods:
+    - @property `name`
     - `assert_passing`
     - `visualize_task`
     - `visualize_solution`
 
     The `run_and_score` and `run_tests` methods are automatically implemented from the above.
+
+    It is recommended that `Task` constrution should be very fast, and that
+    data-loading or other slow operations should be performed lazily
+    at the first call to `get_task_specification` or `visualize_task`.
     """
 
     ### CORE TASK INTERFACE ###
@@ -38,6 +43,13 @@ def score(self, solution, **kwargs) -> "Any":
         """
         raise NotImplementedError()
 
+    @property
+    def name(self) -> str:
+        """
+        Returns the name of the task.
+        """
+        return self.__class__.__name__
+    
     def assert_passing(self, scores, **kwargs) -> None:
         """
         Takes the output of `score` and makes assertions about the scores,
@@ -57,7 +69,7 @@ def assert_passing(self, scores, **kwargs) -> None:
         according to the default tolerances).
         """
         raise NotImplementedError()
-    
+
     def visualize_task(self):
         """
         Visualize the task (but not the solution).
@@ -93,6 +105,7 @@ def run_and_score(self, solver, viz=False, **kwargs) -> dict:
         metrics = self.score(task_output, **kwargs)
         if viz:
             self.visualize_solution(task_output, metrics)
+            solver.visualize_solver_state(task_spec)
         return metrics
 
     def run_tests(self, solver, viz=False, **kwargs) -> None:
@@ -111,3 +124,18 @@ def run_tests(self, solver, viz=False, **kwargs) -> None:
         # Score and assess whether passing.
         metrics = self.run_and_score(solver, viz=viz, **kwargs)
         self.assert_passing(metrics, **kwargs)
+
+    def run_solver_and_make_all_visualizations(self, solver):
+        """
+        Run the solver and make all visualizations.
+        """
+        task_spec = self.get_task_specification()
+        self.visualize_task()
+        task_output = solver.solve(task_spec)
+        metrics = self.score(task_output)
+        solver.visualize_solver_state(task_spec)
+        self.visualize_solution(task_output, metrics)
+        return metrics
+
+    def __repr__(self):
+        return f"{self.name}()"

tests/dense_model_unit_tests/triangle_depth_posterior/test_triangle_depth_posterior.py

@@ -10,7 +10,9 @@
     for triangle_color in [jnp.array([1., 0., 0.]), jnp.array([0., 1., 0.])]
 ]
 
-@pytest.mark.parametrize("task_spec", task_specs[:3])
+# Only run one test for now, to prevent issues due to the current
+# memory leak in the renderer.
+@pytest.mark.parametrize("task_spec", task_specs[:1])
 def test(task_spec):
     task = TrianglePosteriorGridApproximationTask.default_scene_using_colors(*task_spec)
     task.run_tests(

tests/dense_model_unit_tests/triangle_depth_posterior/visualize.ipynb

@@ -59,17 +59,19 @@ def get_cheezitbox_scene_loader(n_frames=30):
 
 # Scene manually constructed in Python: rotating cheezit box
 def ftd_from_rotating_cheezit_box(n_frames=30):
-    (r, centers_2D_frame_0, centers_3D_W_over_time, Xs_WC, observed_rgbds) = load_rotating_cheezit_box_data(n_frames)
+    (r, centers_2D_frame_0, centers_3D_W_over_time, poses_WC, observed_rgbds) = load_rotating_cheezit_box_data(n_frames)
     return b3d.io.FeatureTrackData(
         observed_keypoints_positions=jax.vmap(lambda positions_3D_W, X_WC: b3d.xyz_to_pixel_coordinates(
                 X_WC.inv().apply(positions_3D_W), r.fx, r.fy, r.cx, r.cy
-        ), in_axes=(0, 0))(centers_3D_W_over_time, Xs_WC),
+        ), in_axes=(0, 0))(centers_3D_W_over_time, poses_WC),
         keypoint_visibility=jnp.ones((n_frames, centers_2D_frame_0.shape[0]), dtype=bool),
         camera_intrinsics=r.get_intrinsics_object().as_array(),
         rgbd_images=observed_rgbds,
         latent_keypoint_positions=centers_3D_W_over_time,
-        camera_position=Xs_WC.pos,
-        camera_quaternion=Xs_WC.xyzw
+        camera_position=poses_WC.pos,
+        camera_quaternion=poses_WC.xyzw,
+        # Every point is assigned to one object (the cheez-it box)
+        object_assignments=jnp.zeros(centers_2D_frame_0.shape[0], dtype=int)
     )
 
 def load_rotating_cheezit_box_data(n_frames=30):
@@ -119,9 +121,9 @@ def load_rotating_cheezit_box_data(n_frames=30):
         lambda X_W_Bt: X_W_Bt.apply(centers_3D_B0)
     )(box_poses_W)
 
-    Xs_WC = jax.vmap(lambda x: X_WC)(jnp.arange(n_frames))
+    poses_WC = jax.vmap(lambda x: X_WC)(jnp.arange(n_frames))
 
-    return (renderer, centers_2D_frame_0, centers_3D_W_over_time, Xs_WC, observed_rgbds)
+    return (renderer, centers_2D_frame_0, centers_3D_W_over_time, poses_WC, observed_rgbds)
 
 ### Utils ###
 

tests/sama4d/data_curation.py

@@ -0,0 +1,14 @@
+import jax.numpy as jnp
+from tests.common.solver import Solver
+
+class DummyTracksToSegmentationSolver(Solver):
+    def solve(self, task_spec):
+        # assign every keypoint to the same object
+        # (called object 0)
+        return jnp.zeros(
+            task_spec["keypoint_tracks_2D"].shape[1],
+            dtype=int
+        )
+    
+    def visualize_solver_state(self, task_spec):
+        pass