Add likelihood experiments to main repo

.github/workflows/ruff.yml

@@ -0,0 +1,22 @@
+name: Ruff linter
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  ruff:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v4
+      with:
+        python-version: 3.11.5
+
+    - uses: chartboost/ruff-action@v1
+      with:
+        version: 0.5.4
+        args: check --output-format github

.gitignore

@@ -13,4 +13,8 @@ assets/ycbineoat/*
 **/**.json
 **/**.npy
 **/**.txt
+!src/b3d/renderer/**/**.txt
 assets/kitti/*
+**/**.pth
+**/**.obj
+**/**.r3d

.pre-commit-config.yaml

@@ -1,3 +1,6 @@
+ci:
+    skip: [ruff]
+
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
@@ -8,8 +11,11 @@ repos:
       - id: trailing-whitespace
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.2
+    rev: v0.5.5
     hooks:
+      - id: ruff
+        types_or: [ python, pyi ]
+        args: [ --fix ]
+
       - id: ruff-format
-        args: [ --exclude, kitti.ipynb ]
-        types_or: [ python, pyi, jupyter ]
+        types_or: [ python, pyi ]

README.md

@@ -54,7 +54,7 @@ export XLA_PYTHON_CLIENT_ALLOCATOR="platform"
 - Check if you can access the [GenJAX Users
   Group](https://groups.google.com/u/1/a/chi-fro.org/g/genjax-users), and if not, run `\invite-genjax <google-account-email>` in any channel in the the probcomp Slack
 - [Install the Google Cloud command line tools](https://cloud.google.com/sdk/docs/install).
-- Run `gcloud auth application-default login`.  (This command needs to be rerun ever time your machine reboots.)
+- Run `gcloud auth application-default login`.  (This command needs to be rerun every time your machine reboots.)
 
 Next, on the machine, close the `b3d` repo:
 

demo.py

@@ -1,18 +1,16 @@
 #!/usr/bin/env python
-import jax.numpy as jnp
-import jax
-import matplotlib.pyplot as plt
-import numpy as np
 import os
-import trimesh
+
 import b3d
 import b3d.bayes3d as bayes3d
-from jax.scipy.spatial.transform import Rotation as Rot
-from b3d import Pose
+import fire
 import genjax
+import jax
+import jax.numpy as jnp
+import numpy as np
 import rerun as rr
+from b3d import Pose
 from tqdm import tqdm
-import fire
 
 
 def test_demo():
@@ -53,11 +51,9 @@ def test_demo():
         1.0,
     )
 
-    num_layers = 2048
     renderer = b3d.Renderer(image_width, image_height, fx, fy, cx, cy, near, far)
     model = bayes3d.model_multiobject_gl_factory(renderer)
     importance_jit = jax.jit(model.importance)
-    update_jit = jax.jit(model.update)
 
     # Arguments of the generative model.
     # These control the inlier / outlier decision boundary for color error and depth error.
@@ -118,7 +114,7 @@ def test_demo():
 
     # `make_mesh_from_point_cloud_and_resolution` takes a 3D positions, colors, and sizes of the boxes that we want
     # to place at each position and create a mesh
-    vertices, faces, vertex_colors, face_colors = (
+    vertices, faces, vertex_colors, _face_colors = (
         b3d.make_mesh_from_point_cloud_and_resolution(
             point_cloud,
             colors,
@@ -167,7 +163,7 @@ def test_demo():
         bayes3d.rerun_visualize_trace_t(trace, T_observed_image)
 
     # Outliers are AND of the RGB and Depth outlier masks
-    inliers, color_inliers, depth_inliers, outliers, undecided, valid_data_mask = (
+    _inliers, _color_inliers, _depth_inliers, outliers, _undecided, _valid_data_mask = (
         bayes3d.get_rgb_depth_inliers_from_trace(trace)
     )
     outlier_mask = outliers
@@ -186,7 +182,7 @@ def test_demo():
     point_cloud_colors = point_cloud_colors.reshape(-1, 3)[assignment == 0]
 
     # Create new mesh.
-    vertices, faces, vertex_colors, face_colors = (
+    vertices, faces, vertex_colors, _face_colors = (
         b3d.make_mesh_from_point_cloud_and_resolution(
             point_cloud, point_cloud_colors, point_cloud[:, 2] / fx * 2.0
         )

demos/detector_segmenter.py

@@ -1,10 +1,23 @@
-import requests
-from PIL import Image
-import torch
-import b3d
+import io
 import os
 
-from transformers import OwlViTProcessor, OwlViTForObjectDetection
+import b3d
+import jax
+import jax.numpy as jnp
+import numpy
+import numpy as np
+import rerun as rr
+import torch
+from PIL import Image
+from transformers import (
+    AutoModelForZeroShotObjectDetection,
+    AutoProcessor,
+    DetrFeatureExtractor,
+    DetrForSegmentation,
+    OwlViTForObjectDetection,
+    OwlViTProcessor,
+)
+from transformers.models.detr.feature_extraction_detr import rgb_to_id
 
 processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
 model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
@@ -17,12 +30,10 @@
     )
 )
 
-import numpy as np
 
 image = Image.fromarray(np.array(video_input.rgb[0]))
 image.save("test.png")
 
-from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
 
 model_id = "IDEA-Research/grounding-dino-base"
 
@@ -45,34 +56,16 @@
 )
 
 
-import rerun as rr
-
 PORT = 8812
 rr.init("real")
 rr.connect(addr=f"127.0.0.1:{PORT}")
 
 rr.log("image", rr.Image(np.array(image)))
 
-import io
-import requests
-from PIL import Image
-import torch
-import numpy
-import jax.numpy as jnp
-import numpy as np
-import jax
-import rerun as rr
-
-from transformers import DetrFeatureExtractor, DetrForSegmentation
-from transformers.models.detr.feature_extraction_detr import rgb_to_id
-
 
 rr.init("segmentation")
 rr.connect("127.0.0.1:8812")
 
-import os
-import b3d
-
 # Load date
 # path = os.path.join(b3d.get_root_path(),
 # "assets/shared_data_bucket/input_data/shout_on_desk.r3d.video_input.npz")

demos/differentiable_renderer/dense_fitting/standalone_nishad_demo.py

@@ -1,17 +1,14 @@
 # Note: this uses a custom differentiable renderer defined inline in this script.
 
-import jax.numpy as jnp
-import jax
-import matplotlib.pyplot as plt
-import numpy as np
+import functools
 import os
-import trimesh
+
 import b3d
-from jax.scipy.spatial.transform import Rotation as Rot
-from b3d import Pose
+import jax
+import jax.numpy as jnp
+import numpy as np
 import rerun as rr
-import functools
-import genjax
+from b3d import Pose
 
 rr.init("gradients")
 rr.connect("127.0.0.1:8812")
@@ -156,7 +153,7 @@ def render(particle_centers, particle_widths, particle_colors):
     vertex_colors = vertex_colors.reshape(-1, 3)
     triangle_index_to_particle_index = triangle_index_to_particle_index.reshape(-1)
 
-    uvs, _, triangle_id_image, depth_image = renderer.rasterize(
+    _uvs, _, triangle_id_image, depth_image = renderer.rasterize(
         Pose.identity()[None, ...], vertices, faces, jnp.array([[0, len(faces)]])
     )
 
@@ -221,7 +218,7 @@ def render(particle_centers, particle_widths, particle_colors):
 
 
 def loss_func(a, b, c, gt_rgb, gt_depth):
-    rgb, depth, _ = render(a, b, c)
+    _rgb, depth, _ = render(a, b, c)
     return jnp.mean(jnp.abs(depth - gt_depth))