Update app.py

app.py

@@ -1,48 +1,37 @@
-from keras.models import load_model
-from PIL import Image, ImageOps
-import numpy as np
+import os
+import glob
+import json
+import warnings
+
+warnings.filterwarnings("ignore")
+
+import torch
+from PIL import Image
 import gradio as gr
 
-# Load the model
-classify_model = load_model('keras_model.h5')
+import models
+
+with open("../prepare_data/index_to_species.json", "r") as file:
+    index_to_species_data = file.read()
+index_to_species = json.loads(index_to_species_data)
 
-def format_label(label):
-  """
-  From '0 class 1\n' to 'class 1'
-  """
-  return label[:-1]
+num_classes = len(list(index_to_species.keys()))
+
+# Load the model
+classify_model = models.DinoVisionTransformerClassifier(num_classes)
+classify_model.load_state_dict(torch.load("best_dinov2_both_2023-11-21_07-44-35.pth"))
+classify_model.eval()
 
 
 def classify(image):
-  # Create the array of the right shape to feed into the keras model
-  # The 'length' or number of images you can put into the array is
-  # determined by the first position in the shape tuple, in this case 1.
-  data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
-
-  #resize the image to a 224x224 with the same strategy as in TM2:
-  #resizing the image to be at least 224x224 and then cropping from the center
-  size = (224, 224)
-  image = ImageOps.fit(image, size, Image.LANCZOS)
-
-  #turn the image into a numpy array
-  image_array = np.asarray(image)
-  # Normalize the image
-  normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1
-  # Load the image into the array
-  data[0] = normalized_image_array
-
-  # run the inference
-  pred = classify_model.predict(data)
-  pred = pred.tolist()
-
-  with open('labels.txt','r') as f:
-    labels = f.readlines()
-
-  result = {format_label(labels[i]): round(pred[0][i],2) for i in range(len(pred[0]))}
-  sorted_result = {k: v for k, v in sorted(result.items(), key=lambda item: item[1], reverse=True) if v > 0}
-  
-
-  return sorted_result
+    output = classify_model(image)[0]
+    tops = torch.topk(output, k=k).indices
+    scores = torch.softmax(output, dim=0)[tops]
+    
+    result = {index_to_species[str(tops[i].item())]: round(scores[i].item(), 2) for i in range(len(tops))}
+    sorted_result = {k: v for k, v in sorted(result.items(), key=lambda item: item[1], reverse=True) if v > 0}
+    
+    return sorted_result
 
     
 title = "🐢"