swap in previous pose at best location in array

run.py

@@ -1,15 +1,15 @@
 import os
+import uuid
 from os import listdir
 from os.path import isfile, join
-import uuid
 
 data_path = "/home/haoliangwang/data/"
 hdf5_file_path = os.path.join(
     data_path,
     "physion_hdf5",
 )
 
-for scenario in ['collide', 'drop', 'roll', 'dominoes', 'support', 'link', 'contain']:
+for scenario in ["collide", "link", "drop", "roll", "dominoes", "support", "contain"]:
     scenario_path = join(hdf5_file_path, scenario + "_all_movies")
     onlyhdf5 = [
         f
@@ -26,8 +26,8 @@
     viz_index = 0
     for trial_index, hdf5_file in enumerate(onlyhdf5):
         trial_name = hdf5_file[:-5]
-        if scenario != 'roll' or trial_index < 20:
-            continue
         print(trial_index + 1, "\t", trial_name)
-        os.system(f"python /home/haoliangwang/b3d/test_b3d_tracking_hmm_single.py --scenario {scenario} --trial_name {trial_name} --recording_id {recording_id} --viz_index {viz_index}")
-        viz_index += FINAL_T+1
+        os.system(
+            f"python /home/haoliangwang/b3d/test_b3d_tracking_hmm_single.py --scenario {scenario} --trial_name {trial_name} --recording_id {recording_id} --viz_index {viz_index}"
+        )
+        viz_index += FINAL_T + 1

src/b3d/chisight/dense/dense_model.py

@@ -128,10 +128,18 @@ def dense_multiobject_model(
             latent_rgbd = jnp.flip(latent_rgbd, 1)
             bg_rgb = background[..., :3]
             bg_d = background[..., 3:]
-            latent_rgb = jnp.where(bg_rgb == jnp.array([jnp.inf, jnp.inf, jnp.inf]), latent_rgbd[..., 0:3], bg_rgb)
-            latent_d = jnp.minimum(jnp.where(latent_rgbd[..., 3:] == 0.0, 10, latent_rgbd[..., 3:]), bg_d)
-
-            likelihood_args["latent_rgbd"] = jnp.concatenate([latent_rgb, latent_d], axis=-1)
+            latent_rgb = jnp.where(
+                bg_rgb == jnp.array([jnp.inf, jnp.inf, jnp.inf]),
+                latent_rgbd[..., 0:3],
+                bg_rgb,
+            )
+            latent_d = jnp.minimum(
+                jnp.where(latent_rgbd[..., 3:] == 0.0, 10, latent_rgbd[..., 3:]), bg_d
+            )
+
+            likelihood_args["latent_rgbd"] = jnp.concatenate(
+                [latent_rgb, latent_d], axis=-1
+            )
             likelihood_args["rasterize_results"] = rasterize_results
 
         image = image_likelihood(likelihood_args) @ "rgbd"

src/b3d/chisight/dense/likelihoods/laplace_likelihood_blur.py

@@ -1,8 +1,8 @@
 import functools
 
+import genjax
 import jax
 import jax.numpy as jnp
-import genjax
 
 from b3d.modeling_utils import get_interpenetration
 
@@ -99,9 +99,8 @@
         )
         return log_probabilities
 
     scores = likelihood_per_pixel(observed_rgbd, latent_rgbd, likelihood_args["blur"])
 
-
     valid_window = latent_rgbd[..., 0:3].sum(axis=2) > 0.0  # latent_rgbd[..., 3] > 0.0
     if likelihood_args["masked"].unwrap():
         observed_window = observed_rgbd[..., 0:3].sum(axis=2) > 0.0

src/b3d/chisight/gen3d/dataloading.py

@@ -55,7 +55,9 @@ def load_trial(hdf5_file_path, FINAL_T):
         )
         assert len(object_ids) == len(object_segmentation_colors)
 
-        background_areas = np.zeros((image_arr.shape[0], image_arr.shape[1], image_arr.shape[2]))
+        background_areas = np.zeros(
+            (image_arr.shape[0], image_arr.shape[1], image_arr.shape[2])
+        )
 
         distractors = (
             np.array(f["static"]["distractors"])
@@ -69,9 +71,14 @@ def load_trial(hdf5_file_path, FINAL_T):
         )
         distractors_occluders = np.concatenate([distractors, occluders])
         if len(distractors_occluders):
-            background_areas = jnp.asarray([get_mask_area(seg, object_segmentation_colors[
-                -len(distractors_occluders):
-            ]) for seg in seg_arr])
+            background_areas = jnp.asarray(
+                [
+                    get_mask_area(
+                        seg, object_segmentation_colors[-len(distractors_occluders) :]
+                    )
+                    for seg in seg_arr
+                ]
+            )
             object_ids = object_ids[: -len(distractors_occluders)]
             object_segmentation_colors = object_segmentation_colors[
                 : -len(distractors_occluders)
@@ -178,13 +185,13 @@ def get_initial_state(
         )
         initial_state[f"object_scale_{o_id}_0"] = jnp.array(pred[i]["scale"][0])
         hyperparams["meshes"][int(o_id)] = [
-                b3d.Mesh(
-                    meshes[pred[i]["type"][0]].vertices,
-                    meshes[pred[i]["type"][0]].faces,
-                    jnp.ones(meshes[pred[i]["type"][0]].vertices.shape)
-                    * mean_object_colors,
-                )
-            ]
+            b3d.Mesh(
+                meshes[pred[i]["type"][0]].vertices,
+                meshes[pred[i]["type"][0]].faces,
+                jnp.ones(meshes[pred[i]["type"][0]].vertices.shape)
+                * mean_object_colors,
+            )
+        ]
 
     hyperparams["object_ids"] = Pytree.const([o_id for o_id in object_ids])
     return initial_state, hyperparams

src/b3d/chisight/gen3d/datawriting.py

@@ -1,19 +1,20 @@
-from scipy.spatial.transform import Rotation
+import json
+import os
 from copy import deepcopy
-import trimesh
+
 import numpy as np
-import os
-import json
+import trimesh
+from scipy.spatial.transform import Rotation
 
 import b3d
 
-
 NUM_SAMPLE_FROM_POSTERIOR = 20
 SMOOTHING_WINDOW_SIZE = 3
 FPS = 100
 STATIC_POSITION_THRESHHOLD = 0.007
 STATIC_ROTATION_THRESHHOLD = 0.001
 
+
 def mkdir(path):
     if not os.path.isdir(path):
         os.makedirs(path)
@@ -87,9 +88,9 @@ def compute_angular_velocity(q1, q2, dt):
 
 def scale_mesh(vertices, scale_factor):
     vertices_copy = deepcopy(vertices)
-    vertices_copy[:, 0] *= scale_factor['x']
-    vertices_copy[:, 1] *= scale_factor['y']
-    vertices_copy[:, 2] *= scale_factor['z']
+    vertices_copy[:, 0] *= scale_factor["x"]
+    vertices_copy[:, 1] *= scale_factor["y"]
+    vertices_copy[:, 2] *= scale_factor["z"]
     return vertices_copy
 
 
@@ -135,41 +136,47 @@ def scale_mesh(vertices, scale_factor):
 #                 * composite_scales[composite_id][i]["y"]
 #             )
 
+
 def sample_from_posterior(log_probs_categories, option="rank"):
-        log_probs = [item[0] for item in log_probs_categories]
-        categories = [item[1] for item in log_probs_categories]
-        num_categories = len(log_probs)
-
-        if option == "uniform":
-            def draw_single_sample():
-                index = np.random.choice(num_categories)
-                return categories[index]
-        elif option == "veridical":
-            def draw_single_sample():
-                # see this: https://stackoverflow.com/questions/58339083/how-to-sample-from-a-log-probability-distribution
-                gumbels = np.random.gumbel(size=num_categories)
-                index = np.argmax(log_probs + gumbels)
-                return categories[index]
-        elif option == "rank":
-            def draw_single_sample():
-                weights = np.array([1 / (n + 1) for n in range(num_categories)])
-                weights_norm = weights / weights.sum()
-                index = np.random.choice(num_categories, p=weights_norm)
-                return categories[index]
-        elif option == "mix":
-            def draw_single_sample():
-                t = 0.5
-                t * np.array(log_probs) + (1 - t) * (1 / num_categories)
-                return
-        else:
-            raise NotImplementedError
+    log_probs = [item[0] for item in log_probs_categories]
+    categories = [item[1] for item in log_probs_categories]
+    num_categories = len(log_probs)
+
+    if option == "uniform":
+
+        def draw_single_sample():
+            index = np.random.choice(num_categories)
+            return categories[index]
+    elif option == "veridical":
+
+        def draw_single_sample():
+            # see this: https://stackoverflow.com/questions/58339083/how-to-sample-from-a-log-probability-distribution
+            gumbels = np.random.gumbel(size=num_categories)
+            index = np.argmax(log_probs + gumbels)
+            return categories[index]
+    elif option == "rank":
+
+        def draw_single_sample():
+            weights = np.array([1 / (n + 1) for n in range(num_categories)])
+            weights_norm = weights / weights.sum()
+            index = np.random.choice(num_categories, p=weights_norm)
+            return categories[index]
+    elif option == "mix":
+
+        def draw_single_sample():
+            t = 0.5
+            t * np.array(log_probs) + (1 - t) * (1 / num_categories)
+            return
+    else:
+        raise NotImplementedError
+
+    samples = []
+    np.random.seed(42)
+    for _ in range(NUM_SAMPLE_FROM_POSTERIOR):
+        sample = draw_single_sample()
+        samples.append(sample)
+    return samples
 
-        samples = []
-        np.random.seed(42)
-        for _ in range(NUM_SAMPLE_FROM_POSTERIOR):
-            sample = draw_single_sample()
-            samples.append(sample)
-        return samples
 
 def get_posterior_poses_for_frame(frame, posterior_across_frames):
     pose_samples_from_posterior = {}
@@ -189,7 +196,15 @@ def get_posterior_poses_for_frame(frame, posterior_across_frames):
     return pose_samples_from_posterior
 
 
-def write_json(pred_file, hyperparams, posterior_across_frames, save_path, scenario, trial_name, debug=False):
+def write_json(
+    pred_file,
+    hyperparams,
+    posterior_across_frames,
+    save_path,
+    scenario,
+    trial_name,
+    debug=False,
+):
     pred = pred_file["scene"][0]["objects"]
 
     # prepare the json file to write
@@ -199,8 +214,7 @@ def write_json(pred_file, hyperparams, posterior_across_frames, save_path, scena
 
     for i, o_id in enumerate(hyperparams["object_ids"].unwrap()):
         json_file["model"][int(o_id)] = [
-            pred[i]["type"][0]
-            for _ in range(NUM_SAMPLE_FROM_POSTERIOR)
+            pred[i]["type"][0] for _ in range(NUM_SAMPLE_FROM_POSTERIOR)
         ]
         json_file["scale"][int(o_id)] = [
             {
@@ -230,9 +244,7 @@ def write_json(pred_file, hyperparams, posterior_across_frames, save_path, scena
                 [
                     {
                         "x": pose._position[0].astype(float).item(),
-                        "y": pose._position[1].astype(float).item()
-                        if pose._position[1].astype(float).item() >= 0
-                        else 0,
+                        "y": pose._position[1].astype(float).item(),
                         "z": pose._position[2].astype(float).item(),
                     }
                     for pose in poses[0]
@@ -262,48 +274,73 @@ def write_json(pred_file, hyperparams, posterior_across_frames, save_path, scena
     linear_velocity_dict = {}
     linear_velocity_dict_optim = {}
     for o_id in pose_samples_from_posterior_last_frame.keys():
-        if np.allclose(pose_samples_from_posterior_last_frame[o_id][-1]._position, pose_samples_from_posterior_window_frame[o_id][-1]._position, atol=STATIC_POSITION_THRESHHOLD*SMOOTHING_WINDOW_SIZE):
-            linear_velocity_dict[int(o_id)] = [{"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)]
-            linear_velocity_dict_optim[int(o_id)] = [{"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)]
+        if np.allclose(
+            pose_samples_from_posterior_last_frame[o_id][-1]._position,
+            pose_samples_from_posterior_window_frame[o_id][-1]._position,
+            atol=STATIC_POSITION_THRESHHOLD * SMOOTHING_WINDOW_SIZE,
+        ):
+            linear_velocity_dict[int(o_id)] = [
+                {"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)
+            ]
+            linear_velocity_dict_optim[int(o_id)] = [
+                {"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)
+            ]
         else:
             linear_velocity_dict[int(o_id)] = [
                 compute_linear_velocity(
                     hyperparams["meshes"][int(o_id)][0],
                     json_file["scale"][o_id][i],
                     pose_samples_from_posterior_last_frame[o_id][0][i],
-                    pose_samples_from_posterior_window_frame[o_id][-1],  # using optim pose for window frame
+                    pose_samples_from_posterior_window_frame[o_id][
+                        -1
+                    ],  # using optim pose for window frame
                     SMOOTHING_WINDOW_SIZE / FPS,
                 )
-                for i in range(NUM_SAMPLE_FROM_POSTERIOR)]
+                for i in range(NUM_SAMPLE_FROM_POSTERIOR)
+            ]
             linear_velocity_dict_optim[int(o_id)] = [
                 compute_linear_velocity(
                     hyperparams["meshes"][int(o_id)][0],
                     json_file["scale"][o_id][i],
                     pose_samples_from_posterior_last_frame[o_id][-1],
-                    pose_samples_from_posterior_window_frame[o_id][-1],  # using optim pose for window frame
+                    pose_samples_from_posterior_window_frame[o_id][
+                        -1
+                    ],  # using optim pose for window frame
                     SMOOTHING_WINDOW_SIZE / FPS,
                 )
-                for i in range(NUM_SAMPLE_FROM_POSTERIOR)]
-
+                for i in range(NUM_SAMPLE_FROM_POSTERIOR)
+            ]
 
     angular_velocity_dict = {}
     angular_velocity_dict_optim = {}
     for o_id in pose_samples_from_posterior_last_frame.keys():
-        if np.allclose(pose_samples_from_posterior_last_frame[o_id][-1]._quaternion, pose_samples_from_posterior_window_frame[o_id][-1]._quaternion, atol=STATIC_ROTATION_THRESHHOLD*SMOOTHING_WINDOW_SIZE):
-            angular_velocity_dict[int(o_id)] = [{"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)]
-            angular_velocity_dict_optim[int(o_id)] = [{"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)]
+        if np.allclose(
+            pose_samples_from_posterior_last_frame[o_id][-1]._quaternion,
+            pose_samples_from_posterior_window_frame[o_id][-1]._quaternion,
+            atol=STATIC_ROTATION_THRESHHOLD * SMOOTHING_WINDOW_SIZE,
+        ):
+            angular_velocity_dict[int(o_id)] = [
+                {"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)
+            ]
+            angular_velocity_dict_optim[int(o_id)] = [
+                {"x": 0, "y": 0, "z": 0} for _ in range(NUM_SAMPLE_FROM_POSTERIOR)
+            ]
         else:
             angular_velocity_dict[int(o_id)] = [
                 compute_angular_velocity(
-                    pose_samples_from_posterior_window_frame[o_id][-1]._quaternion,  # using optim pose for window frame
+                    pose_samples_from_posterior_window_frame[o_id][
+                        -1
+                    ]._quaternion,  # using optim pose for window frame
                     pose_samples_from_posterior_last_frame[o_id][0][i]._quaternion,
                     SMOOTHING_WINDOW_SIZE / FPS,
                 )
                 for i in range(NUM_SAMPLE_FROM_POSTERIOR)
             ]
             angular_velocity_dict_optim[int(o_id)] = [
                 compute_angular_velocity(
-                    pose_samples_from_posterior_window_frame[o_id][-1]._quaternion,  # using optim pose for window frame
+                    pose_samples_from_posterior_window_frame[o_id][
+                        -1
+                    ]._quaternion,  # using optim pose for window frame
                     pose_samples_from_posterior_last_frame[o_id][-1]._quaternion,
                     SMOOTHING_WINDOW_SIZE / FPS,
                 )
@@ -314,14 +351,16 @@ def write_json(pred_file, hyperparams, posterior_across_frames, save_path, scena
     json_file["velocity"] = linear_velocity_dict
     json_file["angular_velocity"] = angular_velocity_dict
 
-    missing = find_missing_values(np.array([int(item) for item in hyperparams["object_ids"].unwrap()]))
+    missing = find_missing_values(
+        np.array([int(item) for item in hyperparams["object_ids"].unwrap()])
+    )
     for feature, val in json_file.items():
         for o_id in missing:
             json_file[feature][o_id] = val[int(hyperparams["object_ids"].unwrap()[0])]
 
     json_file_optim = deepcopy(json_file)
-    json_file_optim['velocity'] = linear_velocity_dict_optim
-    json_file_optim['angular_velocity'] = angular_velocity_dict_optim
+    json_file_optim["velocity"] = linear_velocity_dict_optim
+    json_file_optim["angular_velocity"] = angular_velocity_dict_optim
 
     mkdir(f"{save_path}/{scenario}/")
     with open(f"{save_path}/{scenario}/{trial_name}.json", "w") as f:
@@ -343,8 +382,8 @@ def write_json(pred_file, hyperparams, posterior_across_frames, save_path, scena
                         rank[0].astype(float).item(),
                         rank[1]._position.astype(float).tolist(),
                         rank[1]._quaternion.astype(float).tolist(),
-                        )
+                    )
         mkdir(f"{save_path}/{scenario}_verbose/")
         with open(f"{save_path}/{scenario}_verbose/{trial_name}.json", "w") as f:
             json.dump(posterior_across_frames, f)
-    return
+    return