| --- |
| license: mit |
| --- |
| |
| # Torch Scripts |
|
|
| ## VitMatte |
| The demo files are [here](https://github.com/hustvl/ViTMatte/tree/main/demo) |
|
|
|
|
|
|
| ```python |
| 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: |
| <img width=400 src="https://cdn-uploads.huggingface.co/production/uploads/62ff65702979d8fc339b0905/2LdOFka5RXucDCWzzVxwl.png"/> |
