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