File size: 583 Bytes
3aad4a9 |
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 |
---
license: mit
---
# Torch Scripts
## VitMatte
The demo files are [here](https://github.com/hustvl/ViTMatte/tree/main/demo)
```python
import torch
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")
``` |