Update app.py

app.py

@@ -16,17 +16,17 @@ import spaces
 from typing import List, Dict, Tuple, Optional, Union
 
 # Load model/processor
-model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
+model: CLIPModel = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
+processor: CLIPProcessor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
 model.eval()
 
-DATASET_DIR = Path("dataset")
-CACHE_FILE = "cache.pkl"
+DATASET_DIR: Path = Path("dataset")
+CACHE_FILE: str = "cache.pkl"
 
 # Define supported image formats
-IMAGE_EXTENSIONS = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.gif", "*.webp", "*.tiff", "*.tif"]
+IMAGE_EXTENSIONS: List[str] = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.gif", "*.webp", "*.tiff", "*.tif"]
 
-def get_all_image_files()  -> List[Path]:
+def get_all_image_files() -> List[Path]:
     """
     Get all image files from the dataset directory.
     
@@ -35,7 +35,7 @@ def get_all_image_files()  -> List[Path]:
     Returns:
         List[Path]: List of Path objects for all found image files
     """
-    image_files = []
+    image_files: List[Path] = []
     for ext in IMAGE_EXTENSIONS:
         image_files.extend(DATASET_DIR.glob(ext))
         image_files.extend(DATASET_DIR.glob(ext.upper()))  # Also check uppercase
@@ -59,7 +59,7 @@ def get_embedding(image: Image.Image, device: str = "cpu") -> torch.Tensor:
     inputs = processor(images=image, return_tensors="pt").to(device)
     model_device = model.to(device)
     with torch.no_grad():
-        emb = model_device.get_image_features(**inputs)
+        emb: torch.Tensor = model_device.get_image_features(**inputs)
     # L2 normalize the embeddings
     emb = emb / emb.norm(p=2, dim=-1, keepdim=True)
     return emb
@@ -80,29 +80,29 @@ def get_reference_embeddings() -> Dict[str, torch.Tensor]:
         PermissionError: If unable to write cache file
     """
     # Get all current image files
-    current_image_files = get_all_image_files()
-    current_images = set(img_path.name for img_path in current_image_files)
+    current_image_files: List[Path] = get_all_image_files()
+    current_images: set = set(img_path.name for img_path in current_image_files)
     
     # Load existing cache if it exists
-    cached_embeddings = {}
+    cached_embeddings: Dict[str, torch.Tensor] = {}
     if os.path.exists(CACHE_FILE):
         with open(CACHE_FILE, "rb") as f:
             cached_embeddings = pickle.load(f)
     
     # Check if cache is up to date
-    cached_images = set(cached_embeddings.keys())
+    cached_images: set = set(cached_embeddings.keys())
     
     # If cache is missing images or has extra images, rebuild
     if current_images != cached_images:
         print(f"Cache outdated. Current: {len(current_images)}, Cached: {len(cached_images)}")
-        embeddings = {}
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        embeddings: Dict[str, torch.Tensor] = {}
+        device: str = "cuda" if torch.cuda.is_available() else "cpu"
         
         for img_path in current_image_files:
             print(f"Processing {img_path.name}...")
             try:
-                img = Image.open(img_path).convert("RGB")
-                emb = get_embedding(img, device=device)
+                img: Image.Image = Image.open(img_path).convert("RGB")
+                emb: torch.Tensor = get_embedding(img, device=device)
                 embeddings[img_path.name] = emb.cpu()
             except Exception as e:
                 print(f"Error processing {img_path.name}: {e}")
@@ -117,7 +117,8 @@ def get_reference_embeddings() -> Dict[str, torch.Tensor]:
         print(f"Using cached embeddings for {len(cached_embeddings)} images")
         return cached_embeddings
 
-reference_embeddings = get_reference_embeddings()
+# Initialize reference embeddings
+reference_embeddings: Dict[str, torch.Tensor] = get_reference_embeddings()
 
 @spaces.GPU
 def search_similar(query_img: Image.Image) -> List[Tuple[str, str]]:
@@ -138,21 +139,21 @@ def search_similar(query_img: Image.Image) -> List[Tuple[str, str]]:
     global reference_embeddings
     reference_embeddings = get_reference_embeddings()
     
-    query_emb = get_embedding(query_img, device="cuda")
-    results = []
+    query_emb: torch.Tensor = get_embedding(query_img, device="cuda")
+    results: List[Tuple[str, float]] = []
     
     for name, ref_emb in reference_embeddings.items():
         # Move reference embedding to same device as query
-        ref_emb_gpu = ref_emb.to("cuda")
+        ref_emb_gpu: torch.Tensor = ref_emb.to("cuda")
         # Compute cosine similarity
-        sim = torch.nn.functional.cosine_similarity(query_emb, ref_emb_gpu, dim=1).item()
+        sim: float = torch.nn.functional.cosine_similarity(query_emb, ref_emb_gpu, dim=1).item()
         results.append((name, sim))
     
     results.sort(key=lambda x: x[1], reverse=True)
     
     # Filter out low similarity results (adjust threshold as needed)
-    SIMILARITY_THRESHOLD = 0.2  # Only show results above 20% similarity
-    filtered_results = [(name, score) for name, score in results if score > SIMILARITY_THRESHOLD]
+    SIMILARITY_THRESHOLD: float = 0.2  # Only show results above 20% similarity
+    filtered_results: List[Tuple[str, float]] = [(name, score) for name, score in results if score > SIMILARITY_THRESHOLD]
     
     if not filtered_results:
         return [("No similar images found", "No matches above similarity threshold")]
@@ -181,12 +182,12 @@ def add_image(name: str, image: Image.Image) -> str:
         return "Please provide a valid image name."
     
     # Save as PNG to preserve quality for all input formats
-    path = DATASET_DIR / f"{name}.png"
+    path: Path = DATASET_DIR / f"{name}.png"
     image.save(path, "PNG")
     
     # Use GPU for consistency if available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    emb = get_embedding(image, device=device)
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    emb: torch.Tensor = get_embedding(image, device=device)
     
     # Add to current embeddings and save cache
     reference_embeddings[f"{name}.png"] = emb.cpu()
@@ -196,15 +197,37 @@ def add_image(name: str, image: Image.Image) -> str:
     
     return f"Image '{name}' added to dataset. Total images: {len(reference_embeddings)}"
 
-search_interface = gr.Interface(fn=search_similar,
-                                inputs=gr.Image(type="pil", label="Query Image"),
-                                outputs=gr.Gallery(label="Top Matches", columns=5),
-                                allow_flagging="never")
+# Create Gradio interfaces
+search_interface: gr.Interface = gr.Interface(
+    fn=search_similar,
+    inputs=gr.Image(type="pil", label="Query Image"),
+    outputs=gr.Gallery(label="Top Matches", columns=5),
+    allow_flagging="never",
+    title="Image Similarity Search",
+    description="Upload an image to find similar images in the dataset"
+)
 
-add_interface = gr.Interface(fn=add_image,
-                             inputs=[gr.Text(label="Image Name"), gr.Image(type="pil", label="Product Image")],
-                             outputs="text",
-                             allow_flagging="never")
+add_interface: gr.Interface = gr.Interface(
+    fn=add_image,
+    inputs=[
+        gr.Text(label="Image Name", placeholder="Enter a unique name for your image"),
+        gr.Image(type="pil", label="Product Image")
+    ],
+    outputs="text",
+    allow_flagging="never",
+    title="Add Image to Dataset",
+    description="Add a new image to the searchable dataset"
+)
 
-demo = gr.TabbedInterface([search_interface, add_interface], tab_names=["Search", "Add Product"])
-demo.launch(mcp_server=True)
+# Create main application
+demo: gr.TabbedInterface = gr.TabbedInterface(
+    [search_interface, add_interface], 
+    tab_names=["Search", "Add Product"],
+    title="CLIP Image Search System",
+    theme=gr.themes.Soft()
+)
+
+if __name__ == "__main__":
+    # Ensure dataset directory exists
+    DATASET_DIR.mkdir(exist_ok=True)
+    demo.launch(share=True)