Update README.md
Browse files
README.md
CHANGED
@@ -107,13 +107,15 @@ import torch, os
|
|
107 |
from PIL import Image
|
108 |
from briarmbg import BriaRMBG
|
109 |
from utilities import preprocess_image, postprocess_image
|
|
|
110 |
|
111 |
-
model_path =
|
112 |
im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
|
113 |
|
114 |
net = BriaRMBG()
|
115 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
116 |
net.load_state_dict(torch.load(model_path, map_location=device))
|
|
|
117 |
net.eval()
|
118 |
|
119 |
# prepare input
|
|
|
107 |
from PIL import Image
|
108 |
from briarmbg import BriaRMBG
|
109 |
from utilities import preprocess_image, postprocess_image
|
110 |
+
from huggingface_hub import hf_hub_download
|
111 |
|
112 |
+
model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
|
113 |
im_path = f"{os.path.dirname(os.path.abspath(__file__))}/example_input.jpg"
|
114 |
|
115 |
net = BriaRMBG()
|
116 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
117 |
net.load_state_dict(torch.load(model_path, map_location=device))
|
118 |
+
net.to(device)
|
119 |
net.eval()
|
120 |
|
121 |
# prepare input
|