lecongthuong/u2net

U^2-Net

The code for our newly accepted paper in Pattern Recognition 2020:

U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection, Xuebin Qin, Zichen Zhang, Chenyang Huang, Masood Dehghan, Osmar R. Zaiane and Martin Jagersand.

Contact: xuebin[at]ualberta[dot]ca

U^2-Net Results (173.6 MB)

Our previous work: BASNet (CVPR 2019)

Required libraries

Python 3.6
numpy 1.15.2
scikit-image 0.14.0
PIL 5.2.0
PyTorch 0.4.0
torchvision 0.2.1
glob

Usage

1. Clone this repo

git clone https://github.com/NathanUA/U-2-Net.git

2. Download the pre-trained model u2net.pth (173.6 MB) or u2netp.pth (4.7 MB) and put it into the dirctory './saved_models/u2net/' and './saved_models/u2netp/'

3. Cd to the directory 'U-2-Net', run the train or inference process by command: python u2net_train.py or python u2net_test.py respectively. The 'model_name' in both files can be changed to 'u2net' or 'u2netp' for using different models.

We also provide the predicted saliency maps (u2net results,u2netp results) for datasets SOD, ECSSD, DUT-OMRON, PASCAL-S, HKU-IS and DUTS-TE.

U^2-Net Architecture

Quantitative Comparison

Qualitative Comparison

Citation

@InProceedings{Qin_2020_PR,
author = {Qin, Xuebin and Zhang, Zichen and Huang, Chenyang and Dehghan, Masood and Zaiane, Osmar and Jagersand, Martin},
title = {U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection},
booktitle = {Pattern Recognition},
year = {2020}
}

u2net's Issues

train my own datasets

when I train my own dataset, the results is bad，Can you give me some suggestions?