timm
PyTorch
medical
Image Feature Extraction
Ege Oezsoy commited on
Commit
9841873
1 Parent(s): ccaaa60

Adjustments

Browse files
Files changed (1) hide show
  1. endovit_online.py +0 -43
endovit_online.py DELETED
@@ -1,43 +0,0 @@
1
- import torch
2
- from pathlib import Path
3
- from timm.models.vision_transformer import VisionTransformer
4
- from functools import partial
5
- from torch import nn
6
- from huggingface_hub import snapshot_download
7
-
8
- def load_model_from_huggingface(repo_id, model_filename):
9
- # Download model files
10
- model_path = snapshot_download(repo_id=repo_id, revision="main")
11
- model_weights_path = Path(model_path) / model_filename
12
-
13
- # Load model weights
14
- model_weights = torch.load(model_weights_path)['model']
15
-
16
- # Define the model (ensure this matches your model's architecture)
17
- 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)).eval()
18
-
19
- # Load the weights into the model
20
- loading = model.load_state_dict(model_weights, strict=False)
21
-
22
- return model, loading
23
- 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]):
24
- # Define the transformations
25
- transform = T.Compose([
26
- T.Resize((input_size, input_size)),
27
- T.ToTensor(),
28
- T.Normalize(mean=dataset_mean, std=dataset_std)
29
- ])
30
-
31
- # Open the image
32
- image = Image.open(image_path).convert('RGB')
33
-
34
- # Apply the transformations
35
- processed_image = transform(image)
36
-
37
- return processed_image
38
-
39
- device = "cuda"
40
- dtype = torch.float16
41
- model, loading_info = load_model_from_huggingface("egeozsoy/EndoViT", "endovit.pth")
42
- model = model.to(device, dtype)
43
- print(loading_info)