Switch to diffdrr's wrapper of pytorch3d

diffpose/calibration.py

@@ -10,9 +10,10 @@
 from typing import Optional
 
 from beartype import beartype
+from diffdrr.utils import Transform3d
 from diffdrr.utils import convert as convert_so3
+from diffdrr.utils import se3_exp_map, se3_log_map
 from jaxtyping import Float, jaxtyped
-from pytorch3d.transforms import Transform3d, se3_exp_map, se3_log_map
 
 # %% ../notebooks/api/02_calibration.ipynb 7
 @beartype

diffpose/metrics.py

@@ -63,14 +63,14 @@ def __init__(self, patch_size=None):
 # %% ../notebooks/api/04_metrics.ipynb 9
 import torch
 from beartype import beartype
-from diffdrr.utils import convert
-from jaxtyping import Float, jaxtyped
-from pytorch3d.transforms import (
+from diffdrr.utils import (
+    convert,
     so3_log_map,
     so3_relative_angle,
     so3_rotation_angle,
     standardize_quaternion,
 )
+from jaxtyping import Float, jaxtyped
 
 from .calibration import RigidTransform
 

environment.yml

@@ -2,13 +2,11 @@ name: diffpose
 channels:
   - conda-forge
   - pytorch
-  - pytorch3d
   - nvidia
 dependencies:
   - pip
   - pytorch
   - torchvision
-  - pytorch3d
   - pip:
     - diffdrr>=0.3.8
     - h5py

notebooks/api/02_calibration.ipynb

@@ -97,9 +97,10 @@
     "from typing import Optional\n",
     "\n",
     "from beartype import beartype\n",
+    "from diffdrr.utils import Transform3d\n",
     "from diffdrr.utils import convert as convert_so3\n",
-    "from jaxtyping import Float, jaxtyped\n",
-    "from pytorch3d.transforms import Transform3d, se3_exp_map, se3_log_map"
+    "from diffdrr.utils import se3_exp_map, se3_log_map\n",
+    "from jaxtyping import Float, jaxtyped"
    ]
   },
   {

notebooks/api/04_metrics.ipynb

@@ -169,14 +169,14 @@
     "#| export\n",
     "import torch\n",
     "from beartype import beartype\n",
-    "from diffdrr.utils import convert\n",
-    "from jaxtyping import Float, jaxtyped\n",
-    "from pytorch3d.transforms import (\n",
+    "from diffdrr.utils import (\n",
+    "    convert,\n",
     "    so3_log_map,\n",
     "    so3_relative_angle,\n",
     "    so3_rotation_angle,\n",
     "    standardize_quaternion,\n",
     ")\n",
+    "from jaxtyping import Float, jaxtyped\n",
     "\n",
     "from diffpose.calibration import RigidTransform"
    ]

settings.ini

@@ -39,6 +39,6 @@ status = 3
 user = eigenvivek
 
 ### Optional ###
-requirements = diffdrr h5py scikit-image seaborn torch torchvision pytorch3d timm pytorch-transformers torchmetrics tqdm beartype jaxtyping
+requirements = diffdrr h5py scikit-image seaborn torch torchvision timm pytorch-transformers torchmetrics tqdm beartype jaxtyping
 dev_requirements = jupyterlab_code_formatter black flake8 isort nbdev ipykernel jupyter-server-proxy 
 optional_requirements = submitit