""" Builds the FAISS index for the visual similarity search. Uses Trendyol DinoV2 (trendyol-dino-v2-ecommerce-256d) to compute 256-dimensional visual embeddings for each product image in artifacts/images/. The resulting index allows: given a query image, finding the products with the most similar visual appearance (color, shape, texture, style). Prerequisite: first run pipeline/7_download_images.py to download the images. Output: - artifacts/faiss_visual/index.faiss (FAISS L2 index, 256d) - artifacts/faiss_visual/asin_map.json (FAISS index -> ASIN map) """ import json import os import sys import time import warnings warnings.filterwarnings("ignore") os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" import numpy as np import torch import faiss from PIL import Image from pathlib import Path from transformers import AutoModel, AutoImageProcessor # --- CONFIGURATION --- PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) IMAGES_DIR = os.path.join(PROJECT_ROOT, "artifacts", "images") METADATA_PATH = os.path.join(PROJECT_ROOT, "artifacts", "metadata_cleaned.json") OUTPUT_DIR = os.path.join(PROJECT_ROOT, "artifacts", "faiss_visual") VISUAL_MODEL_NAME = "Trendyol/trendyol-dino-v2-ecommerce-256d" BATCH_SIZE = 64 # Batch for GPU/CPU inference def load_model(): """Loads the DinoV2 model and the preprocessor.""" print(f"--- Loading model {VISUAL_MODEL_NAME}... ---") processor = AutoImageProcessor.from_pretrained(VISUAL_MODEL_NAME, trust_remote_code=True) model = AutoModel.from_pretrained(VISUAL_MODEL_NAME, trust_remote_code=True) model.eval() # Use MPS (Apple Silicon) if available, otherwise CPU if torch.backends.mps.is_available(): device = torch.device("mps") print(" Device: Apple Silicon (MPS)") elif torch.cuda.is_available(): device = torch.device("cuda") print(" Device: CUDA GPU") else: device = torch.device("cpu") print(" Device: CPU") model = model.to(device) return model, processor, device def compute_embeddings_batch(model, processor, device, image_paths): """Computes embeddings for a batch of images. Returns numpy array (N, 256).""" images = [] valid_indices = [] for i, path in enumerate(image_paths): try: img = Image.open(path).convert("RGB") images.append(img) valid_indices.append(i) except Exception: continue if not images: return np.array([]), valid_indices inputs = processor(images=images, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) # Trendyol DinoV2: output already projected to (batch, 256) and L2-normalized # last_hidden_state is 2D, no need to extract the CLS token embeddings = outputs.last_hidden_state return embeddings.cpu().numpy(), valid_indices def main(): # Check that the images exist if not os.path.exists(IMAGES_DIR): print(f"ERROR: Images folder not found: {IMAGES_DIR}") print(f"Run first: python pipeline/7_download_images.py") return # Load the ASINs of the current catalog (filter to avoid "orphans") if not os.path.exists(METADATA_PATH): print(f"ERROR: metadata not found: {METADATA_PATH}") return with open(METADATA_PATH, "r", encoding="utf-8") as f: products = json.load(f) current_asins = set(p["asin"] for p in products) print(f"ASINs in the current catalog: {len(current_asins):,}") # Collect all the image files and filter only those of the current catalog all_image_files = sorted(Path(IMAGES_DIR).glob("*.*")) all_image_files = [f for f in all_image_files if f.suffix.lower() in (".jpg", ".jpeg", ".png", ".gif")] image_files = [f for f in all_image_files if f.stem in current_asins] skipped = len(all_image_files) - len(image_files) if not image_files: print(f"ERROR: No image (from the current catalog) found in {IMAGES_DIR}") return print(f"Total images in folder: {len(all_image_files):,}") print(f"Images to process: {len(image_files):,}") print(f"Images ignored (orphans): {skipped:,}") # If intermediate embeddings already exist, skip the computation (resume from crash) cached_emb = os.path.join(OUTPUT_DIR, "embeddings.npy") cached_map = os.path.join(OUTPUT_DIR, "asin_map_tmp.json") if os.path.exists(cached_emb) and os.path.exists(cached_map): print(f"\n--- Intermediate embeddings found, loading from cache... ---") embeddings_matrix = np.load(cached_emb) with open(cached_map, "r") as f: asin_map = json.load(f) print(f"Matrix: {embeddings_matrix.shape}, ASIN: {len(asin_map)}") # Jump directly to building the FAISS index _build_and_save_index(embeddings_matrix, asin_map, start_time=time.time()) return # Load model model, processor, device = load_model() # Compute embeddings in batches all_embeddings = [] asin_map = [] # FAISS index -> ASIN start_time = time.time() print(f"--- Computing embeddings (batch_size={BATCH_SIZE})... ---\n") for batch_start in range(0, len(image_files), BATCH_SIZE): batch_end = min(batch_start + BATCH_SIZE, len(image_files)) batch_files = image_files[batch_start:batch_end] batch_paths = [str(f) for f in batch_files] batch_asins = [f.stem for f in batch_files] embeddings, valid_indices = compute_embeddings_batch(model, processor, device, batch_paths) if len(embeddings) > 0: all_embeddings.append(embeddings) for idx in valid_indices: asin_map.append(batch_asins[idx]) # Progress done = min(batch_end, len(image_files)) elapsed = time.time() - start_time rate = done / elapsed if elapsed > 0 else 0 eta = (len(image_files) - done) / rate if rate > 0 else 0 pct = done / len(image_files) * 100 sys.stdout.write( f"\r [{pct:5.1f}%] {done}/{len(image_files)} | " f"{rate:.1f} img/s | ETA: {eta/60:.1f}min" ) sys.stdout.flush() print(f"\n\n--- Embeddings completed: {len(asin_map)} images ---") if not all_embeddings: print("ERROR: No embedding computed.") return # Concatenate all the embeddings and ensure C-contiguous (required by FAISS) embeddings_matrix = np.ascontiguousarray(np.vstack(all_embeddings), dtype="float32") print(f"Embedding matrix: {embeddings_matrix.shape}") # (N, 256) # Save intermediate embeddings to disk (for resuming in case of a crash in the FAISS phase) os.makedirs(OUTPUT_DIR, exist_ok=True) np.save(os.path.join(OUTPUT_DIR, "embeddings.npy"), embeddings_matrix) with open(os.path.join(OUTPUT_DIR, "asin_map_tmp.json"), "w") as f: json.dump(asin_map, f) print(f"Intermediate embeddings saved to {OUTPUT_DIR}") _build_and_save_index(embeddings_matrix, asin_map, start_time) def _build_and_save_index(embeddings_matrix, asin_map, start_time): """Normalizes, builds the FAISS index and saves to disk.""" import faiss as faiss_lib # Ensure C-contiguous embeddings_matrix = np.ascontiguousarray(embeddings_matrix, dtype="float32") # Normalize for cosine similarity (FAISS IndexFlatIP) norms = np.linalg.norm(embeddings_matrix, axis=1, keepdims=True) embeddings_matrix = embeddings_matrix / (norms + 1e-8) print(f"Normalization completed (numpy)") # Build FAISS index (Inner Product = cosine similarity on normalized vectors) dim = embeddings_matrix.shape[1] index = faiss_lib.IndexFlatIP(dim) index.add(embeddings_matrix) print(f"FAISS index: {index.ntotal} vectors, {dim}d") # Save os.makedirs(OUTPUT_DIR, exist_ok=True) faiss_path = os.path.join(OUTPUT_DIR, "index.faiss") faiss_lib.write_index(index, faiss_path) print(f"Index saved: {faiss_path}") map_path = os.path.join(OUTPUT_DIR, "asin_map.json") with open(map_path, "w") as f: json.dump(asin_map, f) print(f"ASIN map saved: {map_path} ({len(asin_map)} products)") # Clean up intermediate files for tmp in ["embeddings.npy", "asin_map_tmp.json"]: tmp_path = os.path.join(OUTPUT_DIR, tmp) if os.path.exists(tmp_path): os.remove(tmp_path) elapsed = time.time() - start_time print(f"\n--- DONE in {elapsed/60:.1f} minutes ---") if __name__ == "__main__": main()