Implementation of U^2-Net(U Square Net) architecture in TensorFlow for salient object detection and segmentation.

This repository contains the tensorflow implementation of the below paper. The original implementation was done on PyTorch.

U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection

The model has been trained and tested on the DUTS and ECSSD dataset using tf v2.4.1 and keras v2.4.3. The weights can be downloaded from here : DUTS_main, DUTS_light, ECSSD_main, ECSSD_light.

Predicted mask | Original mask:

If you find this useful please cite the original paper using:

@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}
}

Update (29/11/2020):

Added Scripts used for training on larger datasets like DUTS which could cause memory leak. The updated scripts can be modified according to the dataset used.