SMITE: Segment Me In TimE

Amir Alimohammadi¹  Sauradip Nag¹  Saeid Asgari Taghanaki^{1, 2}  Andrea Tagliasacchi^{1, 3, 4} 
Ghassan Hamarneh¹  Ali Mahdavi Amiri¹ 

^{1Simon Fraser University 2Autodesk Research 3University of Toronto 4Google DeepMind}

Paper | Project Page

SMITE is an advanced open-source framework for temporally consistent video segmentation, designed to predict and segment objects across video frames using one or few reference images. With its ability to track and generalize unseen video sequences based on flexible granularity, it ensures precision and efficiency in segmentation, even when faced with occlusions, varying poses, and lighting conditions.

Note: We will release the full code soon. Stay tuned for updates!

Why SMITE?

• SMITE minimizes dataset needs by leveraging pre-trained diffusion models and a few reference images for segmentation.
• It ensures consistent segmentation across video frames with its advanced tracking and temporal voting mechanism.
• It offers flexible segmentation at various granularities, making it ideal for tasks requiring different levels of detail.

We invite you to explore the full potential of SMITE and look forward to your feedback and contributions. If you find SMITE valuable, kindly star the repository on GitHub!

Overview

The main parts of the framework are as follows:

SMITE
├── run.py                            -- script to train the models, runs factory.trainer.py                      
├── models                    
│   ├── unet.py                       -- inflated unet definition
|   ├── attention.py                  -- FullyFrameAttention to attend to all frames
│   ├── controlnet3d.py               -- ControlNet3D model definition (it is not part of the main model but we support it)
|   ├── ...
├── src
|   ├── pipeline_smite.py             -- main pipeline containing all the important functions
|   ├── train.py                     
|   ├── inference.py                     
|   ├── slicing.py                    -- slices frames and latents for efficient attention processing across video sequences
|   ├── tracking.py                   -- tracker initialization, applies tracking to each frame, and uses feature voting
|   ├── frequency_filter.py           -- DCT filter for low-pass regularization
|   ├── metric.py                     
|   ├── latent_optimization.py        -- spatio-temporal guidance          
|   ├── ...   
├── scripts
|   ├── train.sh                      --script for model training
|   ├── inference.sh                  --script for model inference on videos
|   ├── test_on_images.sh             --script for testing the model on image datasets
|   ├── ...   
├── utils
|   ├── setup.py                     
|   ├── args.py                       -- define, parse, and update command-line arguments
|   ├── transfer_weights.py           -- transfer the 2D Unet weights to inflated Unet
|   ├── ...   

TO-DO Checklist

• Add dataset
• Add training script
• Add inference script
• Support for XMEM++ dataset
• Enable multi-gpu training

Citation

If you find this project useful for your research, please use the following BibTeX entry.

@misc{alimohammadi2024smitesegmenttime,
      title={SMITE: Segment Me In TimE}, 
      author={Amirhossein Alimohammadi and Sauradip Nag and Saeid Asgari Taghanaki and Andrea Tagliasacchi and Ghassan Hamarneh and Ali Mahdavi Amiri},
      year={2024},
      eprint={2410.18538},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2410.18538}, 
}