Update app.py

app.py

@@ -30,8 +30,16 @@ class NewheadDinov2ForImageClassification(Dinov2ForImageClassification):
 
         # Classifier head
         self.classifier = create_head(config.hidden_size * 2, config.num_labels)
-        
-# IMPORT CLASSIFICATION MODEL
+
+model = NewheadDinov2ForImageClassification.from_pretrained(checkpoint_name)
+
+# IMPORT MODEL CONFIG PARAMETERS
+hf_hub_download(repo_id=checkpoint_name, filename="config.json")
+id2label = config["id2label"]
+label2id = config["label2id"]
+image_size = config["image_size"]
+classes_names = list(label2id.keys())
+'''
 # import labels
 classes_names = ["Acropore_branched", "Acropore_digitised", "Acropore_tabular", "Algae_assembly", 
                "Algae_limestone", "Algae_sodding", "Dead_coral", "Fish", "Human_object",
@@ -43,8 +51,7 @@ classes_names = ["Acropore_branched", "Acropore_digitised", "Acropore_tabular",
 classes_nb = list(np.arange(len(classes_names)))
 id2label = {int(classes_nb[i]): classes_names[i] for i in range(len(classes_nb))}
 label2id = {v: k for k, v in id2label.items()}
-
-model = NewheadDinov2ForImageClassification.from_pretrained(checkpoint_name)
+'''
 
 def sigmoid(_outputs):
     return 1.0 / (1.0 + np.exp(-_outputs))