test large swin

predict.py

@@ -1,59 +1,61 @@
 import torch
-from torchvision import transforms
+import torch.nn as nn
+import torch.nn.init as init
+from transformers import SwinForImageClassification
+from huggingface_hub import hf_hub_download
 from PIL import Image
 import json
-import numpy as np
-# from model import load_model
-from transformers import AutoImageProcessor, SwinForImageClassification, ViTForImageClassification
-import torch.nn as nn
 import os
-import pandas as pd
 import random
-from huggingface_hub import hf_hub_download
+from torchvision import transforms
 
 # Load labels
 with open("labels.json", "r") as f:
     class_names = json.load(f)
 print("class_names:", class_names)
 
-class DeiT(nn.Module):
-    def __init__(self, model_name="facebook/deit-base-patch16-224", num_classes=40):
-        super(DeiT, self).__init__()
-        self.model = ViTForImageClassification.from_pretrained(model_name)
+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)
         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.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')
 
     def forward(self, images):
         outputs = self.model(images)
         return outputs.logits
 
-# Load model
-model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_base_best_model.pth")
+# 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")
 print("Model path:", model_path)
-model = DeiT(model_name="facebook/deit-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()
 
-
-#Swin
+# 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])
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-
 def predict(image_path):
-    # Load and prepare image
     image = Image.open(image_path).convert("RGB")
     x = transform(image).unsqueeze(0)
-
     with torch.no_grad():
         outputs = model(x)
         print("Logits:", outputs)
@@ -69,7 +71,6 @@ def predict(image_path):
     class_folder = f"sample_images/{str(top1_label).replace(' ', '_')}"
     reference_image = None
     if os.path.isdir(class_folder):
-        # List all image files in the folder
         image_files = [f for f in os.listdir(class_folder) if f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp"))]
         if image_files:
             chosen_file = random.choice(image_files)
@@ -81,7 +82,6 @@ def predict(image_path):
     else:
         print(f"[DEBUG] Class folder does not exist: {class_folder}")
 
-    # Format Top-5 for gr.Label with class names
     top5_probs = {class_names[int(idx)]: float(score) for idx, score in zip(top5.indices, top5.values)}
     print(f"image path: {image_path}")
     print(f"top1_label: {top1_label}")
@@ -89,6 +89,4 @@ def predict(image_path):
     print(f"[DEBUG] Top-5 labels: {[class_names[int(idx)] for idx in top5.indices]}")
     print(f"[DEBUG] Top-5 probs: {top5_probs}")
 
-    return image, reference_image, top5_probs
-
-
+    return image, reference_image, top5_probs