Initial model commit

endovit.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec41888fd928eb404e518d61344c25822b3c4a776f997b876db040b86a5ba21a
+size 1341188699

endovit_demo.py

@@ -0,0 +1,39 @@
+import torch
+import torchvision.transforms as T
+from PIL import Image
+from pathlib import Path
+from timm.models.vision_transformer import VisionTransformer
+from functools import partial
+from torch import nn
+
+# requires: pytorch 2.0.1, timm 0.9.16
+def process_single_image(image_path, input_size=224, dataset_mean=[0.3464, 0.2280, 0.2228], dataset_std=[0.2520, 0.2128, 0.2093]):
+    # Define the transformations
+    transform = T.Compose([
+        T.Resize((input_size, input_size)),
+        T.ToTensor(),
+        T.Normalize(mean=dataset_mean, std=dataset_std)
+    ])
+
+    # Open the image
+    image = Image.open(image_path).convert('RGB')
+
+    # Apply the transformations
+    processed_image = transform(image)
+
+    return processed_image
+
+
+image_paths = sorted(Path('demo_images').glob('*.png'))
+images = torch.stack([process_single_image(image_path) for image_path in image_paths])
+
+device = "cuda"
+dtype = torch.float16
+
+model_weights = torch.load('endovit_seg.pth')['model']
+
+model = VisionTransformer(patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6)).to(device, dtype).eval()
+loading = model.load_state_dict(model_weights, strict=False)
+print(loading)
+output = model.forward_features(images.to(device, dtype))
+print(output.shape)

requirements.txt

@@ -0,0 +1,2 @@
+torch==2.0.1
+timm==0.9.16