How to specify path to custom u2net model and disable pre-multiplied alpha

[alastairpatrick]

In case it helps anyone, this is the code I wrote to drive rembg. The reasons I didn't use the command line tool are because I needed to specify the path to a custom u2net model and because I wanted output without pre-multiplied alpha.

import glob
import os
import rembg
import sys

from PIL import Image

input_dir = sys.argv[1]
output_dir = sys.argv[2]
model_path = '/U-2-Net/saved_models/u2net/u2net.onnx'

session = rembg.sessions.U2netCustomSession("u2net", {}, model_path=model_path)

os.makedirs(output_dir, exist_ok=True)

for input_path in glob.glob(os.path.join(input_dir, "*.png")):
    output_base = os.path.basename(input_path)
    output_path = os.path.join(output_dir, output_base)

    print(output_base)
    input = Image.open(input_path)
    output = rembg.remove(input, session=session, post_process_mask=True, alpha_matting=True)

    # Reverse premultiplied alpha
    input.putalpha(output.getchannel('A'))

    input.save(output_path)