WaterMask

---

:speech_balloon: Updates

[2026-03] Upgrade to MMDetection 3.0+: We have upgraded the code from MMDetection 2.0 to MMDetection 3.0+! The main reasons for this upgrade are:

• MMDet 2.0 is outdated and no longer actively maintained
• MMDet 3.0 provides better performance, cleaner APIs, and improved modularity
• Better compatibility with latest PyTorch versions (2.0+) and CUDA
• Access to newer model architectures and training techniques

The original MMDet 2.0 version is still available on the mmdet2 branch.

[2025-07] UIIS10K:

UIIS10K upgrades the original UIIS dataset with 10 048 underwater images annotated at pixel level for 10 object classes (fish · reptiles · artiodactyla · mollusks · corals · plants · garbage · ruins · divers · robots).

• Largest underwater benchmark – more images and masks than any existing general underwater instance segmentation dataset.
• Scene diversity – shallow & deep water, clear & turbid conditions, multiple resolutions, complex backgrounds, heavy occlusions.
• Crowded scenes – 23 % of images contain ≥ 5 instances. (up to 80+ per image).
• Multi-task ready – pixel masks + bounding boxes support instance segmentation, object detection, and semantic segmentation.

---

This repo is the official implementation of "WaterMask: Instance Segmentation for Underwater Imagery". By Shijie Lian, Hua Li, Runmin Cong, Suqi Li, Wei Zhang, Sam Kwong, and has been accepted by ICCV2023! 🎉🎉🎉

:rocket: Highlights:

• UIIS dataset: We construct the first general Underwater Image Instance Segmentation (UIIS) dataset containing 4,628 images for 7 categories with pixel-level annotations for underwater instance segmentation task.

  

• SOTA performance: We propose the first underwater instance segmentation model, WaterMask, as far as we know, which achieves good performance in qualitative and quantitative comparisons with natural image instance segmentation SOTA methods.

  

Installation

Requirements

• Python 3.7+
• PyTorch 2.0+ (we use PyTorch 2.4.1)
• CUDA 12.1 or other version
• mmengine
• mmpretrain
• mmcv>=2.0.0
• MMDetection 3.0+

Environment Installation

Step 0: Download and install Miniconda from the official website.

Step 1: Create a conda environment and activate it.

conda create -n watermask python=3.9 -y
conda activate watermask

Step 2: Install PyTorch. If you have experience with PyTorch and have already installed it, you can skip to the next section.

pip install torch==2.4.1 torchvision==0.19.1 --index-url https://download.pytorch.org/whl/cu121

Step 3: Install MMEngine, MMCV, and MMDetection using MIM.

pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0"
mim install mmpretrain
mim install mmdet

Datasets

Please create a data folder in your working directory and put the UIIS dataset in it for training or testing. UIIS is divided into two parts, with 3937 images for training and 691 images for testing.

data
  ├── UDW
  |   ├── annotations
  │   │   │   ├── train.json
  │   │   │   ├── val.json
  │   ├── train
  │   │   ├── L_1.jpg
  │   │   ├── ......
  │   ├── ......

you can get our UIIS dataset in Hugging Face、Baidu Disk (pwd:fiuk) or Google Drive

Main Results

We provide some results on UIIS dataset with pretrained models. These model are trained on an NVIDIA A5000 GPU. Note that all models and logs are available at Baidu Netdisk and Google Drive link.

model	Schedule	mAP	AP50	AP75	config	download
WaterMask R101-FPN	1x	25.6	41.7	27.9	config	log / pth
WaterMask R50-FPN	3x	26.4	43.6	28.8	config	log / pth
WaterMask R101-FPN	3x	27.2	43.7	29.3	config	log / pth
Cascade WaterMask R101-FPN	3x	27.1	42.9	30.4	config	log / pth

Training

python tools/train.py configs/_our_/water_r50_fpn_1x.py --work-dir you_dir_to_save_logs_and_models

or

bash tools/dist_train.sh configs/_our_/water_r50_fpn_1x.py NUM_GPUS --work-dir you_dir_to_save_logs_and_models

Test

python tools/test.py configs/_our_/water_r50_fpn_1x.py model_checkpoint_path --eval segm

or

bash tools/dist_test.sh configs/_our_/water_r50_fpn_1x.py model_checkpoint_path NUM_GPUS --eval segm

Citation

If you find our repo useful for your research, please cite us:

@InProceedings{UIIS_Dataset_2023,
    author    = {Shijie Lian, Hua Li, Runmin Cong, Suqi Li, Wei Zhang, Sam Kwong},
    title     = {WaterMask: Instance Segmentation for Underwater Imagery},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2023},
    pages     = {1305-1315}
}

@article{UIIS10K_Dataset_2025,
    author    = {Hua Li, Shijie Lian, Zhiyuan Li, Runmin Cong, Chongyi Li},
    title     = {Taming SAM for Underwater Instance Segmentation and Beyond},
    year      = {2025},
    journal   = {arXiv preprint arXiv:2505.15581},
}

License

This project is released under the Apache 2.0 license.

Acknowledgement

This software repository is implemented based on the MMDetection framework. Thanks to them for their excellent work.