Learn how to use the awesome Rembg python library to remove the background of any image with machine learning.

How to remove the background of a picture with Machine Learning using Rembg (Python 3) in Ubuntu 20.04

There are multiple graphical tools out there that you can use to remove everything that surrounds an object in an image. However, doing this automatically it's quite difficult to do, normally an user will always need to interact with the tool to remove it so good as possible. What if I told you that there's an awesome open source project that aims to remove the background of an image automatically using machine learning?

In this article, I will explain to you how to easily install and configure the Rembg library for Python to remove the background from an image.

1. Install PyTorch

If you are using anaconda, you can install Pytorch with the following command (you can learn how to install Anaconda following this tutorial just in case that you don't have it installed and want to use it):

conda install pytorch torchvision torchaudio cpuonly -c pytorch

Alternatively, you may install Pytorch using pip with the following command:

pip install torch==1.7.1+cpu torchvision==0.8.2+cpu -f https://download.pytorch.org/whl/torch_stable.html

Note: for this tutorial and explained as well in the github repository, we are using Pytorch for CPU. If you have a graphic card available in the server, you may install Pytorch with support for CUDA. For more information about Pytorch, please visit the official website here.

2. Install RemBG

RemBG is a tool to remove the background of any image. This library uses the U2-Net under the hood for pattern recognition ("U^2-Net: Going Deeper with Nested U-Structure for Salient Object Detection."). To install this package in your system, use pip:

pip install rembg

For more information about this awesome open source tool, please visit the official repository at Github here.

3. Using the library

After installing the Rembg library, you will be able to choose either directly from the command line or inside a simple python script. Note that when running for first time, the model will be downloaded so it will take a while, once it finishes, the background removal doesn't take that much for every image. Rembg will be available globally in your environment.

Background removal through the CLI

The fastest way to test if the library is working is to simply use a remote image, the following command will use a Wikipedia image of Julian Casablancas:

curl -s https://upload.wikimedia.org/wikipedia/commons/thumb/e/ec/Julian_Casablancas_%40_Lollapalooza_2014_%2813657827913%29.jpg/440px-Julian_Casablancas_%40_Lollapalooza_2014_%2813657827913%29.jpg | rembg > output.png

The generated output.png will look like this (considering the original image):

Julian Casablancas PNG

Really awesome isn't? In some servers, specially the virtualized ones (avx nor avx2 instructions), you may get the following warning:

[W NNPACK.cpp:79] Could not initialize NNPACK! Reason: Unsupported hardware.

But you can ignore it as the output image hasn't background anymore. If you need to remove the mentioned warning, you may need to compile Pytorch from source and disable NNPack support. You can of course do the same with local images:

rembg -o ./output-image.png ./input-image.jpg

Using the library inside a Python script

You can as well use the library along with your pipeline embedding it into your scripts. For example, the following script will do the same as the previous example, take a local file and create the version without background:

# example.py
from rembg.bg import remove
import numpy as np
import io
from PIL import Image

# Uncomment the following lines if working with trucated image formats (ex. JPEG / JPG)
# In my case I do give JPEG images as input, so i'll leave it uncommented
from PIL import ImageFile

input_image = 'input-image.jpg'
output_image = 'output-image.png'

f = np.fromfile(input_image)
result = remove(f)
img = Image.open(io.BytesIO(result)).convert("RGBA")

Give it a try and start removing the background of your images automatically using this wonderful tool!

In case of error when using rembg

In my case, when using rembg directly from the command line in production after the installation, it throwed the following distutils.errors.CompileError:

Failed to import ahead-of-time-compiled modules. This is expected on first import.
Compiling modules and trying again (this might take a minute).
Traceback (most recent call last):
  File "/usr/local/bin/rembg", line 5, in <module>
    from rembg.cli import main
  File "/usr/local/lib/python3.9/dist-packages/rembg/cli.py", line 9, in <module>
    from .bg import remove
  File "/usr/local/lib/python3.9/dist-packages/rembg/bg.py", line 6, in <module>
    from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
  File "/usr/local/lib/python3.9/dist-packages/pymatting/__init__.py", line 2, in <module>
    import pymatting_aot.cc
  File "/usr/local/lib/python3.9/dist-packages/pymatting_aot/cc.py", line 28, in <module>
  File "/usr/local/lib/python3.9/dist-packages/pymatting_aot/cc.py", line 16, in compile_modules
  File "/usr/local/lib/python3.9/dist-packages/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/numba/pycc/cc.py", line 217, in compile
    objects += self._compile_mixins(build_dir)
  File "/usr/local/lib/python3.9/dist-packages/numba/pycc/cc.py", line 187, in _compile_mixins
    objects = self._toolchain.compile_objects(sources, build_dir,
  File "/usr/local/lib/python3.9/dist-packages/numba/pycc/platform.py", line 137, in compile_objects
    objects = self._compiler.compile(sources,
  File "/usr/local/lib/python3.9/dist-packages/numpy/distutils/ccompiler.py", line 89, in <lambda>
    m = lambda self, *args, **kw: func(self, *args, **kw)
  File "/usr/local/lib/python3.9/dist-packages/numpy/distutils/ccompiler.py", line 360, in CCompiler_compile
    pool.map(single_compile, build_items)
  File "/usr/lib/python3.9/multiprocessing/pool.py", line 364, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/usr/lib/python3.9/multiprocessing/pool.py", line 771, in get
    raise self._value
  File "/usr/lib/python3.9/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/usr/lib/python3.9/multiprocessing/pool.py", line 48, in mapstar
    return list(map(*args))
  File "/usr/local/lib/python3.9/dist-packages/numpy/distutils/ccompiler.py", line 325, in single_compile
    self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
  File "/usr/local/lib/python3.9/dist-packages/numpy/distutils/ccompiler.py", line 89, in <lambda>
    m = lambda self, *args, **kw: func(self, *args, **kw)
  File "/usr/local/lib/python3.9/dist-packages/numpy/distutils/unixccompiler.py", line 58, in UnixCCompiler__compile
    raise CompileError(msg) from None
distutils.errors.CompileError: Command "x86_64-linux-gnu-gcc -pthread -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -fPIC -DPYCC_MODULE_NAME=aot -DPYCC_USE_NRT=1 -I/usr/include/python3.9 -I/usr/local/lib/python3.9/dist-packages/numpy/core/include -c /usr/local/lib/python3.9/dist-packages/numba/pycc/modulemixin.c -o /tmp/pycc-build-aot-2b86dl9a/usr/local/lib/python3.9/dist-packages/numba/pycc/modulemixin.o" failed with exit status 1

To solve this problem, I simply had to install the python3-dev package (in my case Python 3.9):

sudo apt-get install python3.9-dev

Happy coding ❤️!

Senior Software Engineer at Software Medico. Interested in programming since he was 14 years old, Carlos is a self-taught programmer and founder and author of most of the articles at Our Code World.