Official PyTorch Implementation of:
Link to published paper here
Micro-ultrasound is a new modality for accurate, low-cost prostate cancer imaging, but its acquisition produces oblique slices that do not align with axial MRI or histopathology. This geometric mismatch complicates interpretation and prevents direct registration to histopathology, which is necessary to map ground-truth cancer outlines onto micro-ultrasound for training machine learning models for automated cancer detection. We address this challenge with a geometry-aware reconstruction framework that converts oblique micro-ultrasound slices into axial 3D volumes. Our method includes: (i) a coordinate-based sampling scheme that uses cylindrical geometry to accurately map each voxel into Cartesian space, and (ii) a generalized implicit neural representation that models the continuous intensity field between slices, preserving high-frequency speckle texture that traditional interpolation blurs. The reconstructed volumes achieve a 9% relative SSIM improvement over a coordinate-matched trilinear baseline while maintaining ultrasound-specific texture and boundary detail. This framework produces high-quality axial micro-ultrasound volumes suitable for reliable histopathology registration and for creating pathology-informed datasets to train cancer detection models.
Our model utilizes a Dual-Path Hybrid Attention Transformer (HAT) augmented with Arc Length Embeddings to encode physical acquisition geometry.
Our dataset is not publicaly available.
We include data/data_build.py so one can easily create their own dataset given micro ultra-sound scans.
All configurations for training, evalatuion, and inference are managed through the hydra configs.
Training supports Torch Distributed Data Parallel.
# Train on 2 GPUs
torchrun --standalone --nproc_per_node=2 train.pyReconstruct a full 3D NIfTI volume from raw DICOM data.
python inference.pyCode adapted from HAT: Hybrid Attention Transformer.
To do when created