switch to deit

predict.py

@@ -4,7 +4,7 @@ from PIL import Image
 import json
 import numpy as np
 # from model import load_model
-from transformers import AutoImageProcessor, SwinForImageClassification
+from transformers import AutoImageProcessor, SwinForImageClassification, ViTForImageClassification
 import torch.nn as nn
 import os
 import pandas as pd
@@ -16,30 +16,40 @@ with open("labels.json", "r") as f:
     class_names = json.load(f)
 print("class_names:", class_names)
 
-# Load model
-model = SwinForImageClassification.from_pretrained("microsoft/swin-base-patch4-window7-224", num_labels=1000)
+class DeiT(nn.Module):
+    def __init__(self, model_name="facebook/deit-small-patch16-224", num_classes=None):
+        super(DeiT, self).__init__()
+        self.model = ViTForImageClassification.from_pretrained(model_name)
+        in_features = self.model.classifier.in_features
+        self.model.classifier = nn.Sequential(
+            nn.Linear(in_features, num_classes)
+        )
+
+    def forward(self, images):
+        outputs = self.model(images)
+        return outputs.logits
 
-# Download the model file from the Hub if not present
+# Load model
 model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="best_model.pth")
 print("Model path:", model_path)
-state_dict = torch.load(model_path, map_location="cpu")
-model.load_state_dict(state_dict)
+model = DeiT(num_classes=len(class_names))
+checkpoint = torch.load(model_path, map_location="cpu")
+model.load_state_dict(checkpoint["model_state_dict"])
 model.eval()
 
-
-# Image transform
-# transform = transforms.Compose([
-#     transforms.Resize((224, 224)),
-#     transforms.ToTensor(),
-#     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
-# ])
-#Swin
+#deit transform
 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([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
 ])
+#Swin
+# 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):