""" Gradio app: Text-to-Image ranking using OpenCLIP (open-source) Features: - Accepts a text query and multiple images (100+). - Encodes text and images with OpenCLIP (ViT-B-32 by default). - Computes cosine similarity, normalizes scores to 0-100. - Returns a ranked CSV and a visual grid image annotated with scores. - GPU optional (will use CUDA if available). """ import os import io import math import time from typing import List, Tuple, Optional import torch import open_clip from PIL import Image, ImageDraw, ImageFont import numpy as np import pandas as pd import gradio as gr # ------------------------- # Configuration / Globals # ------------------------- MODEL_NAME = "ViT-B-32" # MODEL_PRETRAIN = "laion2b_s32b_b79k" MODEL_PRETRAIN = "openai" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" BATCH_SIZE = 64 TOP_K_DEFAULT = 20 THUMB_SIZE = (256, 256) FONT_PATH = None NORMALIZE_SCORE_TO = 100 # ------------------------- _model_data = {"loaded": False} def load_model(device: str = DEVICE): """ Loads OpenCLIP model and transforms. Cached on first call. Returns model, preprocess function, tokenizer, and embedding dimension. """ if _model_data.get("loaded", False): return _model_data["model"], _model_data["preprocess"], _model_data["tokenizer"], _model_data["dim"] print(f"Loading OpenCLIP {MODEL_NAME} ({MODEL_PRETRAIN}) to {device} ...") model, _, preprocess = open_clip.create_model_and_transforms(MODEL_NAME, MODEL_PRETRAIN) tokenizer = open_clip.get_tokenizer(MODEL_NAME) model.to(device) model.eval() dim = model.text_projection.shape[1] if hasattr(model, "text_projection") else model.projection.shape[1] _model_data.update({ "loaded": True, "model": model, "preprocess": preprocess, "tokenizer": tokenizer, "dim": dim }) print("Model loaded.") return model, preprocess, tokenizer, dim # ------------------------- # Utilities # ------------------------- def load_pil_image(file_obj) -> Image.Image: """ Given a file-like object from Gradio (or path), return a PIL image in RGB. """ if isinstance(file_obj, str): img = Image.open(file_obj) else: file_obj.seek(0) img = Image.open(io.BytesIO(file_obj.read())) return img.convert("RGB") def batchify(iterable, batch_size): """Yield successive batches from iterable""" it = list(iterable) for i in range(0, len(it), batch_size): yield it[i:i + batch_size] def encode_text(text: str, model, tokenizer, device: str = DEVICE) -> torch.Tensor: """ Encode text to a normalized embedding tensor (1 x dim) """ texts_tokenized = tokenizer([text]) with torch.no_grad(): text_tokens = texts_tokenized.to(device) text_feats = model.encode_text(text_tokens) # (1, dim) text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True) return text_feats def encode_images(images: List[Image.Image], model, preprocess, device: str = DEVICE, batch_size: int = BATCH_SIZE) -> torch.Tensor: """ Encode a list of PIL images into normalized embeddings (N x dim). Uses batching to avoid memory blowups. Returns CPU tensor. """ all_feats = [] model_device = next(model.parameters()).device for batch in batchify(images, batch_size): batch_tensors = torch.stack([preprocess(img) for img in batch]).to(device) with torch.no_grad(): feats = model.encode_image(batch_tensors) feats = feats / feats.norm(dim=-1, keepdim=True) all_feats.append(feats.cpu()) all_feats = torch.cat(all_feats, dim=0) return all_feats def cosine_similarity_matrix(text_feat: torch.Tensor, image_feats: torch.Tensor) -> np.ndarray: """ Given text_feat (1 x dim) and image_feats (N x dim), compute cosine similarities in numpy. Returns ndarray shape (N,) """ if isinstance(text_feat, torch.Tensor): text_feat = text_feat.cpu() sims = (image_feats @ text_feat.squeeze(0).cpu().T).numpy().squeeze() sims = np.clip(sims, -1.0, 1.0) return sims def normalize_scores_to_range(scores: np.ndarray, low=0.0, high=NORMALIZE_SCORE_TO) -> np.ndarray: """ Maps scores from [-1,1] (cosine) to [low,high] (e.g., 0..100). If all scores equal, map to mid-range to avoid divide-by-zero. """ min_s, max_s = float(scores.min()), float(scores.max()) if math.isclose(min_s, max_s): mid = (low + high) / 2.0 return np.full_like(scores, fill_value=mid, dtype=float) scores_clipped = np.clip(scores, -1.0, 1.0) norm01 = (scores_clipped - (-1.0)) / (2.0) mapped = low + norm01 * (high - low) return mapped def make_visual_grid(images: List[Image.Image], scores: List[float], top_k: int = 12, thumb_size: Tuple[int, int] = THUMB_SIZE, columns: int = 4, font_path: Optional[str] = FONT_PATH) -> Image.Image: """ Create a single PIL image that arranges top_k thumbnails in a grid with score captions. """ top_k = min(top_k, len(images)) rows = math.ceil(top_k / columns) w, h = thumb_size caption_height = 28 grid_w = columns * w grid_h = rows * (h + caption_height) grid_img = Image.new("RGB", (grid_w, grid_h), color=(255, 255, 255)) draw = ImageDraw.Draw(grid_img) try: if font_path and os.path.exists(font_path): font = ImageFont.truetype(font_path, 16) else: font = ImageFont.load_default() except Exception: font = ImageFont.load_default() for idx in range(top_k): img = images[idx].copy().resize(thumb_size, Image.Resampling.LANCZOS) col = idx % columns row = idx // columns x = col * w y = row * (h + caption_height) grid_img.paste(img, (x, y)) caption = f"{scores[idx]:.1f}" bbox = draw.textbbox((0, 0), caption, font=font) text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1] rect_x0 = x rect_y0 = y + h rect_x1 = x + w rect_y1 = rect_y0 + caption_height draw.rectangle([rect_x0, rect_y0, rect_x1, rect_y1], fill=(255, 255, 255)) text_x = x + 6 text_y = rect_y0 + (caption_height - text_h) // 2 draw.text((text_x, text_y), caption, fill=(0, 0, 0), font=font) return grid_img def rank_images_by_text(query: str, files: List[gr.File], top_k: int = TOP_K_DEFAULT, use_gpu: bool = (DEVICE == "cuda")) -> Tuple[pd.DataFrame, Image.Image]: """ Main pipeline: - load model (if not) - read images from files - encode text and images - compute cosine similarity - produce ranked DataFrame and visual grid image Returns: (pandas.DataFrame with columns ['filename','score_cosine','score_normalized'], PIL.Image grid) """ start_time = time.time() if not query or (not files): raise ValueError("Please provide both a text query and at least one image file.") model, preprocess, tokenizer, dim = load_model(DEVICE if use_gpu else "cpu") device = DEVICE if use_gpu else "cpu" images = [] filenames = [] for f in files: try: pil = load_pil_image(f) images.append(pil) name = getattr(f, "name", None) if name: fname = os.path.basename(name) else: fname = getattr(f, "filename", "uploaded_image") filenames.append(fname) except Exception as e: print(f"Skipping a file due to load error: {e}") if len(images) == 0: raise ValueError("No valid images could be loaded from uploads.") text_feat = encode_text(query, model, tokenizer, device=device) image_feats = encode_images(images, model, preprocess, device=device, batch_size=BATCH_SIZE) sims = cosine_similarity_matrix(text_feat, image_feats) # range [-1,1] scores_norm = normalize_scores_to_range(sims, low=0.0, high=float(NORMALIZE_SCORE_TO)) # Rank results order = np.argsort(-sims) sims_sorted = sims[order] scores_sorted = scores_norm[order] filenames_sorted = [filenames[i] for i in order] images_sorted = [images[i] for i in order] df = pd.DataFrame({ "filename": filenames_sorted, "score_cosine": sims_sorted, f"score_{int(NORMALIZE_SCORE_TO)}": scores_sorted }) top_k = min(top_k, len(images_sorted)) top_images = images_sorted[:top_k] top_scores = scores_sorted[:top_k].tolist() grid_img = make_visual_grid(top_images, top_scores, top_k=top_k, thumb_size=THUMB_SIZE, columns=4) elapsed = time.time() - start_time print(f"Query processed in {elapsed:.2f}s. Images: {len(images)}. Top-K: {top_k}") return df, grid_img # ------------------------- # Gradio app UI # ------------------------- def gradio_rank_fn(query: str, image_files: List[gr.File], top_k: int = TOP_K_DEFAULT, use_gpu: bool = (DEVICE == "cuda")): """ Wrapper for Gradio. Returns (ranked table as CSV string / DataFrame, grid image as PIL, optionally downloadable CSV). """ if not image_files: return "No images uploaded.", None, None try: df, grid_img = rank_images_by_text(query, image_files, top_k=top_k, use_gpu=use_gpu) except Exception as e: return f"Error: {e}", None, None csv_buffer = io.StringIO() df.to_csv(csv_buffer, index=False) csv_bytes = csv_buffer.getvalue().encode("utf-8") csv_buffer.close() summary = f"Ranked {len(df)} images for query: '{query}'. Top score: {df['score_cosine'].max():.4f}" return summary, grid_img, ("rankings.csv", csv_bytes, "text/csv") def build_interface(): title = "Text → Image Ranking" description = """ Enter any text query (e.g., "red chinos") and upload multiple product images (100+ supported). The app uses an OpenCLIP model (open-source) to compute embeddings for text and images, then ranks images by cosine similarity. """ with gr.Blocks(title=title) as demo: gr.Markdown(f"# {title}") gr.Markdown(description) with gr.Row(): with gr.Column(scale=3): query = gr.Textbox(label="Text query", placeholder="e.g. 'red chinos' or 'floral kurta with pockets'", lines=1) image_files = gr.File(label="Upload product images (multiple)", file_count="multiple", file_types=["image"], interactive=True) top_k = gr.Slider(minimum=1, maximum=64, value=TOP_K_DEFAULT, step=1, label="Top-K to visualize") use_gpu = gr.Checkbox(label=f"Use GPU (detected device: {DEVICE}). Uncheck to force CPU.", value=(DEVICE == "cuda")) run_btn = gr.Button("Rank images") status_output = gr.Textbox(label="Status", interactive=False) with gr.Column(scale=2): gallery = gr.Image(type="pil", label="Top results grid (annotated)") download = gr.File(label="Download CSV rankings") summary = gr.Textbox(label="Summary", interactive=False) def wrapped_run(q, files, topk, use_gpu_flag): status = "Processing..." try: summary_text, grid_img, csv_tuple = gradio_rank_fn(q, files, topk, use_gpu_flag) if csv_tuple: fname, content_bytes, mime = csv_tuple tmp_path = os.path.join(os.getcwd(), fname) with open(tmp_path, "wb") as f: f.write(content_bytes) csv_path = tmp_path else: csv_path = None return summary_text, grid_img, csv_path except Exception as e: return f"Error: {e}", None, None run_btn.click(fn=wrapped_run, inputs=[query, image_files, top_k, use_gpu], outputs=[summary, gallery, download]) gr.Markdown("## Notes") gr.Markdown("- This uses an **open-source** OpenCLIP model. No paid API calls.") gr.Markdown("- The app is slow because every time it runs it creates embeddings of the text and the images . The speed of the app can be increased if we use already stored images so we don't have to create embeddings everytime.") gr.Markdown("The accuracy of this app can be increased if we used different models of open clip , but for computational efficiency i have utilized one of the efficient models . Also if we finetune this model , the accuracy of the model can be hugely increased, But since this is just a asssignment , i have created a demo prototype only.") return demo if __name__ == "__main__": demo = build_interface() # Start Gradio demo.launch()