Update app.py

app.py

@@ -15,6 +15,44 @@ import models
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 # print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# DINOv2
+
+# Select checkpoint
+dinov2_ckpt = ['dinov2_vits14', 'dinov2_vitb14', 'dinov2_vitl14', 'dinov2_vitg14'][1]
+dinov2 = torch.hub.load('facebookresearch/dinov2', dinov2_ckpt)
+
+dinov2.to(device)
+print()
+
+transform_image = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406],
+                        std=[0.229, 0.224, 0.225])
+])
+
+def extract_embedding(image):
+    """
+    Predict the identity of an image.
+    Args:
+    image: A PIL Image object.
+    Returns:
+    A string representing the predicted identity of the image.
+    """
+    
+    # Convert the image to a tensor.
+    transformed_img = transform_image(image)[:3].unsqueeze(0).to(device)
+    
+    # Get the embedding of the image.
+    with torch.no_grad():
+    embedding = dinov2(transformed_img)
+    # print(embedding.shape)
+    embedding = embedding[0].cpu().numpy().tolist()
+    return {
+      "embedding": embedding
+    }
+    
 with open("index_to_species.json", "r") as file:
     index_to_species_data = file.read()
 index_to_species = json.loads(index_to_species_data)