|
|
--- |
|
|
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") |
|
|
``` |
