change to base deit

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 Tesing using DeiT version 1.0",
+    title="Scene Classification with Reference Image Testing using DeiT version 1.0",
     description="Upload an image to get the predicted class with a sample image and top-5 prediction chart."
 )
 

predict.py

@@ -17,7 +17,7 @@ with open("labels.json", "r") as f:
 print("class_names:", class_names)
 
 class DeiT(nn.Module):
-    def __init__(self, model_name="facebook/deit-small-patch16-224", num_classes=40):
+    def __init__(self, model_name="facebook/deit-base-patch16-224", num_classes=40):
         super(DeiT, self).__init__()
         self.model = ViTForImageClassification.from_pretrained(model_name)
         in_features = self.model.classifier.in_features
@@ -30,7 +30,7 @@ class DeiT(nn.Module):
         return outputs.logits
 
 # Load model
-model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_best_model_1.pth")
+model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="deit_base_best_model.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")