Upload indexer.py

indexer.py

@@ -0,0 +1,53 @@
+import torch
+import clip
+from torch.utils.data import DataLoader
+from dataset import FlickrStreamer
+from torchvision import transforms
+from tqdm import tqdm
+
+# --- SETTINGS ---
+INDEX_LIMIT = 5000  # How many images to make searchable
+BATCH_SIZE = 32
+SAVE_PATH = "flickr_embeddings.pt"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# ----------------
+
+def build_index():
+    print(f"Loading CLIP on {DEVICE}...")
+    model, preprocess = clip.load("ViT-B/32", device=DEVICE)
+    
+    # We need the images resized for CLIP (224x224)
+    tf = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor()
+    ])
+    
+    # Load Dataset (Stream mode)
+    print(f"Connecting to Deep Lake (Limit: {INDEX_LIMIT})...")
+    dataset = FlickrStreamer(limit=INDEX_LIMIT, transform=tf)
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
+    
+    all_embeddings = []
+    
+    print("Indexing images (calculating fingerprints)...")
+    with torch.no_grad():
+        for images in tqdm(loader):
+            images = images.to(DEVICE)
+            
+            # Calculate features
+            features = model.encode_image(images)
+            
+            # Normalize features (Crucial for Cosine Similarity later)
+            features /= features.norm(dim=-1, keepdim=True)
+            
+            all_embeddings.append(features.cpu())
+            
+    # Concatenate all batches into one big list
+    final_index = torch.cat(all_embeddings)
+    
+    # Save to file
+    torch.save(final_index, SAVE_PATH)
+    print(f"\nSuccess! Saved {len(final_index)} image fingerprints to '{SAVE_PATH}'")
+
+if __name__ == "__main__":
+    build_index()