BEN / inference.py
MaxwellMeyer's picture
Created inference file, moved torch no grad to model.py, removed timm user warning and used a different photo from the demo for default inference image.
2ca249a
raw
history blame
384 Bytes
import model
from PIL import Image
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
file = "./image.png" # input image
model = model.BEN_Base().to(device).eval() #init pipeline
model.loadcheckpoints("./BEN_Base.pth")
image = Image.open(file)
mask, foreground = model.inference(image)
mask.save("./mask.png")
foreground.save("./foreground.png")