How to use own pre trained model with rembg?

[OverKilL013]

Hlello, need some help here. I have own model u2net_epoch.pth and converted one u2net_epoch.onnx
I was tried to load model like this:
**
model_name = "u2net_epoch"
session = new_session(model_name)
output = remove(input_image,session=session)
**
Nothing happend, as i understand I need to create file like sessions\u2net_custom.py but I have no idea what to do with download part to make it work localy.
I was tried cli version:
**
rembg i -m u2net_custom -x '{"model_path": "/.u2net/u2net_epoch.onnx"}' path/to/input.png path/to/output.png
rembg i -m u2net_custom -x '{"model_path": "/.u2net/u2net_epoch.pth"}' path/to/input.png path/to/output.png
**
got error "ValueError: model_path is required"
tried to hard code it inside "sessions\u2net_custom.py":
**
model_path = r"J:/my_model/u2net_epoch.onnx" or model_path = r"J:/my_model/u2net_epoch.pth"
**
got another error "NO_SUCHFILE : Load model from None failed:Load model None failed. File doesn't exist"

also tryed to place model files to C:\Users{user}.u2net nothing helps

So what clearly do with u2net_custom.py and how to use it for own models 🤔 would be apriciete for help 🤗

p.s. - rembg works perfectly with no errors with u2net.onnx inside C:\Users{user}.u2net if try to replace it =) it downloads new one of course

win 11
python 3.10 -m venv

[OverKilL013]

Just create file like this, and now you can run rembg with session named "U2netYour".
session = new_session(U2netYour)
remove(input_image, session=session)
Dont forget to import rembg new_session

=== new session file in \rembg\sessions ===
import os
from typing import List

import numpy as np
import onnxruntime as ort
import pooch
from PIL import Image
from PIL.Image import Image as PILImage

from .base import BaseSession

class U2netYourSession(BaseSession):
def init(
self,
model_name: str,
sess_opts: ort.SessionOptions,
providers=None,
*args,
**kwargs
):
#model_path = kwargs.get("model_path")
model_path = r"your\path\to\model\mode.onnx"
print(f"Model path: {model_path}")

    if model_path is None:
        raise ValueError("model_path is required")

    super().__init__(model_name, sess_opts, providers, *args, **kwargs)

def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
    ort_outs = self.inner_session.run(
        None,
        self.normalize(
            img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
        ),
    )

    pred = ort_outs[0][:, 0, :, :]

    ma = np.max(pred)
    mi = np.min(pred)

    pred = (pred - mi) / (ma - mi)
    pred = np.squeeze(pred)

    mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
    mask = mask.resize(img.size, Image.LANCZOS)

    return [mask]

@classmethod
def download_models(cls, *args, **kwargs):
    model_path = r"your\path\to\model\mode.onnx"
    if model_path is None:
        return

    return os.path.abspath(os.path.expanduser(model_path))

@classmethod
def name(cls, *args, **kwargs):
    return "u2net_your"

[jurandfantom]

When you write about session, what exactly you mean? Readme not include such informations so not sure where dig into.

Second thing, I downloaded yolov4 and struggle to load it. I created U2netYour.py file inside of rembg\sessions with content of above reply, yet still not sure what next.

"ust create file like this, and now you can run rembg with session named "U2netYour".
session = new_session(U2netYour)
remove(input_image, session=session)
Dont forget to import rembg new_session"