Update app.py

app.py

@@ -0,0 +1,367 @@
+"""
+VisualRAG — Multi-Modal AI System
+==================================
+Stack : YOLOv8n · CLIP ViT-B/32 · FAISS · Zephyr-7B · Gradio
+Deploy: HuggingFace Spaces (CPU Basic — free tier)
+
+Pipeline
+  Index  : Image → YOLOv8 detection → CLIP embedding → FAISS vector store
+  Query  : Text → CLIP text embedding → cosine k-NN → LLM answer generation
+"""
+
+import json
+import os
+from datetime import datetime
+
+import faiss
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from transformers import CLIPModel, CLIPProcessor
+from ultralytics import YOLO
+from huggingface_hub import InferenceClient
+
+# ── Model loading (runs once at Space start-up) ────────────────────────────────
+print("⏳ Loading CLIP ViT-B/32 ...")
+CLIP_MODEL_ID = "openai/clip-vit-base-patch32"
+clip_model     = CLIPModel.from_pretrained(CLIP_MODEL_ID)
+clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID)
+clip_model.eval()
+
+print("⏳ Loading YOLOv8n ...")
+yolo = YOLO("yolov8n.pt")       # auto-downloads ~6 MB on first run
+
+print("⏳ Initialising LLM inference client ...")
+# Uses HF Serverless Inference — free for open models, no GPU in Space needed.
+# Add HF_TOKEN as a Space secret for higher rate limits.
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+llm = InferenceClient(
+    model="HuggingFaceH4/zephyr-7b-beta",
+    token=HF_TOKEN,
+)
+
+print("✅ All models ready.")
+
+# ── FAISS vector store (in-memory, session-scoped) ─────────────────────────────
+EMBED_DIM = 512                          # CLIP ViT-B/32 output dimension
+faiss_index = faiss.IndexFlatIP(EMBED_DIM)   # inner-product → cosine after L2 norm
+image_store: list[dict] = []            # parallel list to FAISS index rows
+
+
+# ── Embedding helpers ──────────────────────────────────────────────────────────
+
+def embed_image(pil_img: Image.Image) -> np.ndarray:
+    """Return L2-normalised 512-d CLIP image embedding."""
+    inputs = clip_processor(images=pil_img, return_tensors="pt")
+    with torch.no_grad():
+        features = clip_model.get_image_features(**inputs)
+    emb = features.numpy().astype("float32")
+    faiss.normalize_L2(emb)
+    return emb
+
+
+def embed_text(text: str) -> np.ndarray:
+    """Return L2-normalised 512-d CLIP text embedding."""
+    inputs = clip_processor(text=[text], return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        features = clip_model.get_text_features(**inputs)
+    emb = features.numpy().astype("float32")
+    faiss.normalize_L2(emb)
+    return emb
+
+
+# ── Computer-vision detection pipeline ────────────────────────────────────────
+
+def run_detection(pil_img: Image.Image):
+    """
+    Run YOLOv8 on a PIL image.
+    Returns: (annotated_pil, detections_list, summary_str)
+    """
+    results = yolo(np.array(pil_img))[0]
+    annotated = Image.fromarray(results.plot())
+
+    detections = []
+    if results.boxes is not None:
+        for box in results.boxes:
+            cls_id = int(box.cls[0])
+            conf   = float(box.conf[0])
+            detections.append({
+                "label":      yolo.names[cls_id],
+                "confidence": round(conf, 3),
+            })
+
+    # Object count summary, e.g. "2 person, 1 car"
+    counts: dict[str, int] = {}
+    for d in detections:
+        counts[d["label"]] = counts.get(d["label"], 0) + 1
+    summary = ", ".join(f"{v} {k}" for k, v in counts.items()) or "no objects detected"
+
+    return annotated, detections, summary
+
+
+# ── Index pipeline ─────────────────────────────────────────────────────────────
+
+def index_image(image_np, note: str):
+    """
+    Detect objects, compute CLIP embedding, add to FAISS index.
+    Triggered by the 'Detect & Index' button.
+    """
+    if image_np is None:
+        return None, "❌ Please upload an image first.", _store_badge()
+
+    pil_img                     = Image.fromarray(image_np)
+    annotated, detections, summary = run_detection(pil_img)
+    embedding                   = embed_image(pil_img)
+
+    # Persist to vector store
+    faiss_index.add(embedding)
+    record = {
+        "id":        len(image_store),
+        "image":     pil_img.copy(),
+        "annotated": annotated,
+        "detections": detections,
+        "summary":   summary,
+        "note":      note.strip() or "—",
+        "ts":        datetime.now().strftime("%H:%M:%S"),
+    }
+    image_store.append(record)
+
+    msg = f"✅ Image #{record['id']} indexed · Found: {summary}"
+    return annotated, msg, _store_badge()
+
+
+def _store_badge() -> str:
+    return f"���� {len(image_store)} image(s) in vector store"
+
+
+# ── RAG query pipeline ─────────────────────────────────────────────────────────
+
+def query_images(text_query: str, top_k: int):
+    """
+    1. Embed query text with CLIP.
+    2. Cosine k-NN search in FAISS.
+    3. Build RAG prompt from retrieved context.
+    4. Generate answer with Zephyr-7B via HF Inference API.
+    """
+    if not image_store:
+        return None, "❌ No images indexed yet — go to 'Detect & Index' first.", []
+
+    query_emb   = embed_text(text_query)
+    k           = min(int(top_k), len(image_store))
+    scores, idxs = faiss_index.search(query_emb, k)
+
+    # Build retrieval results & context string for the LLM
+    hits      = []
+    ctx_lines = []
+    for rank, (score, idx) in enumerate(zip(scores[0], idxs[0])):
+        if idx < 0:
+            continue
+        item = image_store[int(idx)]
+        hits.append({
+            "rank":    rank + 1,
+            "img_id":  int(idx),
+            "score":   round(float(score), 4),
+            "objects": item["summary"],
+            "note":    item["note"],
+        })
+        ctx_lines.append(
+            f"[Image #{idx}] objects: {item['summary']} | "
+            f"note: {item['note']} | indexed at: {item['ts']} | "
+            f"cosine similarity: {score:.3f}"
+        )
+
+    context = "\n".join(ctx_lines)
+
+    # Zephyr-7B chat prompt (ChatML format)
+    prompt = (
+        "<|system|>\n"
+        "You are a concise visual-AI assistant. "
+        "Answer the user's question using only the retrieved image context below. "
+        "If the context is insufficient, say so.\n"
+        "<|user|>\n"
+        f"Retrieved context:\n{context}\n\n"
+        f"Question: {text_query}\n"
+        "<|assistant|>\n"
+    )
+
+    try:
+        answer = llm.text_generation(
+            prompt,
+            max_new_tokens=300,
+            temperature=0.2,
+            repetition_penalty=1.1,
+            stop_sequences=["<|user|>", "<|system|>"],
+        ).strip()
+    except Exception as exc:
+        answer = (
+            f"⚠️ LLM call failed ({exc}).\n\n"
+            f"Raw retrieval context:\n{context}"
+        )
+
+    best_idx   = int(idxs[0][0]) if len(idxs[0]) > 0 and idxs[0][0] >= 0 else None
+    best_image = image_store[best_idx]["annotated"] if best_idx is not None else None
+
+    return best_image, answer, hits
+
+
+# ── Gradio UI ──────────────────────────────────────────────────────────────────
+
+CUSTOM_CSS = """
+.hero { text-align:center; padding:20px 0 4px; }
+.hero h1 { font-size:28px; margin:0; }
+.hero p  { color:var(--color-subdued); margin:6px 0 0; }
+.badge-row { display:flex; gap:8px; justify-content:center; flex-wrap:wrap; margin-top:10px; }
+.badge {
+  background:var(--color-background-secondary);
+  border:1px solid var(--border-color-primary);
+  border-radius:20px;
+  padding:3px 12px;
+  font-size:12px;
+  color:var(--color-text-body);
+}
+"""
+
+with gr.Blocks(
+    title="VisualRAG – Multi-Modal AI System",
+    theme=gr.themes.Soft(primary_hue="blue"),
+    css=CUSTOM_CSS,
+) as demo:
+
+    # ── Header ─────────────────────────────────────────────────────────────────
+    gr.HTML("""
+    <div class="hero">
+      <h1>🔍 VisualRAG</h1>
+      <p>Multi-Modal AI · Object Detection + Visual Embeddings + RAG + LLM</p>
+      <div class="badge-row">
+        <span class="badge">YOLOv8</span>
+        <span class="badge">CLIP ViT-B/32</span>
+        <span class="badge">FAISS</span>
+        <span class="badge">Zephyr-7B</span>
+        <span class="badge">Gradio</span>
+      </div>
+    </div>
+    """)
+
+    with gr.Tabs():
+
+        # ── TAB 1 : Detect & Index ─────────────────────────────────────────────
+        with gr.Tab("📤 Detect & Index"):
+            gr.Markdown(
+                "Upload any image. YOLOv8 will detect objects, then CLIP encodes it "
+                "into a 512-d embedding stored in FAISS for later retrieval."
+            )
+            with gr.Row():
+                with gr.Column(scale=1):
+                    img_input = gr.Image(
+                        label="Upload image",
+                        type="numpy",
+                        sources=["upload", "webcam"],
+                    )
+                    note_input = gr.Textbox(
+                        label="Context note (optional)",
+                        placeholder="e.g. 'Warehouse entrance camera, aisle 3'",
+                    )
+                    index_btn = gr.Button(
+                        "🔍 Detect & Index", variant="primary", size="lg"
+                    )
+
+                with gr.Column(scale=1):
+                    detection_output = gr.Image(label="Detection result (annotated)")
+                    status_output    = gr.Textbox(label="Status", interactive=False)
+                    store_output     = gr.Textbox(label="Vector store", interactive=False, value=_store_badge())
+
+            index_btn.click(
+                fn=index_image,
+                inputs=[img_input, note_input],
+                outputs=[detection_output, status_output, store_output],
+            )
+
+        # ── TAB 2 : Query (RAG) ────────────────────────────────────────────────
+        with gr.Tab("💬 Query (RAG)"):
+            gr.Markdown(
+                "Type any question about your indexed images. "
+                "CLIP embeds the query, FAISS retrieves the most similar images, "
+                "and Zephyr-7B generates a grounded answer."
+            )
+            with gr.Row():
+                with gr.Column(scale=1):
+                    query_input = gr.Textbox(
+                        label="Your question",
+                        placeholder="e.g. 'How many people are visible?'  or  'Are there any vehicles?'",
+                        lines=3,
+                    )
+                    topk_slider = gr.Slider(
+                        minimum=1, maximum=5, value=3, step=1,
+                        label="Top-K images to retrieve",
+                    )
+                    query_btn = gr.Button(
+                        "🔎 Search & Generate Answer", variant="primary", size="lg"
+                    )
+
+                with gr.Column(scale=1):
+                    match_img  = gr.Image(label="Best matching image")
+                    llm_answer = gr.Textbox(
+                        label="AI Answer (RAG-grounded)",
+                        lines=6,
+                        interactive=False,
+                    )
+
+            hits_json = gr.JSON(label="Retrieval scores (all top-K)")
+
+            query_btn.click(
+                fn=query_images,
+                inputs=[query_input, topk_slider],
+                outputs=[match_img, llm_answer, hits_json],
+            )
+
+        # ── TAB 3 : How it works ───────────────────────────────────────────────
+        with gr.Tab("🏗️ How it works"):
+            gr.Markdown("""
+## System overview
+
+### Index pipeline
+```
+Image → YOLOv8 detection (objects + confidence)
+      → CLIP ViT-B/32 image encoder (512-d embedding)
+      → L2 normalisation
+      → FAISS IndexFlatIP (cosine similarity store)
+```
+
+### Query / RAG pipeline
+```
+Text query → CLIP text encoder (512-d embedding)
+           → L2 normalisation
+           → FAISS k-NN search (cosine similarity, top-K)
+           → Build RAG prompt: query + retrieved context
+           → Zephyr-7B-β via HF Serverless Inference API
+           → Natural language answer
+```
+
+## Stack
+
+| Component | Technology | Why |
+|---|---|---|
+| Object detection | YOLOv8n (Ultralytics) | Fast CPU inference, production-proven |
+| Visual embedding | CLIP ViT-B/32 | Joint image–text space enables text→image search |
+| Vector index | FAISS IndexFlatIP | Exact cosine k-NN, zero infra |
+| LLM | Zephyr-7B-β (HF API) | Free serverless, good instruction following |
+| UI | Gradio 4.x | Native HF Spaces support |
+
+## Skills demonstrated (mapped to job description)
+- **Computer vision pipeline**: detect → annotate → embed → index
+- **RAG architecture**: embed → retrieve → generate
+- **Multi-modal reasoning**: text queries over an image knowledge base
+- **Vector database**: FAISS with cosine similarity (L2-normalised inner product)
+- **LLM integration**: HF Inference API with prompt engineering
+- **Production patterns**: stateless functions, error handling, model caching
+            """)
+
+    # ── Footer ─────────────────────────────────────────────────────────────────
+    gr.HTML("""
+    <div style="text-align:center; padding:16px 0 4px; color:var(--color-subdued); font-size:12px">
+      VisualRAG · YOLOv8 + CLIP + FAISS + LLM · Built for HuggingFace Spaces
+    </div>
+    """)
+
+demo.launch()