Update app.py

app.py

@@ -3,60 +3,43 @@ from tensorflow import keras
 import numpy as np
 import gradio as gr
 
-model = keras.models.load_model("Model.keras")
+# Load PyTorch 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 = []
+    for in_f ,out_f in zip(features_lst[:-1] , features_lst[1:]):
+        layers.append(nn.Linear(in_f , out_f))
+        layers.append(activation_func())
+        layers.append(nn.BatchNorm1d(out_f))
+        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
 
-classnames = [
-    "Acacia",
-    "Adenanthera microsperma",
-    "Adenium species",
-    "Anacardium occidentale",
-    "Annona squamosa",
-    "Artocarpus altilis",
-    "Artocarpus heterophyllus",
-    "Barringtonia acutangula",
-    "Cananga odorata",
-    "Carica papaya",
-    "Casuarina equisetifolia",
-    "Cedrus",
-    "Chrysophyllum cainino",
-    "Citrus aurantiifolia",
-    "Citrus grandis",
-    "Cocos nucifera",
-    "Dalbergia oliveri",
-    "Delonix regia",
-    "Dipterocarpus alatus",
-    "Erythrina fusca",
-    "Eucalyptus",
-    "Ficus microcarpa",
-    "Ficus racemosa",
-    "Gmelina arborea Roxb",
-    "Hevea brasiliensis",
-    "Hopea",
-    "Khaya senegalensis",
-    "Khaya senegalensis A.Juss",
-    "Lagerstroemia speciosa",
-    "Magnolia alba",
-    "Mangifera",
-    "Melaleuca",
-    "Melia azedarach",
-    "Musa",
-    "Nephelium lappaceum",
-    "Persea",
-    "Polyalthia longifolia",
-    "Prunnus",
-    "Prunus salicina",
-    "Psidium guajava",
-    "Pterocarpus macrocarpus",
-    "Senna siamea",
-    "Spondias mombin L",
-    "Syzygium nervosum",
-    "Tamarindus indica",
-    "Tectona grandis",
-    "Terminalia catappa",
-    "Veitchia merrilli",
-    "Wrightia",
-    "Wrightia religiosa",
-]
+class NewheadDinov2ForImageClassification(Dinov2ForImageClassification):
+    def __init__(self, config: Dinov2Config) -> None:
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.dinov2 = Dinov2Model(config)
+
+        # Classifier head
+        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",
+               "Living_coral", "Millepore", "No_acropore_encrusting", "No_acropore_massive",
+              "No_acropore_sub_massive",  "Rock", "Sand",
+               "Scrap", "Sea_cucumber", "Syringodium_isoetifolium",
+               "Thalassodendron_ciliatum",  "Useless"]
+
+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 predict(path):
@@ -65,7 +48,7 @@ def predict(path):
     image = np.expand_dims(image, axis=0)
     pred = model.predict(image, verbose=0)
     pred = pred[0]
-    confidences = {classnames[i]: round(float(pred[i]), 2) for i in range(50)}
+    confidences = {classes_names[i]: round(float(pred[i]), 2) for i in range(50)}
     return confidences