Update README.md

README.md

@@ -1,3 +1,25 @@
----
-license: mit
----
+---
+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")
+```