Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import torch
 import clip
 from PIL import Image
+from typing import List, Tuple
 import numpy as np
 
 # Load CLIP model and preprocessing
@@ -9,17 +10,34 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model, preprocess = clip.load("ViT-B/32", device=device)
 
 # Function to predict descriptions and probabilities
-def predict(image, descriptions):
+def predict(image: Image.Image, descriptions: List[str]) -> Tuple[str, float]:
+    """
+    Predict the best matching description for the provided image based on the given descriptions.
+    Uses the CLIP model to compute similarities between the image and text descriptions.
+
+    Args:
+        image (Image.Image): The input image for which the descriptions are being evaluated.
+        descriptions (List[str]): A list of textual descriptions to compare against the image.
+
+    Returns:
+        Tuple[str, float]: A tuple containing the best-matching description and the corresponding probability.
+    """
+    # Preprocess the image and move it to the appropriate device
     image = preprocess(image).unsqueeze(0).to(device)
+    # Tokenize the descriptions and move them to the appropriate device
     text = clip.tokenize(descriptions).to(device)
 
     with torch.no_grad():
+        # Encode image and text features using the CLIP model
         image_features = model.encode_image(image)
         text_features = model.encode_text(text)
 
+        # Compute the similarity scores (logits) between image and text
         logits_per_image, logits_per_text = model(image, text)
+        # Convert logits to probabilities
         probs = logits_per_image.softmax(dim=-1).cpu().numpy()
 
+    # Return the description with the highest probability and the corresponding probability
     return descriptions[np.argmax(probs)], np.max(probs)
 
 # Streamlit app