Upload 4 files

README.md

@@ -1,3 +1,198 @@
----
-license: cc-by-2.0
----
+# E(3)-Pose
+
+In this repository, we present E(3)-Pose, the first symmetry-aware framework for 6-DoF object pose estimation from volumetric images that uses an E(3)-equivariant convolutional neural network (E(3)-CNN). Although we evaluate the utility of E(3)-Pose on fetal brain MRI, the proposed methods hold potential for broader applications.
+
+<br />
+
+![Teaser](images/teaser.png)
+
+<br />
+
+We rapidly estimate pose from volumes in a two-step process that separately estimates translation and rotation:
+
+1. **Translation Estimation:**
+    * A standard segmentation U-Net localizes the object in the volume.
+    * The center-of-mass (CoM) of the predicted mask is the estimated origin of the canonical object coordinate frame. 
+2. **Rotation Estimation:**
+    * We crop input volumes such that the predicted segmentation mask is scaled to 60% of the cropped dimensions.
+    * The E(3)-CNN takes in the cropped volume as input, and outputs an E(3)-equivariant rotation parametrization consisting of 2 vectors and 1 pseudovector.
+    * The output rotation is computed by choosing the pseudovector direction that ensures right-handedness, and orthonormalizing via support-vector decomposition (SVD).
+
+Our E(3)-CNN architecture builds on prior theoretical work on [3D steerable CNNs](https://proceedings.neurips.cc/paper_files/paper/2018/file/488e4104520c6aab692863cc1dba45af-Paper.pdf)<sup>1</sup> and uses code borrowed from [e3nn-UNet](https://github.com/SCAN-NRAD/e3nn_Unet)<sup>2</sup>, which implements 3D convolutions with the [e3nn](https://e3nn.org/)<sup>3</sup> Python library for building E(3)-equivariant networks.
+
+<br />
+
+![Method overview](images/method_overview.png)
+
+<br />
+
+Overall, E(3)-Pose outperforms state-of-the-art methods for pose estimation in fetal brain MRI volumes representative of clinical applications, 
+including strategies that rely on anatomical landmark detection ([Fetal-Align](https://github.com/mu40/fetal-align)<sup>4</sup>), 
+template registration ([FireANTs](https://github.com/rohitrango/FireANTs)<sup>5</sup> and [EquiTrack](https://github.com/BBillot/EquiTrack)<sup>6</sup>), and direct pose regression with standard CNNs ([3DPose-Net](https://github.com/SadeghMSalehi/DeepRegistration)<sup>7</sup>, [6DRep](https://www.spiedigitallibrary.org/conference-proceedings-of-spie/12464/124640T/Automatic-brain-pose-estimation-in-fetal-MRI/10.1117/12.2647613.full)<sup>8</sup>, [RbR](https://github.com/HuXiaoling/Regre4Regis)<sup>9</sup>). 
+See figure below for example results. Particularly, we show in our paper that regularizing network parameters to conform with physical symmetries mitigates overfitting to research-quality training datasets, and permits better generalization to out-of-distribution clinical data with pose ambiguities.
+
+<br />
+
+![examples](images/examples.png)
+
+<br />
+
+The full article describing this method is available at:
+
+**Equivariant Symmetry-Aware Head Pose Estimation for Fetal MRI** \
+Muthukrishnan, Gagoski, Lee, Grant, Adalsteinsson, Golland, Billot \
+arXiV (2025) \
+[ [arxiv](https://arxiv.org/abs/2512.04890) | [bibtex](bibtex.bib) ]
+
+---
+### Installation
+
+1. Clone this repository.
+2. Edit the environment prefix in `environment.yml` and then install all dependencies:
+```
+    cd E3-Pose/
+    conda env create -f environment.yml
+    conda activate e3pose
+    pip install -r requirements.txt
+```
+3. Install [pytorch3d](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md)
+4. If you want to use our trained model weights for fetal brain MRI, download the model weights [here](https://drive.google.com/drive/folders/1r6FVzXG9VLH-0MtMnD2hwjzdDqss1DSE?usp=sharing).
+5. If you want to train your own network on a [publicly available fetal MRI dataset](https://pubmed.ncbi.nlm.nih.gov/40800813/)<sup>10</sup>, download our manually annotated segmentations and poses [here](https://drive.google.com/file/d/1yO2o2sNNNEfcB_-ZDcVvHyCGxqk6SYyE/view?usp=sharing).
+
+You're now ready to use E(3)-Pose!
+
+<br />
+
+---
+### Usage
+
+This repository contains all the code necessary to train and test your own networks. We provide separate scripts for training the segmentation U-Net and E(3)-CNNs, and a single script to deploy both for full rigid pose estimation.
+
+#### Training a Segmentation U-Net for Translation Estimation
+
+1. Set up separate training/validation dataset directories for images and ground-truth segmentation labels, where file names between image and label directories are the same. Ensure that all image file extensions are .nii or .nii.gz.
+
+2. If you are training a multi-class segmentation network, ensure that the object for which you want to estimate pose has category label 1 in the ground-truth labels.
+
+3. Name the output directory to save all model weights and metrics during network training.
+
+4. To train the segmentation U-Net, run:
+
+    ```
+    python scripts/train_unet.py train_image_dir/ train_label_dir/ val_image_dir/ val_label_dir/ output_dir/
+    ```
+
+    For detailed descriptions of other arguments, run:
+    
+    ```
+    python scripts/train_unet.py -h
+    ```
+
+#### Training an E(3)-CNN for Rotation Estimation
+
+1. Set up separate training/validation dataset directories for images and ground-truth segmentation labels, where file names between image and label directories are the same. Ensure that all image file extensions are .nii or .nii.gz. If your segmentation labels have multiple classes, ensure that the object for which you want to estimate pose has category label 1.
+
+2. Set up separate CSV files for rotation annotations in training and validation datasets, in the following format:
+
+    | frame_id | rot_x | rot_y | rot_z |
+    |----------|-------|-------|-------|
+    | ...      | ...   | ...   | ...   |
+
+    where **frame_id** is the file name of the volume without the file extension, and **rot_x**, **rot_y**, **rot_z** are the Euler angles in degrees of the rotation from the volume to the canonical coordinate frame. The Euler angle rotation assumes the "xyz" ordering convention.
+
+3. Name the output directory to save all model weights and metrics during network training.
+
+4. To train the E(3)-CNN, run:
+
+    ```
+    python scripts/train_e3cnn.py train_image_dir/ train_label_dir/ path_to_train_annotations.csv \
+        val_image_dir/ val_label_dir/ path_to_val_annotations.csv \
+        output_dir/
+    ```
+
+    For detailed descriptions of other arguments, run:
+    
+    ```
+    python scripts/train_e3cnn.py -h
+    ```
+
+#### Running Rigid Pose Estimation with Trained Model Weights
+
+1. Set up an input directory of images (all file extensions must be .nii or .nii.gz) on which to run rigid pose estimation.
+
+2. Name the output directory to save all estimated poses.
+
+3. To estimate pose on all inputs, run:
+
+    ```
+    python scripts.inference.py input_image_dir/ output_dir/ path_to_segmentation_unet.ckpt path_to_e3cnn.pth
+    ```
+
+    For detailed descriptions of other arguments, run:
+    
+    ```
+    python scripts/inference.py -h
+    ```
+
+4. Output poses are saved as 4x4 transform matrices in .npy format in the output directory, where file names are the same as the inputs.
+
+<br />
+
+---
+### Citation/Contact
+
+If you find this work useful for your research, please cite:
+
+**Equivariant Symmetry-Aware Head Pose Estimation for Fetal MRI** \
+Muthukrishnan, Gagoski, Lee, Grant, Adalsteinsson, Golland, Billot \
+arXiV (2025) \
+[ [arxiv](https://arxiv.org/abs/2512.04890) | [bibtex](bibtex.bib) ]
+
+If you have any question regarding the usage of this code, or any suggestions to improve it, please raise an issue
+(preferred) or contact us at:\
+**[email protected]**
+
+
+<br />
+
+---
+### References
+<sup>1</sup> *3D steerable CNNs: Learning rotationally equivariant features in volumetric data* \
+Weiler, Geiger, Welling, Boomsma, Cohen \
+Advances in Neural Information Processing Systems, 2018
+
+<sup>2</sup> *Leveraging SO(3)-steerable convolutions for pose-robust semantic segmentation in 3D medical data* \
+Diaz, Geiger, McKinley \
+Journal of Machine Learning in Biomedical Imaging, 2024
+
+<sup>3</sup> *e3nn: Euclidean neural networks* \
+Geiger and Smidt \
+arXiV, 2022
+
+<sup>4</sup> *Rapid head-pose detection for automated slice prescription of fetal-brain MRI* \
+Hoffmann, Abaci Turk, Gagoski, Morgan, Wighton, Tisdall, Reuter, Adalsteinsson, Grant, Wald, van der Kouwe \
+International Journal of Imaging Systems and Technology, 2021
+
+<sup>5</sup> *FireANTs: Adaptive Riemannian optimization for multi-scale diffeomorphic registration* \
+Jena, Chaudhari, Gee \
+arXiV, 2024
+
+<sup>6</sup> *SE(3)-equivariant and noise-invariant 3D rigid motion tracking in brain MRI* \
+Billot, Dey, Moyer, Hoffmann, Abaci Turk, Gagoski \
+IEEE Transactions on Medical Imaging, 2024
+
+<sup>7</sup> *Real-time deep pose estimation with geodesic loss for image-to-template rigid registration* \
+Salehi, Khan, Erdogmus, Gholipour \
+IEEE Transactions on Medical Imaging, 2019
+
+<sup>8</sup> *Automatic brain pose estimation in fetal MRI* \
+Faghihpirayesh, Karimi, Erdogmus, Gholipour \
+Proceedings of SPIE: Medical Imaging: Image Processing, 2023
+
+<sup>9</sup> *Registration by Regression (RbR): A framework for interpretable and flexible atlas registration* \
+Gopinath, Hu, Hoffmann, Puonti, Iglesias \
+International Workshop on Biomedical Image Registration, 2024
+
+<sup>10</sup> *The developing human connectome project fetal functional MRI release: Methods and data structures* \
+Karolis et al \
+Imaging Neuroscience, 2025

bibtex.bib

@@ -0,0 +1,9 @@
+@misc{muthukrishnan2025equivariantsymmetryawareheadpose,
+      title={Equivariant Symmetry-Aware Head Pose Estimation for Fetal MRI}, 
+      author={Ramya Muthukrishnan and Borjan Gagoski and Aryn Lee and P. Ellen Grant and Elfar Adalsteinsson and Polina Golland and Benjamin Billot},
+      year={2025},
+      eprint={2512.04890},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2512.04890}, 
+}

environment.yaml

@@ -0,0 +1,129 @@
+name: e3pose
+channels:
+  - defaults
+dependencies:
+  - python=3.10
+  - pandas
+  - pip:
+    - fvcore
+    - iopath
+    - aiohttp==3.9.1
+    - aiosignal==1.3.1
+    - alabaster==0.7.16
+    - async-timeout==4.0.3
+    - attrs==23.2.0
+    - babel==2.16.0
+    - batchgenerators==0.25
+    - click==8.1.7
+    - contourpy==1.1.1
+    - cycler==0.12.1
+    - Deprecated==1.2.14
+    - dicom2nifti==2.4.9
+    - docopt==0.6.2
+    - docutils==0.19
+    - e3nn==0.5.5
+    - et_xmlfile==2.0.0
+    - exceptiongroup==1.2.2
+    - filelock==3.13.1
+    - fonttools==4.43.1
+    - frozenlist==1.4.1
+    - fsspec==2024.2.0
+    - humanize==4.9.0
+    - imageio==2.32.0
+    - imagesize==1.4.1
+    - iniconfig==2.1.0
+    - Jinja2==3.1.2
+    - joblib==1.3.2
+    - kiwisolver==1.4.5
+    - lazy_loader==0.3
+    - lightning==2.5.0.post0
+    - lightning-utilities==0.10.1
+    - linecache2==1.0.0
+    - markdown-it-py==3.0.0
+    - MarkupSafe==2.1.3
+    - matplotlib==3.10.0
+    - mdurl==0.1.2
+    - MedPy==0.4.0
+    - monai==1.4.0
+    - mpmath==1.3.0
+    - multidict==6.0.4
+    - networkx==3.2.1
+    - nibabel==5.1.0
+    - nnunet==1.7.1
+    - nose==1.3.7
+    - numpy==1.26.0
+    - nvidia-cublas-cu11==11.10.3.66
+    - nvidia-cublas-cu12==12.1.3.1
+    - nvidia-cuda-cupti-cu12==12.1.105
+    - nvidia-cuda-nvrtc-cu11==11.7.99
+    - nvidia-cuda-nvrtc-cu12==12.1.105
+    - nvidia-cuda-runtime-cu11==11.7.99
+    - nvidia-cuda-runtime-cu12==12.1.105
+    - nvidia-cudnn-cu11==8.5.0.96
+    - nvidia-cudnn-cu12==8.9.2.26
+    - nvidia-cufft-cu12==11.0.2.54
+    - nvidia-curand-cu12==10.3.2.106
+    - nvidia-cusolver-cu12==11.4.5.107
+    - nvidia-cusparse-cu12==12.1.0.106
+    - nvidia-nccl-cu12==2.19.3
+    - nvidia-nvjitlink-cu12==12.4.127
+    - nvidia-nvtx-cu12==12.1.105
+    - opencv-python==4.10.0.84
+    - openpyxl==3.1.5
+    - opt-einsum==3.3.0
+    - opt-einsum-fx==0.1.4
+    - packaging==23.2
+    - pandas==2.1.4
+    - patsy==1.0.1
+    - Pillow==10.0.1
+    - pluggy==1.5.0
+    - pydicom==2.4.3
+    - Pygments==2.17.2
+    - pykwalify==1.8.0
+    - pyparsing==3.1.1
+    - pyradiomics==3.0.1
+    - pytest==8.3.5
+    - python-dateutil==2.8.2
+    - python-gdcm==3.0.22
+    - pytorch-lightning==2.1.3
+    - pytz==2023.3.post1
+    - PyWavelets==1.6.0
+    - PyYAML==6.0.1
+    - rich==13.7.0
+    - ruamel.yaml==0.18.6
+    - ruamel.yaml.clib==0.2.8
+    - scikit-image==0.22.0
+    - scikit-learn==1.6.0
+    - scipy==1.11.2
+    - shellingham==1.5.4
+    - SimpleITK==2.3.1
+    - six==1.16.0
+    - snowballstemmer==2.2.0
+    - Sphinx==6.2.1
+    - sphinxcontrib-applehelp==2.0.0
+    - sphinxcontrib-devhelp==2.0.0
+    - sphinxcontrib-htmlhelp==2.1.0
+    - sphinxcontrib-jsmath==1.0.1
+    - sphinxcontrib-qthelp==2.0.0
+    - sphinxcontrib-serializinghtml==2.0.0
+    - sphstat==1.0.6
+    - statsmodels==0.14.4
+    - sympy==1.13.1
+    - threadpoolctl==3.2.0
+    - tifffile==2023.9.26
+    - tomli==2.2.1
+    - torch==2.2.1
+    - torchaudio==2.2.1
+    - torchmetrics==1.3.0.post0
+    - torchvision==0.17.1
+    - torchio==0.20.6
+    - traceback2==1.4.0
+    - triton==2.2.0
+    - typer==0.9.0
+    - typing==3.7.4.3
+    - typing_extensions==4.12.2
+    - unittest2==1.1.0
+    - webcolors==24.11.1
+    - wrapt==1.16.0
+    - yarl==1.9.4
+prefix: /path/to/envs/e3pose

requirements.txt

@@ -0,0 +1,120 @@
+aiohttp==3.9.1
+aiosignal==1.3.1
+alabaster==0.7.16
+async-timeout==4.0.3
+attrs==23.2.0
+babel==2.16.0
+batchgenerators==0.25
+click==8.1.7
+contourpy==1.1.1
+cycler==0.12.1
+Deprecated==1.2.14
+dicom2nifti==2.4.9
+docopt==0.6.2
+docutils==0.19
+e3nn==0.5.5
+et_xmlfile==2.0.0
+exceptiongroup==1.2.2
+filelock==3.13.1
+fonttools==4.43.1
+frozenlist==1.4.1
+fsspec==2024.2.0
+humanize==4.9.0
+imageio==2.32.0
+imagesize==1.4.1
+iniconfig==2.1.0
+Jinja2==3.1.2
+joblib==1.3.2
+kiwisolver==1.4.5
+lazy_loader==0.3
+lightning==2.5.0.post0
+lightning-utilities==0.10.1
+linecache2==1.0.0
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.10.0
+mdurl==0.1.2
+MedPy==0.4.0
+monai==1.4.0
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.2.1
+nibabel==5.1.0
+nnunet==1.7.1
+nose==1.3.7
+numpy==1.26.0
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.19.3
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.1.105
+opencv-python==4.10.0.84
+openpyxl==3.1.5
+opt-einsum==3.3.0
+opt-einsum-fx==0.1.4
+packaging==23.2
+pandas==2.1.4
+patsy==1.0.1
+Pillow==10.0.1
+pluggy==1.5.0
+pydicom==2.4.3
+Pygments==2.17.2
+pykwalify==1.8.0
+pyparsing==3.1.1
+pyradiomics==3.0.1
+pytest==8.3.5
+python-dateutil==2.8.2
+python-gdcm==3.0.22
+pytorch-lightning==2.1.3
+pytz==2023.3.post1
+PyWavelets==1.6.0
+PyYAML==6.0.1
+rich==13.7.0
+ruamel.yaml==0.18.6
+ruamel.yaml.clib==0.2.8
+scikit-image==0.22.0
+scikit-learn==1.6.0
+scipy==1.11.2
+shellingham==1.5.4
+SimpleITK==2.3.1
+six==1.16.0
+snowballstemmer==2.2.0
+Sphinx==6.2.1
+sphinxcontrib-applehelp==2.0.0
+sphinxcontrib-devhelp==2.0.0
+sphinxcontrib-htmlhelp==2.1.0
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==2.0.0
+sphinxcontrib-serializinghtml==2.0.0
+sphstat==1.0.6
+statsmodels==0.14.4
+sympy==1.13.1
+threadpoolctl==3.2.0
+tifffile==2023.9.26
+tomli==2.2.1
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.2.1
+torchaudio==2.2.1
+torchmetrics==1.3.0.post0
+torchvision==0.17.1
+torchio==0.20.6
+traceback2==1.4.0
+triton==2.2.0
+typer==0.9.0
+typing==3.7.4.3
+typing_extensions==4.12.2
+unittest2==1.1.0
+webcolors==24.11.1
+wrapt==1.16.0
+yarl==1.9.4