Upload folder using huggingface_hub
Browse files
app.py
CHANGED
@@ -27,13 +27,15 @@ def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
|
|
27 |
return im_array
|
28 |
|
29 |
|
30 |
-
def inference(
|
31 |
|
32 |
model_path = "ormbg.pth"
|
33 |
|
34 |
net = ORMBG()
|
35 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
36 |
|
|
|
|
|
37 |
if torch.cuda.is_available():
|
38 |
net.load_state_dict(torch.load(model_path))
|
39 |
net = net.cuda()
|
|
|
27 |
return im_array
|
28 |
|
29 |
|
30 |
+
def inference(image):
|
31 |
|
32 |
model_path = "ormbg.pth"
|
33 |
|
34 |
net = ORMBG()
|
35 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
36 |
|
37 |
+
orig_image = Image.fromarray(image)
|
38 |
+
|
39 |
if torch.cuda.is_available():
|
40 |
net.load_state_dict(torch.load(model_path))
|
41 |
net = net.cuda()
|