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