1

I try to load the pre-trained u2net_human_seg.onnx model in my python program to use it for better background removing. When I try it, I get an error: _pickle.UnpicklingError: invalid load key, '\x08'.

My code:

import rembg
from PIL import Image
import torch
import numpy as np

def remove_background(input_path):
    input = Image.open(input_path)
    output = remove(input)
    output.save(input_path)

def remove(input):
    input_tensor = torch.from_numpy(np.array(input)).float()

    output = model(input_tensor)

    output_image = Image.fromarray(output.detach().numpy())
    return output_image


if __name__ == '__main__':
    model = torch.load("models/u2net_human_seg.onnx")
    remove_background("images/test.jpg")

My input is test, where 2 people are visible. The paths' should be correct... I didn't find similar cases online as no one seems to load a model in a custom application.

0 Answers0