E(3)-Pose

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).

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

1. Install pytorch3d
2. If you want to use our trained model weights for fetal brain MRI, download the model weights here.
3. If you want to train your own network on a publicly available fetal MRI dataset¹⁰, download our manually annotated segmentations and poses here.

Usage

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.

Citation/Contact

References