U^2-Net

✂️ Simplified version of U^2-Net for testing purposes

Motivation

This repository came up after seeing the amazing work that the University of Alberta did with the U^2-Net model for Salient Object detection. I wanted to understand how it actually works and build a easier script for running inference in one or several images to make this more accessible to the community.

Examples

Python requirements

This project is using Python3.7. All these requirements have been specified in the requirements.lock file.

1. numpy 1.15.2
2. scikit-image 0.14.0
3. Pillow 5.2.0
4. torch 0.4.1
5. torchvision 0.2.1

Environment

To create your environment and testing your own images, you must follow the next scripts:

1. Pull models from Git LFS

   git lfs fetch --all

2. Create a virutalenv, where a usage of virtualenv is recommended for package library / runtime isolation.

   python -m pip install --user virtualenv
   python3 -m virtualenv --python=/usr/bin/python3 /opt/venv
   source /opt/venv/bin/activate

3. Install dependencies

   pip3 install -r requirements.lock

Usage

There are two ways of testing this out. Using Python, after building the enviornment mentioned above, or using the Docker image provided.

Python

Simple as running the following script:

python3 -m src.run INPUT_PATH OUTPUT_PATH [--model MODEL] [--gpu] [--show_user_warnings]

• INPUT_PATH: existing path pointing to a single image or a folder containing several images to process.
• OUTPUT_PATH: path to a folder where put the generated images.
• MODEL: model to be used for running inference, which could be u2net (default) or u2netp (a smaller version).
• GPU: if this parameter is enabled, the script will try to find CUDA 8 installed for running use GPU accelaration.
• SHOW_USER_WARNINGS: if this parameter is enabled, user warnings will be showed during the execution.

Docker

Before doing anything, you need to configure your local Docker installation for authenticating to GitHub Packages.

Once that, you will just need to run the following command:

docker run \
	-e "INPUT_PATH=input_path" \
	-e "OUTPUT_PATH=output_path" \
	-v "input_path:input_path" \
	-v "output_path:output_path" \
	--name u_2_net \
	docker.pkg.github.com/albertsuarez/u-2-net/u-2-net

INPUT_PATH and OUTPUT_PATH are required environment variables for running the image. Volumes are needed as well because input files are not inside the Docker base image and for collecting the results from the host.

Merits

Again, kudos to Xuebin Qin, Zichen Zhang, Chenyang Huang, Masood Dehghan, Osmar R. Zaiane and Martin Jagersand for this amazing model that they came up.