Home

Awesome

SAM-RBox

This is an implementation of SAM (Segment Anything Model) for generating rotated bounding boxes with MMRotate, which is a comparison method of H2RBox-v2: Boosting HBox-supervised Oriented Object Detection via Symmetric Learning.

NOTE: This project has been involved into OpenMMLab's new repo PlayGround. For more details, please refer to this.

<div align=center> <img src="https://user-images.githubusercontent.com/79644233/231636420-8b7f81f3-51d2-439c-87cc-6f7eebd32193.png"/> </div>

Recently, SAM has demonstrated strong zero-shot capabilities by training on the largest segmentation dataset to date. Thus, we use a trained horizontal FCOS detector to provide HBoxes into SAM as prompts, so that corresponding Masks can be generated by zero-shot, and finally the rotated RBoxes are obtained by performing the minimum circumscribed rectangle operation on the predicted Masks. Thanks to the powerful zero-shot capability, SAM-RBox based on ViT-B has achieved 63.94%. However, it is also limited to the time-consuming post-processing, only 1.7 FPS during inference.

image image

Prepare Env

The code is based on MMRotate 1.x and official API of SAM.

Here is the installation commands of recommended environment.

pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html

pip install openmim
mim install mmengine 'mmcv>=2.0.0rc0' 'mmrotate>=1.0.0rc0'

pip install git+https://github.com/facebookresearch/segment-anything.git
pip install opencv-python pycocotools matplotlib onnxruntime onnx

Note

  1. Prepare DOTA data set according to MMRotate doc.
  2. Download the detector weight from MMRotate model zoo.
  3. python main_sam_dota.py prompts SAM with HBox obtained from annotation file (such as DOTA trainval).
  4. python main_rdet-sam_dota.py prompts SAM with HBox predicted by a well-trained detector for non-annotated data (such as DOTA test).
  5. Many configs, including pipeline (i.e. transforms), dataset, dataloader, evaluator, visualizer, are set in data.py.
  6. You can change the detector config and the corresponding weight path in main_rdet-sam_dota.py to any detector that can be built with MMRotate.

Citation

@article{yu2023h2rboxv2,
  title={H2RBox-v2: Boosting HBox-supervised Oriented Object Detection via Symmetric Learning},
  author={Yu, Yi and Yang, Xue and Li, Qingyun and Zhou, Yue and Zhang, Gefan and Yan, Junchi and Da, Feipeng},
  journal={arXiv preprint arXiv:2304.04403},
  year={2023}
}

@inproceedings{yang2023h2rbox,
  title={H2RBox: Horizontal Box Annotation is All You Need for Oriented Object Detection},
  author={Yang, Xue and Zhang, Gefan and Li, Wentong and Wang, Xuehui and Zhou, Yue and Yan, Junchi},
  booktitle={International Conference on Learning Representations},
  year={2023}
}

@article{kirillov2023segany,
  title={Segment Anything}, 
  author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
  journal={arXiv:2304.02643},
  year={2023}
}

Other awesome SAM projects: