File size: 1,189 Bytes
4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 524ae85 4cb2cd7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
from skimage import io
import torch, os
from PIL import Image
from briarmbg import BriaRMBG
from utilities import preprocess_image, postprocess_image
def example_inference():
model_path = f"{os.path.dirname(__file__)}/model.pth"
im_path = f"{os.path.dirname(__file__)}/example_input.jpg"
net = BriaRMBG()
if torch.cuda.is_available():
net.load_state_dict(torch.load(model_path)).cuda()
else:
net.load_state_dict(torch.load(model_path,map_location="cpu"))
net.eval()
# prepare input
model_input_size = [1024,1024]
orig_im = io.imread(im_path)
orig_im_size = orig_im.shape[0:2]
image = preprocess_image(orig_im, model_input_size)
if torch.cuda.is_available():
image=image.cuda()
# inference
result=net(image)
# post process
result_image = postprocess_image(result[0][0], orig_im_size)
# save result
pil_im = Image.fromarray(result_image)
no_bg_image = Image.new("RGBA", pil_im.size, (0,0,0,0))
orig_image = Image.open(im_path)
no_bg_image.paste(orig_image, mask=pil_im)
no_bg_image.save("example_image_no_bg.png")
if __name__ == "__main__":
example_inference() |