Skip to content

lyxok1/STM-Training

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

28 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Trainig Script for Space Time Memory Network

This codebase implemented training code for Space Time Memory Network with some cyclic features.

sample results

Update

  • We have post a journal version of our paper here, the modified cycle version of AOT model in our journal paper is also available now.

Requirement

python package

  • torch
  • python-opencv
  • pillow
  • yaml
  • imgaug
  • yacs
  • progress
  • nvidia-dali (optional)

GPU support

  • GPU Memory >= 12GB
  • CUDA >= 10.0

Data

See the doc DATASET.md for more details on data organization of our prepared dataset.

Release

We provide pre-trained model with different backbone in our codebase, results are validated on DAVIS17-val with gradient correction.

model backbone data backend J F J & F link FPS
STM-Cycle Resnet18 DALI 65.3 70.8 68.1 Google Drive 14.8
STM-Cycle Resnet50 PIL 70.5 76.3 73.4 Google Drive 9.3

Runing

Appending the root folder to the search path of python interpreter

export PYTHONPATH=${PYTHONPATH}:./

To train the STM network, run following command.

python3 train.py --cfg config.yaml OPTION_KEY OPTION_VAL

To test the STM network, run following command

python3 test.py --cfg config.yaml initial ${PATH_TO_MODEL} OPTION_KEY OPTION_VAL

The test results will be saved as indexed png file at ${ROOT}/${output_dir}/${valset}.

To run a segmentation demo, run following command

python3 demo/demo.py --cfg demo/demo.yaml OPTION_KEY OPTION_VAL

The segmentation results will be saved at ${output_dir}.

Acknowledgement

This codebase borrows the code and structure from official STM repository

Reference

The codebase is built based on following works

@InProceedings{Oh_2019_ICCV,
author = {Oh, Seoung Wug and Lee, Joon-Young and Xu, Ning and Kim, Seon Joo},
title = {Video Object Segmentation Using Space-Time Memory Networks},
booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
month = {October},
year = {2019}
}

@InProceedings{Li_2020_NeurIPS,
author = {Li, Yuxi and Xu, Ning and Peng Jinlong and John See and Lin Weiyao},
title = {Delving into the Cyclic Mechanism in Semi-supervised Video Object Segmentation},
booktitle = {Neural Information Processing System (NeurIPS)},
year = {2020}
}

@article{li2022exploring,
  title={Exploring the Semi-Supervised Video Object Segmentation Problem from a Cyclic Perspective},
  author={Li, Yuxi and Xu, Ning and Yang, Wenjie and See, John and Lin, Weiyao},
  journal={International Journal of Computer Vision},
  pages={1--17},
  year={2022},
  publisher={Springer}
}

About

training script for space time memory network

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages