Use my own model path when using as a library

[Garyyyyyyyyyyyyy]

Hi guys!

I'm a big fan of rembg. However, every time I use it, I have to download the U2NET model from GitHub, which adds some delay. I'd love to specify a local path for the model to avoid downloading it repeatedly, but I'm not sure how to do this when using rembg as a library.

I installed rembg using pip and I'm currently on last version.

Here's the code I'm using:

def remove_bg(image_url):
    response = requests.get(image_url)
    input_image = response.content
    output_image = rembg.remove(input_image)
    image = Image.open(io.BytesIO(output_image))
    image_array = np.array(image)
    image_array = image_array[:, :, :3]
    return(image_array)

Any guidance on this would be greatly appreciated. Thanks in advance!

[AlfredUFY]

To use rembg with a locally stored U2NET model, one approach is to set the model path using the environment variable U2NET_PATH. This will direct rembg to utilize the U2NET model at the specified local path instead of downloading it again from the Internet.

You can set environment variables in Python using the os module. Here's an example of how you can modify your code to include this:

import os
import rembg
import requests
import numpy as np
from PIL import Image
import io

def remove_bg(image_url):
    os.environ['U2NET_PATH'] = 'path/to/your/local/model'  # Set the U2NET model path
    response = requests.get(image_url)
    input_image = response.content
    output_image = rembg.remove(input_image)
    image = Image.open(io.BytesIO(output_image))
    image_array = np.array(image)
    image_array = image_array[:, :, :3]
    return image_array

Please replace 'path/to/your/local/model' with the actual local path to your U2NET model. Remember, if the path changes or the program is restarted, you will need to reset it. Also make sure that the local U2NET model is compatible with the rembg version you’re currently using.

Note: The path needs to point to the .pth file directly. Also, check if your environment has write permissions to the folder where the model file exists because rembg might do some caching operations.

[rapid537]

according to #568 (comment)
the environment variable should actually be U2NET_HOME

i would recommend getting a reference to your project directory first and then build the path from there.
then just drop the folder containing your model into that directory.
e.g.

import os

def remove_bg(image_url):
    basedir = os.path.abspath(os.path.dirname(__file__))
    os.environ['U2NET_HOME'] = basedir + '/u2net/u2net.onnx'
    ...