U²-Net on tinygrad

This repository contains the tinygrad-ported version of the U²-Net model.

Setup

pip3 install -r requirements.txt

Download weights

The repo contains the small models' weights in weights folder. The large (176 mb) weights can be downloaded from the following links:

foreground segmentation

portrait model

Run inference

python3 u2net_run.py -i /path/to/image_or_folder -m selected_model

Where selected_model is one of:

fg: human segmentation (large variant, 176 mb)

fg_small: human segmentation (small variant, 4.7 mb)

sky_small: sky segmentation (small variant, 4.7 mb)

portrait: converts photo to pencil drawing (only large variant available, 176 mb)

Perform sky replacement

By providing one photo that contains a visible sky and another photo showing only a sky, you can replace the sky in the original photo with the new one:

python3 u2net_run.py -i /path/to/image_with_sky -m sky_small -sky /path/to/new_sky

Examples

Foreground segmentation

Sky segmentation

Sky replacement

License

Apache-2.0 license