danielgatis/rembg

[BUG] ... KeyError: 'in_ch' RMBG-1.4

f-amerehi opened this issue · 2 comments

Describe the bug

Following the usage from HF. When loading the model via net = BriaRMBG.from_pretrained("briaai/RMBG-1.4") it returns a KeyError for 'in_ch' despite the key is available.

from skimage import io
import torch
from PIL import Image
from briarmbg import BriaRMBG
from utilities import preprocess_image, postprocess_image
import matplotlib.pyplot as plt



net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.to(device)

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[<ipython-input-6-08fbcf1bcc4f>](https://localhost:8080/#) in <cell line: 1>()
----> 1 net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
      2 
      3 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
      4 net.to(device)

3 frames
[/content/RMBG-1.4/briarmbg.py](https://localhost:8080/#) in __init__(self, config)
    350     def __init__(self,config:dict={"in_ch":3,"out_ch":1}):
    351         super(BriaRMBG,self).__init__()
--> 352         in_ch=config["in_ch"]
    353         out_ch=config["out_ch"]
    354         self.conv_in = nn.Conv2d(in_ch,64,3,stride=2,padding=1)

KeyError: 'in_ch'

But the dictionary has the key of 3. How to fix the error?
image

Thanks!

This issue is stale because it has been open for 30 days with no activity.

This issue was closed because it has been inactive for 14 days since being marked as stale.