melmass/pytorch-scripts

Torch Scripts

VitMatte

The demo files are here

import torch
from torchvision.transforms import functional as F

image = Image.open("./retriever_rgb.png").convert("RGB")
image = F.to_tensor(image).unsqueeze(0).to("cuda").half()

trimap = Image.open("./retriever_trimap.png").convert("L")
trimap = F.to_tensor(trimap).unsqueeze(0).to("cuda").half()

input = {"image": image, "trimap": trimap}

model = torch.jit.load("./vitmatte_b_dis.pt").to("cuda")
alpha = model(input)

output = F.to_pil_image(predictions)
output.save("./predicted.png")

This is the output: