roll back to deit

predict.py

@@ -3,7 +3,8 @@ from torchvision import transforms
 from PIL import Image
 import json
 import numpy as np
-from transformers import AutoImageProcessor, SwinForImageClassification, ViTForImageClassification, ViTImageProcessor
+# from model import load_model
+from transformers import AutoImageProcessor, SwinForImageClassification, ViTForImageClassification
 import torch.nn as nn
 import os
 import pandas as pd
@@ -15,29 +16,31 @@ with open("labels.json", "r") as f:
     class_names = json.load(f)
 print("class_names:", class_names)
 
-model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="vit_best_model.pth")
-
-processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
-
-class ViTCustom(nn.Module):
-    def __init__(self, model_name="google/vit-base-patch16-224", num_classes=40):
-        super(ViTCustom, self).__init__()
+class DeiT(nn.Module):
+    def __init__(self, model_name="facebook/deit-small-patch16-224", num_classes=40):
+        super(DeiT, self).__init__()
         self.model = ViTForImageClassification.from_pretrained(model_name)
         in_features = self.model.classifier.in_features
-        self.model.classifier = nn.Linear(in_features, num_classes)
+        self.model.classifier = nn.Sequential(
+            nn.Linear(in_features, num_classes)
+        )
 
     def forward(self, images):
         outputs = self.model(images)
         return outputs.logits
 
 # Load model
-model = ViTCustom(model_name="google/vit-base-patch16-224", num_classes=40)
+model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_best_model_1.pth")
+print("Model path:", model_path)
+model = DeiT(model_name="facebook/deit-tiny-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()
 
+
+#Swin
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -45,7 +48,9 @@ transform = transforms.Compose([
                          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)
 
@@ -64,6 +69,7 @@ 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)
@@ -75,6 +81,7 @@ 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}")
@@ -85,4 +92,3 @@ def predict(image_path):
     return image, reference_image, top5_probs
 
 
-