Rembg (AWS Lambda)

This is a stripped-down fork of danielgatis/rembg designed for AWS Lambda environments.

rembg-aws-lambda is a tool to remove images background.

Check out my similar project, profile-photo, which can create a headshot from an image.

Requirements

python: >3.7, <3.11

Installation

CPU support:

pip install rembg-aws-lambda

GPU support:

First of all, you need to check if your system supports the onnxruntime-gpu.

Go to https://onnxruntime.ai and check the installation matrix.

If yes, just run:

pip install rembg-aws-lambda[gpu]

Usage as a library

Input and output as bytes

from rembg import remove

input_path = 'input.png'
output_path = 'output.png'

with open(input_path, 'rb') as i:
    with open(output_path, 'wb') as o:
        input = i.read()
        output = remove(input)
        o.write(output)

Input and output as a PIL image

from rembg import remove
from PIL import Image

input_path = 'input.png'
output_path = 'output.png'

input = Image.open(input_path)
output = remove(input)
output.save(output_path)

Input and output as a numpy array

from rembg import remove
import cv2

input_path = 'input.png'
output_path = 'output.png'

input = cv2.imread(input_path)
output = remove(input)
cv2.imwrite(output_path, output)

How to iterate over files in a performatic way

from pathlib import Path
from rembg import remove, new_session

session = new_session()

for file in Path('path/to/folder').glob('*.png'):
    input_path = str(file)
    output_path = str(file.parent / (file.stem + ".out.png"))

    with open(input_path, 'rb') as i:
        with open(output_path, 'wb') as o:
            input = i.read()
            output = remove(input, session=session)
            o.write(output)

Models

All models are downloaded and saved in the user home folder in the .u2net directory.

The available models are:

• u2net (download, source): A pre-trained model for general use cases.
• u2netp (download, source): A lightweight version of u2net model.
• u2net_human_seg (download, source): A pre-trained model for human segmentation.
• u2net_cloth_seg (download, source): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body.
• silueta (download, source): Same as u2net but the size is reduced to 43Mb.

How to train your own model

If You need more fine tunned models try this: danielgatis/rembg#193 (comment)

Some video tutorials

• https://www.youtube.com/watch?v=3xqwpXjxyMQ
• https://www.youtube.com/watch?v=dFKRGXdkGJU
• https://www.youtube.com/watch?v=Ai-BS_T7yjE
• https://www.youtube.com/watch?v=dFKRGXdkGJU
• https://www.youtube.com/watch?v=D7W-C0urVcQ

References

• https://arxiv.org/pdf/2005.09007.pdf
• https://github.com/NathanUA/U-2-Net
• https://github.com/pymatting/pymatting

Buy me a coffee

Liked some of my work? Buy me a coffee (or more likely a beer)

License

Copyright:

• (c) 2020-present Daniel Gatis
• (c) 2023-present Ritvik Nag

Licensed under MIT License