Update app.py

app.py

@@ -5,7 +5,9 @@ from transformers import Dinov2Config, Dinov2Model, Dinov2ForImageClassification
 import torch.nn as nn
 import os
 
-# Load PyTorch model
+model_name = "dinov2-large-2024_01_24-with_data_aug_batch-size32_epochs93_freeze"
+
+# CREATE CUSTOM MODEL
 def create_head(num_features , number_classes ,dropout_prob=0.5 ,activation_func =nn.ReLU):
     features_lst = [num_features , num_features//2 , num_features//4]
     layers = []
@@ -16,6 +18,7 @@ def create_head(num_features , number_classes ,dropout_prob=0.5 ,activation_func
         if dropout_prob !=0 : layers.append(nn.Dropout(dropout_prob))
     layers.append(nn.Linear(features_lst[-1] , number_classes))
     return nn.Sequential(*layers)
+from transformers import Dinov2Config, Dinov2Model
 
 class NewheadDinov2ForImageClassification(Dinov2ForImageClassification):
     def __init__(self, config: Dinov2Config) -> None:
@@ -28,7 +31,6 @@ class NewheadDinov2ForImageClassification(Dinov2ForImageClassification):
         self.classifier = create_head(config.hidden_size * 2, config.num_labels)
         
 # IMPORT CLASSIFICATION MODEL
-checkpoint_name = "lombardata/dino-base-2023_11_27-with_custom_head"
 # import labels
 classes_names = ["Acropore_branched", "Acropore_digitised", "Acropore_tabular", "Algae_assembly", 
                "Algae_limestone", "Algae_sodding", "Dead_coral", "Fish", "Human_object",