/rembg-trainer

Code to train U2Net model for use with rembg tool

Primary LanguagePythonMIT LicenseMIT

rembg trainer

This code allows you to easily train U2-Net model in ONNX format to use with rembg tool.

This work is based off U2Net repo, which is under Apache licence. The derivative work is loicensed under MIT; do as you please with it.

A couple of notes on performance:

  • Default parameters are fine-tuned for maximum performance on systems with 32gb of processing memory, like the Apple M1 Pro. Adjust accordingly.
  • Computations are performed in float32, because float16 support on Metal is a bit undercooked at the moment.
  • If this is your first time using CUDA on Windows, you'd have to install CUDA Toolkit.
  • For CUDA, you can easily rewrite this code with half precision calculations for increased performance. Apex library can help you with that; I don't have such plans at the moment.
  • For acceleration on AMD GPUs, please refer to installation guide of AMD ROCm platform. No code changes will be required.

If the training is interrupted for any reason, don't worry — the program saves its state regularly, allowing you to resume from where you left off. Frequency of saving can be adjusted.

Fancy a go?

  • Download the latest release
  • Install requirements.txt
  • Put your images into images folder
  • Put their masks into masks folder; or see below
  • Launch python3 u2net_train.py --help for more details on supported command line flags
  • Launch script with your desired configuration
  • Go grab yourself a nice latte and wait........... and wait.....
  • Once you've had your fill of waiting, here's how you use resulting model with rembg:
rembg p -w input output -m u2net_custom -x '{"model_path": "/saved_models/u2net/27.onnx"}'
# input — folder with images to have their backgrounds removed
# output — folder for resulting images processed with custom model
# adjust path(s) as necessary!

Mask extraction

If you already have a bunch of images with removed background, then you can create masks off them using the provided alpha.py script. Create a directory called clean, put your pngs there, and launch the script.

But fair warning mate: the script is very CPU-heavy. Oh, and you'll need the ImageMagick tool installed and present in your PATH.

So, at the end of the day, you will end up with the following folder structure:

  • images — source images, will be needed for training
  • masks — required for training, to teach model where the background was
  • clean — images with removed background, to extract masks (they're not used for actual training)

Leave your mark 👉👈🥺

Buy me a coffee an alcohol-free cider here