3D Mitochondria Instance Segmentation with Spatio-Temporal Transformers

Omkar Thawakar,Rao Muhammad Anwer, Jorma Laaksonen, Orly Reiner, Mubarak Shah, Fahad Shahbaz Khan

• (Feb 23, 2023): Code uploaded V1.0. (in progress)
• (Feb 21, 2023): Arxiv Release.

Accurate 3D mitochondria instance segmentation in electron microscopy (EM) is a challenging problem and serves as a prerequisite to empirically analyze their distributions and morphology. Most existing approaches employ 3D convolutions to obtain representative features. However, these convolution-based approaches struggle to effectively capture long-range dependencies in the volume mitochondria data, due to their limited local receptive field. To address this, we propose a hybrid encoder-decoder framework based on a split spatio-temporal attention module that efficiently computes spatial and temporal selfattentions in parallel, which are later fused through a deformable convolution. Further, we introduce a semantic foreground-background adversarial loss during training that aids in delineating the region of mitochondria instances from the background clutter. Our extensive experiments on three benchmarks, Lucchi, MitoEM-R and MitoEM-H, reveal the benefits of the proposed contributions achieving state-of-theart results on all three datasets.

First, clone the repository locally:

git clone https://github.com/OmkarThawakar/STT-UNET.git

Then, create environment from yml or create conda env and install dependencies:

conda env create -f mito.yml

OR

conda create -n mito python=3.8
conda activate mito
pip install -r requirements.txt

Install dependencies and pycocotools for MitoEM:

pip install torchsummary waterz malis

If you find our work useful, please consider citing:

    @article{thawakar2023sttunet,
      title={3D Mitochondria Instance Segmentation with Spatio-Temporal Transformers},
      author={Omkar Thawakar and Rao Muhammad Anwer and Jorma Laaksonen and Orly Reiner and Mubarak Shah and Fahad Shahbaz Khan},
      journal={arXiv:2303.12073},
      year={2023},
}

Should you have any question, please contact at [email protected].