Update app.py

app.py

@@ -15,27 +15,28 @@ import models
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 # print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
-with open("index_to_species.json", "r") as file:
-    index_to_species_data = file.read()
-index_to_species = json.loads(index_to_species_data)
+# with open("index_to_species.json", "r") as file:
+#     index_to_species_data = file.read()
+# index_to_species = json.loads(index_to_species_data)
 
 num_classes = len(list(index_to_species.keys()))
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the model
-classify_model = models.DinoVisionTransformerClassifier(num_classes)
-classify_model = classify_model.to(device)
-classify_model.load_state_dict(torch.load("best_dinov2_both_2023-11-21_07-44-35.pth", map_location=torch.device(device)))
-classify_model.eval()
+linear_model_name = 'linear_2025-07-08.pt'
+classify_model = models.LinearClassifier(input_dim=768, output_dim=num_classes)
+classify_model.load_state_dict(torch.load(os.path.join('models', linear_model_name)))
 
 k = 5
 
 def classify(image):
-    output = classify_model(image)[0]
+    embedding = extract_embedding(image)
+    embedding = embedding['embedding']
+    output = classify_model(torch.Tensor(embedding).to(device))
     tops = torch.topk(output, k=k).indices
     scores = torch.softmax(output, dim=0)[tops]
-    
+
     result = {index_to_species[str(tops[i].item())].replace("_", " "): 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}
     # Get the current time
@@ -56,4 +57,5 @@ gr.Interface(
     inputs=gr.Image(type="pil", label="Input Image"),
     outputs=[gr.JSON()],
     title=title,
+    debug=True
 ).launch()