Update app.py

app.py

@@ -6,10 +6,11 @@ import torch.nn as nn
 import os
 from huggingface_hub import hf_hub_download
 
+# DEFINE MODEL NAME
 model_name = "dinov2-large-2024_01_24-with_data_aug_batch-size32_epochs93_freeze"
 checkpoint_name = "lombardata/" + model_name
 
-# CREATE CUSTOM MODEL
+# IMPORT CLASSIFICATION 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 = []
@@ -35,7 +36,12 @@ class NewheadDinov2ForImageClassification(Dinov2ForImageClassification):
 model = NewheadDinov2ForImageClassification.from_pretrained(checkpoint_name)
 
 # IMPORT MODEL CONFIG PARAMETERS
-hf_hub_download(repo_id=checkpoint_name, filename="config.json")
+config_path = hf_hub_download(repo_id=checkpoint_name, filename="config.json")
+# Opening JSON file
+config_file = open(config_path)
+# returns JSON object as 
+config = json.load(config_file)
+# import parameters
 id2label = config["id2label"]
 label2id = config["label2id"]
 image_size = config["image_size"]