change to final best model

app.py

@@ -9,7 +9,7 @@ demo = gr.Interface(
         gr.Image(label="Top-1 Class Example"),
         gr.Label(label="Top-5 Probabilities")
     ],
-    title="Scene Classification with Reference Image Testing using VIT version 1.0",
+    title="Scene Classification with Reference Image Testing using Large SWIN version 1.0",
     description="Upload an image to get the predicted class with a sample image and top-5 prediction chart."
 )
 

predict.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.init as init
-from transformers import AutoImageProcessor, AutoModelForImageClassification
+from transformers import SwinForImageClassification
 from huggingface_hub import hf_hub_download
 from PIL import Image
 import json
@@ -14,50 +14,51 @@ with open("labels.json", "r") as f:
     class_names = json.load(f)
 print("class_names:", class_names)
 
-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)
+MODEL_NAME = "microsoft/swin-large-patch4-window7-224"  
+
+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)
         in_features = self.model.classifier.in_features
         self.model.classifier = nn.Sequential(
-            nn.Dropout(p=dropout_rate),
+            nn.Linear(in_features, in_features),
+            nn.LeakyReLU(),
+            nn.Dropout(0.3),
             nn.Linear(in_features, num_classes)
         )
-        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)
+        # 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')
 
     def forward(self, images):
-        outputs = self.model(pixel_values=images)
+        outputs = self.model(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="vit_best_model.pth")
+model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="large_swin_best_model.pth")
 print("Model path:", model_path)
-model = ViT(model_name="google/vit-base-patch16-224", num_classes=40)
+model = SwinCustom(model_name=MODEL_NAME, 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()
 
-transform = model.get_test_transforms()
+# Preprocessing
+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])
+])
 
 def predict(image_path):
     image = Image.open(image_path).convert("RGB")
     x = transform(image).unsqueeze(0)
     with torch.no_grad():
         outputs = model(x)
-        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
         print("Logits:", outputs)
+        probs = torch.nn.functional.softmax(outputs, dim=1)[0]
         print("Probs:", probs)
         print("Sum of probs:", probs.sum())
         top5 = torch.topk(probs, k=5)