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.