Cómo eliminar el fondo de una imagen con Machine Learning usando Rembg (Python 3) en Ubuntu 20.04

Aprende a usar la increíble biblioteca de Python de Rembg para eliminar el fondo de cualquier imagen con aprendizaje automático.

Existen múltiples herramientas gráficas que puede utilizar para eliminar todo lo que rodea a un objeto en una imagen. Sin embargo, hacer esto automáticamente es bastante difícil de hacer, normalmente un usuario siempre necesitará interactuar con la herramienta para eliminarla lo mejor posible. ¿Qué pasaría si te dijera que hay un proyecto de código abierto increíble que tiene como objetivo eliminar el fondo de una imagen automáticamente mediante el aprendizaje automático?

En este artículo, le explicaré cómo instalar y configurar fácilmente la biblioteca Rembg para Python para eliminar el fondo de una imagen.

1. Instale PyTorch

Si está usando anaconda, puede instalar Pytorch con el siguiente comando (puede aprender cómo instalar Anaconda siguiendo este tutorial en caso de que no lo tenga instalado y quiera usarlo):

conda install pytorch torchvision torchaudio cpuonly -c pytorch

Alternativamente, puedes instalar Pytorch usando pip con el siguiente comando:

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

Nota: para este tutorial y también explicado en el repositorio de github, estamos usando Pytorch para CPU. Si tiene una tarjeta gráfica disponible en el servidor, puede instalar Pytorch con soporte para CUDA. Para obtener más información sobre Pytorch, visite el sitio web oficial aquí.

2. Instala RemBG

RemBG es una herramienta para eliminar el fondo de cualquier imagen. Esta biblioteca utiliza U2-Net bajo el capó para el reconocimiento de patrones (" U ^ 2-Net: Profundizando con estructura en U anidada para la detección de objetos salientes "). Para instalar este paquete en su sistema, use pip:

pip install rembg

Para obtener más información sobre esta increíble herramienta de código abierto, visite el repositorio oficial de Github aquí .

3. Uso de la biblioteca

Después de instalar la biblioteca Rembg, podrá elegir directamente desde la línea de comandos o dentro de un simple script de Python. Tenga en cuenta que cuando se ejecuta por primera vez, el modelo se descargará, por lo que tomará un tiempo, una vez que finalice, la eliminación del fondo no requiere tanto para cada imagen. Rembg estará disponible globalmente en su entorno.

Eliminación de antecedentes a través de la CLI

La forma más rápida de probar si la biblioteca está funcionando es simplemente usar una imagen remota, el siguiente comando usará una imagen de Wikipedia de 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

El archivo generado output.png se verá así ( considerando la imagen original ):

Realmente asombroso, ¿no? En algunos servidores, especialmente los virtualizados (instrucciones avx ni avx2), puede aparecer la siguiente advertencia:

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

Pero puedes ignorarlo ya que la imagen de salida ya no tiene fondo. Si necesitas eliminar la advertencia mencionada, es posible que debas compilar Pytorch desde la fuente y deshabilitar la compatibilidad con NNPack . Por supuesto, puedes hacer lo mismo con imágenes locales:

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

Usando la biblioteca dentro de una secuencia de comandos de Python

También puedes usar la biblioteca junto con su canalización incrustándola en sus scripts. Por ejemplo, el siguiente script hará lo mismo que el ejemplo anterior, tomará un archivo local y creará la versión sin fondo:

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

# Descomente las siguientes líneas si trabaja con formatos de imagen trucados (por ejemplo, JPEG / JPG)
# En mi caso, doy imágenes JPEG como entrada, así que lo dejaré sin comentar.
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

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")
img.save(output_image)

Pruébalo y comienza a eliminar el fondo de tus imágenes automáticamente con esta maravillosa herramienta.

En caso de error usando rembg desde la terminal

En mi case, cuando use por primera vez rembg directamente desde la linea de comandos en producción luego de la instalación, este lanzo el siguiente error 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>
    compile_modules()
  File "/usr/local/lib/python3.9/dist-packages/pymatting_aot/cc.py", line 16, in compile_modules
    cc.compile()
  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

Para solucionarlo simplemente tuve que instalar el paquete python3-dev (en mi caso para Python 3.9):