danielgatis / rembg

Rembg is a tool to remove images background
MIT License
15.46k stars 1.76k forks source link

[BUG] Failed to achieve varied output results with different inputs using SAM model #550

Closed HuangMason320 closed 6 months ago

HuangMason320 commented 7 months ago

Describe the bug Unable to have different result output with different input using SAM model

To Reproduce

import rembg
from PIL import Image
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import onnxruntime

input_points_list =np.array([[1241, 620],[1494, 954],[1490,800]]) 
input_label_list = np.array([[1,1,1]])

input_path = '/Users/masonhuang/Pictures/Download/15257490252132.jpg'
output_path = '/Users/masonhuang/Code/Python/rembg-main/examples/test.png'

image = Image.open(input_path)

model_name = "sam_prompt"

session = rembg.new_session(model_name)
output = rembg.remove(image, session=session, input_points = np.array(input_points_list), input_labels = np.array(input_label_list))

print(f'The image is saved')
output.save(output_path)

i've got this when running code.

[W:onnxruntime:, coreml_execution_provider.cc:81 GetCapability] CoreMLExecutionProvider::GetCapability, number of partitions supported by CoreML: 67 number of nodes in the graph: 371 number of nodes supported by CoreML: 300

Images Input 15257490252132 Output test

OS Version: MacOS 14.1.2 Python 3.11 Pytorch 2.0.1

Rembg version: v2.0.53

Additional context

I have a custom SAM model(.onnx file) but i'm not sure where to modify

Besides, I've tried to change the code below:

sam.py

from

model_name = kwargs.get("sam_model", "sam_vit_b_01ec64")

to

model_name = kwargs.get("sam_model", "sam_vit_h_4b8939")

session_factory.py

from

model_name: str = "sam_prompt", providers=None, *args, **kwargs
session_class: Type[BaseSession] = U2netSession

to

model_name: str = "u2net", providers=None, *args, **kwargs
session_class: Type[BaseSession] = SamSession

bg.py

from

if session is None:
        session = new_session("u2net", *args, **kwargs)

to

if session is None:
        session = new_session("sam_prompt", *args, **kwargs)

when I apply above change, I got error message:

[W:onnxruntime:, coreml_execution_provider.cc:81 GetCapability] CoreMLExecutionProvider::GetCapability, number of partitions supported by CoreML: 107 number of nodes in the graph: 607 number of nodes supported by CoreML: 358
2023-12-07 10:57:00.072353
 [W:onnxruntime:, coreml_execution_provider.cc:81 GetCapability] CoreMLExecutionProvider::GetCapability, number of partitions supported by CoreML: 26 number of nodes in the graph: 487 number of nodes supported by CoreML: 112

jsonschema.exce
ptions.ValidationError: '{}' is not of type 'array'

Failed validating 'type' in schema:
    {'items': {'properties': {'data': {'items': {'type': 'number'},
                                       'type': 'array'},
                              'label': {'type': 'integer'},
                              'type': {'type': 'string'}},
               'type': 'object'},
     'type': 'array'}

On instance:
    '{}'

Could you tell me. Thank you

cyqresig commented 5 months ago

[W:onnxruntime:, coreml_execution_provider.cc:81 GetCapability] CoreMLExecutionProvider::GetCapability, number of partitions supported by CoreML: 67 number of nodes in the graph: 371 number of nodes supported by CoreML: 300

I've got this on Mac env, too.

I found that the version of onnxruntime need to be downgrade(e.g. onnxruntime==1.16.3)

FergusKidd commented 4 months ago

[W:onnxruntime:, coreml_execution_provider.cc:81 GetCapability] CoreMLExecutionProvider::GetCapability, number of partitions supported by CoreML: 67 number of nodes in the graph: 371 number of nodes supported by CoreML: 300

I've got this on Mac env, too.

I found that the version of onnxruntime need to be downgrade(e.g. onnxruntime==1.16.3)

Fixed mine, thanks