[pre-commit.ci] pre-commit autoupdate

.pre-commit-config.yaml

@@ -1,14 +1,14 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v6.0.0
     hooks:
       - id: check-yaml
         args: [--unsafe]
       - id: end-of-file-fixer
       - id: trailing-whitespace
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.10
+    rev: v0.16.8
     hooks:
       - id: ruff
         # types_or: [ python, pyi, jupyter ]

bayes3d/colmap/colmap_loader.py

@@ -196,9 +196,9 @@ def read_intrinsics_text(path):
                 elems = line.split()
                 camera_id = int(elems[0])
                 model = elems[1]
-                assert (
-                    model == "PINHOLE"
-                ), "While the loader support other types, the rest of the code assumes PINHOLE"
+                assert model == "PINHOLE", (
+                    "While the loader support other types, the rest of the code assumes PINHOLE"
+                )
                 width = int(elems[2])
                 height = int(elems[3])
                 params = np.array(tuple(map(float, elems[4:])))

bayes3d/colmap/dataset_loader.py

@@ -95,7 +95,9 @@ def readColmapCameras(cam_extrinsics, cam_intrinsics, images_folder):
             FovY = focal2fov(focal_length_y, height)
             FovX = focal2fov(focal_length_x, width)
         else:
-            assert False, "Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!"
+            assert False, (
+                "Colmap camera model not handled: only undistorted datasets (PINHOLE or SIMPLE_PINHOLE cameras) supported!"
+            )
 
         image_path = os.path.join(images_folder, os.path.basename(extr.name))
         image_name = os.path.basename(image_path).split(".")[0]

bayes3d/genjax/model.py

@@ -129,11 +129,11 @@ def get_far_plane(trace):
 def add_object(trace, key, obj_id, parent, face_parent, face_child):
     N = b.get_indices(trace).shape[0] + 1
     choices = trace.get_choices()
-    choices[f"parent_{N-1}"] = parent
-    choices[f"id_{N-1}"] = obj_id
-    choices[f"face_parent_{N-1}"] = face_parent
-    choices[f"face_child_{N-1}"] = face_child
-    choices[f"contact_params_{N-1}"] = jnp.zeros(3)
+    choices[f"parent_{N - 1}"] = parent
+    choices[f"id_{N - 1}"] = obj_id
+    choices[f"face_parent_{N - 1}"] = face_parent
+    choices[f"face_child_{N - 1}"] = face_child
+    choices[f"contact_params_{N - 1}"] = jnp.zeros(3)
     return model.importance(key, choices, (jnp.arange(N), *trace.get_args()[1:]))[1]
 
 

bayes3d/neural/cosypose_baseline/cosypose_utils.py

@@ -206,7 +206,7 @@ def cosypose_interface(rgb_imgs, camera_k):
     all_scores = []
     for i, rgb_img in enumerate(rgb_imgs):
         pred = COSYPOSE_MODEL.inference(rgb_img, camera_k)
-        print(f"{i+1}/{num_imgs} inference done")
+        print(f"{i + 1}/{num_imgs} inference done")
 
         pred_poses = np.asarray(pred.poses.cpu())
         pred_ids = [

bayes3d/neural/dino.py

@@ -189,9 +189,9 @@ def patch_vit_resolution(model: nn.Module, stride: int) -> nn.Module:
             return model
 
         stride = nn_utils._pair(stride)
-        assert all(
-            [(patch_size // s_) * s_ == patch_size for s_ in stride]
-        ), f"stride {stride} should divide patch_size {patch_size}"
+        assert all([(patch_size // s_) * s_ == patch_size for s_ in stride]), (
+            f"stride {stride} should divide patch_size {patch_size}"
+        )
 
         # fix the stride
         model.patch_embed.proj.stride = stride
@@ -415,7 +415,9 @@ def extract_descriptors(
         if not include_cls:
             x = x[:, :, 1:, :]  # remove cls token
         else:
-            assert not bin, "bin = True and include_cls = True are not supported together, set one of them False."
+            assert not bin, (
+                "bin = True and include_cls = True are not supported together, set one of them False."
+            )
         if not bin:
             desc = (
                 x.permute(0, 2, 3, 1).flatten(start_dim=-2, end_dim=-1).unsqueeze(dim=1)
@@ -431,9 +433,9 @@ def extract_saliency_maps(self, batch: torch.Tensor) -> torch.Tensor:
         :param batch: batch to extract saliency maps for. Has shape BxCxHxW.
         :return: a tensor of saliency maps. has shape Bxt-1
         """
-        assert (
-            self.model_type == "dino_vits8"
-        ), "saliency maps are supported only for dino_vits model_type."
+        assert self.model_type == "dino_vits8", (
+            "saliency maps are supported only for dino_vits model_type."
+        )
         self._extract_features(batch, [11], "attn")
         head_idxs = [0, 2, 4, 5]
         curr_feats = self._feats[0]  # Bxhxtxt

bayes3d/renderer.py

@@ -126,7 +126,7 @@ def add_mesh(self, mesh, mesh_name=None, scaling_factor=1.0, center_mesh=True):
         bounding_box_dims, bounding_box_pose = bayes3d.utils.aabb(mesh.vertices)
         if center_mesh:
             if not jnp.isclose(bounding_box_pose[:3, 3], 0.0).all():
-                print(f"Centering mesh with translation {bounding_box_pose[:3,3]}")
+                print(f"Centering mesh with translation {bounding_box_pose[:3, 3]}")
             mesh.vertices = mesh.vertices - bounding_box_pose[:3, 3]
 
         self.meshes.append(mesh)

bayes3d/rendering/nvdiffrast_jax/renderer_matching_pytorch.py

@@ -178,7 +178,7 @@ def xfm_points(points, matrix):
         pos_grads.min().item(),
         pos_grads.max().item(),
     )
-    print(f"JAX rasterization (eval + grad): {(end_time - start_time)*1000} ms")
+    print(f"JAX rasterization (eval + grad): {(end_time - start_time) * 1000} ms")
 
     # save viz
     b.viz.get_depth_image(rast_out[0][:, :, 2]).save("img_jax.png")
@@ -229,7 +229,7 @@ def xfm_points(points, matrix):
         pos_grads.min().item(),
         pos_grads.max().item(),
     )
-    print(f"Torch rasterization (eval + grad): {(end_time - start_time)*1000} ms")
+    print(f"Torch rasterization (eval + grad): {(end_time - start_time) * 1000} ms")
 
     # save viz
     b.viz.get_depth_image(jnp.array(rast_out[0][:, :, 2].cpu())).save("img_torch.png")
@@ -278,7 +278,7 @@ def xfm_points(points, matrix):
     print(
         f"JAX BWD (sum, min, max): g_attr={g_attr.sum().item(), g_attr.min().item(), g_attr.max().item()}\ng_rast={g_rast.sum().item(), g_rast.min().item(), g_rast.max().item()}"
     )
-    print(f"JAX interpolation: {(end_time - start_time)*1000} ms")
+    print(f"JAX interpolation: {(end_time - start_time) * 1000} ms")
 
     # save viz
     b.viz.get_depth_image(gb_pos[0][:, :, 2]).save("interpolate_jax.png")
@@ -316,7 +316,7 @@ def xfm_points(points, matrix):
     print(
         f"TORCH BWD (sum, min, max): g_attr={g_attr.sum().item(), g_attr.min().item(), g_attr.max().item()}\ng_rast={g_rast.sum().item(), g_rast.min().item(), g_rast.max().item()}"
     )
-    print(f"Torch interpolation: {(end_time - start_time)*1000} ms")
+    print(f"Torch interpolation: {(end_time - start_time) * 1000} ms")
 
     # save viz
     b.viz.get_depth_image(jnp.array(gb_pos[0][:, :, 2].cpu())).save(

scripts/_mkl/notebooks/00a - Types.ipynb

@@ -29,18 +29,18 @@
     "Int = Array\n",
     "FaceIndex = int\n",
     "FaceIndices = Array\n",
-    "ArrayN      = Array\n",
-    "Array3      = Array\n",
-    "Array2      = Array\n",
-    "ArrayNx2    = Array\n",
-    "ArrayNx3    = Array\n",
-    "Matrix      = jaxlib.xla_extension.ArrayImpl\n",
-    "PrecisionMatrix  = Matrix\n",
+    "ArrayN = Array\n",
+    "Array3 = Array\n",
+    "Array2 = Array\n",
+    "ArrayNx2 = Array\n",
+    "ArrayNx3 = Array\n",
+    "Matrix = jaxlib.xla_extension.ArrayImpl\n",
+    "PrecisionMatrix = Matrix\n",
     "CovarianceMatrix = Matrix\n",
-    "CholeskyMatrix   = Matrix\n",
-    "SquareMatrix     = Matrix\n",
-    "Vector     = Array\n",
-    "Direction  = Vector\n",
+    "CholeskyMatrix = Matrix\n",
+    "SquareMatrix = Matrix\n",
+    "Vector = Array\n",
+    "Direction = Vector\n",
     "BaseVector = Vector"
    ]
   },

scripts/_mkl/notebooks/00b - Utils.ipynb

@@ -24,7 +24,7 @@
    "source": [
     "#|export\n",
     "import matplotlib.pyplot as plt\n",
-    "from   matplotlib.collections import LineCollection\n",
+    "from matplotlib.collections import LineCollection\n",
     "import numpy as np\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
@@ -45,7 +45,7 @@
    "outputs": [],
    "source": [
     "#|export\n",
-    "key       = jax.random.PRNGKey(0)\n",
+    "key = jax.random.PRNGKey(0)\n",
     "logsumexp = jax.scipy.special.logsumexp"
    ]
   },
@@ -57,16 +57,19 @@
    "source": [
     "#|export\n",
     "def keysplit(key, *ns):\n",
-    "    if len(ns) == 0:  \n",
+    "    if len(ns) == 0:\n",
     "        return jax.random.split(key, 1)[0]\n",
     "    elif len(ns) == 1:\n",
-    "        n, = ns\n",
-    "        if n == 1: return keysplit(key)\n",
-    "        else:      return jax.random.split(key, ns[0])\n",
+    "        (n,) = ns\n",
+    "        if n == 1:\n",
+    "            return keysplit(key)\n",
+    "        else:\n",
+    "            return jax.random.split(key, ns[0])\n",
     "    else:\n",
     "        keys = []\n",
-    "        for n in ns: keys.append(keysplit(key, n))\n",
-    "        return keys\n"
+    "        for n in ns:\n",
+    "            keys.append(keysplit(key, n))\n",
+    "        return keys"
    ]
   },
   {
@@ -125,10 +128,12 @@
     "#|export\n",
     "def bounding_box(arr, pad=0):\n",
     "    \"\"\"Takes a euclidean-like arr (`arr.shape[-1] == 2`) and returns its bounding box.\"\"\"\n",
-    "    return jnp.array([\n",
-    "        [jnp.min(arr[...,0])-pad, jnp.min(arr[...,1])-pad],\n",
-    "        [jnp.max(arr[...,0])+pad, jnp.max(arr[...,1])+pad]\n",
-    "    ])"
+    "    return jnp.array(\n",
+    "        [\n",
+    "            [jnp.min(arr[..., 0]) - pad, jnp.min(arr[..., 1]) - pad],\n",
+    "            [jnp.max(arr[..., 0]) + pad, jnp.max(arr[..., 1]) + pad],\n",
+    "        ]\n",
+    "    )"
    ]
   },
   {
@@ -140,21 +145,24 @@
     "#|export\n",
     "def argmax_axes(a, axes=None):\n",
     "    \"\"\"Argmax along specified axes\"\"\"\n",
-    "    if axes is None: return jnp.argmax(a)\n",
-    "    \n",
-    "    n = len(axes)        \n",
-    "    axes_  = set(range(a.ndim))\n",
+    "    if axes is None:\n",
+    "        return jnp.argmax(a)\n",
+    "\n",
+    "    n = len(axes)\n",
+    "    axes_ = set(range(a.ndim))\n",
     "    axes_0 = axes\n",
-    "    axes_1 = sorted(axes_ - set(axes_0))    \n",
-    "    axes_  = axes_0 + axes_1\n",
+    "    axes_1 = sorted(axes_ - set(axes_0))\n",
+    "    axes_ = axes_0 + axes_1\n",
     "\n",
     "    b = jnp.transpose(a, axes=axes_)\n",
     "    c = b.reshape(np.prod(b.shape[:n]), -1)\n",
     "\n",
     "    I = jnp.argmax(c, axis=0)\n",
-    "    I = jnp.array([jnp.unravel_index(i, b.shape[:n]) for i in I]).reshape(b.shape[n:] + (n,))\n",
+    "    I = jnp.array([jnp.unravel_index(i, b.shape[:n]) for i in I]).reshape(\n",
+    "        b.shape[n:] + (n,)\n",
+    "    )\n",
     "\n",
-    "    return  I"
+    "    return I"
    ]
   },
   {
@@ -177,7 +185,7 @@
     "test_shape = (3, 99, 5, 9)\n",
     "a = jnp.arange(np.prod(test_shape)).reshape(test_shape)\n",
     "\n",
-    "I = argmax_axes(a, axes=[0,1])\n",
+    "I = argmax_axes(a, axes=[0, 1])\n",
     "I.shape"
    ]
   },
@@ -195,8 +203,12 @@
    "outputs": [],
    "source": [
     "#|export\n",
-    "def cam_to_screen(x): return jnp.array([x[0]/x[2], x[1]/x[2], jnp.linalg.norm(x)])\n",
-    "def screen_to_cam(y): return y[2]*jnp.array([y[0], y[1], 1.0])"
+    "def cam_to_screen(x):\n",
+    "    return jnp.array([x[0] / x[2], x[1] / x[2], jnp.linalg.norm(x)])\n",
+    "\n",
+    "\n",
+    "def screen_to_cam(y):\n",
+    "    return y[2] * jnp.array([y[0], y[1], 1.0])"
    ]
   },
   {
@@ -206,23 +218,25 @@
    "outputs": [],
    "source": [
     "#|export\n",
-    "def rot2d(hd): return jnp.array([\n",
-    "    [jnp.cos(hd), -jnp.sin(hd)], \n",
-    "    [jnp.sin(hd),  jnp.cos(hd)]\n",
-    "    ]);\n",
+    "def rot2d(hd):\n",
+    "    return jnp.array([[jnp.cos(hd), -jnp.sin(hd)], [jnp.sin(hd), jnp.cos(hd)]])\n",
+    "\n",
     "\n",
-    "def pack_2dpose(x,hd): \n",
-    "    return jnp.concatenate([x,jnp.array([hd])])\n",
+    "def pack_2dpose(x, hd):\n",
+    "    return jnp.concatenate([x, jnp.array([hd])])\n",
     "\n",
-    "def apply_2dpose(p, ys): \n",
-    "    return ys@rot2d(p[2] - jnp.pi/2).T + p[:2]\n",
     "\n",
-    "def unit_vec(hd): \n",
+    "def apply_2dpose(p, ys):\n",
+    "    return ys @ rot2d(p[2] - jnp.pi / 2).T + p[:2]\n",
+    "\n",
+    "\n",
+    "def unit_vec(hd):\n",
     "    return jnp.array([jnp.cos(hd), jnp.sin(hd)])\n",
     "\n",
+    "\n",
     "def adjust_angle(hd):\n",
     "    \"\"\"Adjusts angle to lie in the interval [-pi,pi).\"\"\"\n",
-    "    return (hd + jnp.pi)%(2*jnp.pi) - jnp.pi"
+    "    return (hd + jnp.pi) % (2 * jnp.pi) - jnp.pi"
    ]
   },
   {
@@ -243,7 +257,7 @@
     "\n",
     "\n",
     "def argdiffs(args, other=None):\n",
-    "    return tuple(map(lambda v: Diff(v, UnknownChange), args))\n"
+    "    return tuple(map(lambda v: Diff(v, UnknownChange), args))"
    ]
   },
   {
@@ -260,10 +274,10 @@
     "class Args(tuple):\n",
     "    def __new__(cls, *args, **kwargs):\n",
     "        return _tuple.__new__(cls, list(args) + list(kwargs.values()))\n",
-    "    \n",
+    "\n",
     "    def __init__(self, *args, **kwargs):\n",
     "        self._d = dict()\n",
-    "        for k,v in kwargs.items():\n",
+    "        for k, v in kwargs.items():\n",
     "            self._d[k] = v\n",
     "            setattr(self, k, v)\n",
     "\n",
@@ -298,29 +312,34 @@
    "outputs": [],
    "source": [
     "#|export\n",
-    "# \n",
+    "#\n",
     "# Monkey patching `sample` for `BuiltinGenerativeFunction`\n",
-    "# \n",
+    "#\n",
     "cls = genjax._src.generative_functions.static.static_gen_fn.StaticGenerativeFunction\n",
     "\n",
+    "\n",
     "def genjax_sample(self, key, *args, **kwargs):\n",
     "    tr = self.simulate(key, args)\n",
     "    return tr.get_retval()\n",
     "\n",
+    "\n",
     "setattr(cls, \"sample\", genjax_sample)\n",
     "\n",
     "\n",
-    "# \n",
+    "#\n",
     "# Monkey patching `sample` for `DeferredGenerativeFunctionCall`\n",
-    "# \n",
+    "#\n",
     "cls = genjax._src.generative_functions.supports_callees.SugaredGenerativeFunctionCall\n",
     "\n",
+    "\n",
     "def deff_gen_func_call(self, key, **kwargs):\n",
     "    return self.gen_fn.sample(key, *self.args, **kwargs)\n",
     "\n",
+    "\n",
     "def deff_gen_func_logpdf(self, x, **kwargs):\n",
     "    return self.gen_fn.logpdf(x, *self.args, **kwargs)\n",
     "\n",
+    "\n",
     "setattr(cls, \"__call__\", deff_gen_func_call)\n",
     "setattr(cls, \"sample\", deff_gen_func_call)\n",
     "setattr(cls, \"logpdf\", deff_gen_func_logpdf)"