File size: 785 Bytes
3aad4a9 1272389 3aad4a9 1272389 3aad4a9 1272389 |
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 |
---
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"/> |