/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."

Primary LanguagePythonApache License 2.0Apache-2.0

U^2-Net (U square net)

The code for our newly accepted paper U^2-Net (U square net) 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

Updates !!!

(2020-Sep-13) Our U^2-Net based model is the 6th in MICCAI 2020 Thyroid Nodule Segmentation Challenge.

(2020-May-18) The official paper of our U^2-Net (U square net) (PDF in elsevier(free until July 5 2020), PDF in arxiv) is now available. If you are not able to access that, please feel free to drop me an email.

(2020-May-16) We fixed the upsampling issue of the network. Now, the model should be able to handle arbitrary input size. (Tips: This modification is to facilitate the retraining of U^2-Net on your own datasets. When using our pre-trained model on SOD datasets, please keep the input size as 320x320 to guarantee the performance.)

(2020-May-16) We highly appreciate Cyril Diagne for building this fantastic AR project: AR Copy and Paste using our U^2-Net (Qin et al, PR 2020) and BASNet(Qin et al, CVPR 2019). The demo video in twitter has achieved over 5M views, which is phenomenal and shows us more application probabilities of SOD.

U^2-Net Results (176.3 MB)

U^2-Net Results

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

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

  2. 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

U^2-Net architecture

Quantitative Comparison

Quantitative Comparison

Quantitative Comparison

Qualitative Comparison

Qualitative Comparison

Citation

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