---
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"/>