Update inference.py
Browse files- inference.py +1 -2
inference.py
CHANGED
|
@@ -11,7 +11,6 @@ model = BEN2.BEN_Base().to(device).eval() #init pipeline
|
|
| 11 |
|
| 12 |
model.loadcheckpoints("./BEN2_Base.pth")
|
| 13 |
image = Image.open(file)
|
| 14 |
-
|
| 15 |
|
| 16 |
-
mask.save("./mask.png")
|
| 17 |
foreground.save("./foreground.png")
|
|
|
|
| 11 |
|
| 12 |
model.loadcheckpoints("./BEN2_Base.pth")
|
| 13 |
image = Image.open(file)
|
| 14 |
+
foreground = model.inference(image)
|
| 15 |
|
|
|
|
| 16 |
foreground.save("./foreground.png")
|