test VIT

predict.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.init as init
-from transformers import SwinForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 from huggingface_hub import hf_hub_download
 from PIL import Image
 import json
@@ -14,52 +14,50 @@ with open("labels.json", "r") as f:
     class_names = json.load(f)
 print("class_names:", class_names)
 
-MODEL_NAME = "microsoft/swin-large-patch4-window7-224"  # or your preferred Swin model
-
-class SwinCustom(nn.Module):
-    def __init__(self, model_name=MODEL_NAME, num_classes=40):
-        super(SwinCustom, self).__init__()
-        self.model = SwinForImageClassification.from_pretrained(model_name, num_labels=num_classes, ignore_mismatched_sizes=True)
+class ViT(nn.Module):
+   
+    def __init__(self, model_name="google/vit-base-patch16-224", num_classes=40, dropout_rate=0.1):
+        super(ViT, self).__init__()
+        self.extractor = AutoImageProcessor.from_pretrained(model_name)
+        self.model = AutoModelForImageClassification.from_pretrained(model_name)
         in_features = self.model.classifier.in_features
         self.model.classifier = nn.Sequential(
-            nn.Linear(in_features, in_features),
-            nn.LeakyReLU(),
-            nn.Dropout(0.3),
+            nn.Dropout(p=dropout_rate),
             nn.Linear(in_features, num_classes)
         )
-        # Weight initialization
-        for m in self.model.classifier:
-            if isinstance(m, nn.Linear):
-                init.kaiming_uniform_(m.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
+        self.img_size = (self.extractor.size['height'], self.extractor.size['width'])
+        self.normalize = transforms.Normalize(mean=self.extractor.image_mean, std=self.extractor.image_std)
 
     def forward(self, images):
-        outputs = self.model(images)
+        outputs = self.model(pixel_values=images)
         return outputs.logits
 
+    def get_test_transforms(self):
+        return transforms.Compose([
+            transforms.Resize(self.img_size),
+            transforms.ToTensor(),
+            self.normalize
+        ])
+
 # Download your fine-tuned model checkpoint from the Hub
-model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="large_swin_best_model.pth")
+model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="vit_best_model.pth")
 print("Model path:", model_path)
-model = SwinCustom(model_name=MODEL_NAME, num_classes=40)
+model = ViT(model_name="google/vit-base-patch16-224", num_classes=40)
 state_dict = torch.load(model_path, map_location="cpu")
 if "model_state_dict" in state_dict:
     state_dict = state_dict["model_state_dict"]
 model.load_state_dict(state_dict, strict=False)
 model.eval()
 
-# Preprocessing (Swin default)
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
+transform = model.get_test_transforms()
 
 def predict(image_path):
     image = Image.open(image_path).convert("RGB")
     x = transform(image).unsqueeze(0)
     with torch.no_grad():
         outputs = model(x)
-        print("Logits:", outputs)
         probs = torch.nn.functional.softmax(outputs, dim=1)[0]
+        print("Logits:", outputs)
         print("Probs:", probs)
         print("Sum of probs:", probs.sum())
         top5 = torch.topk(probs, k=5)