Update pipeline.py

pipeline.py

@@ -1,336 +1,129 @@
-DATA_ROOT = Path("sat_land_patches")                 # local dataset folder
-EMB_DIR   = Path("embeddings_part3")                 # where Part 3 outputs are
-
-# Load best embeddings + metadata saved as .npy + .csv
-DB_E = np.load(EMB_DIR / "best_embeddings.npy").astype(np.float32)   # (N,D) L2-normalized
-db_meta = pd.read_csv(EMB_DIR / "best_metadata.csv")                # id,label,filename,model_id
-
-DB_labels = db_meta["label"].values
-DB_files  = db_meta["filename"].values
-
-print("DB:", DB_E.shape, "| labels:", len(np.unique(DB_labels)))
+# pipeline.py
+import math
+from pathlib import Path
 
- # ===============================================================================================
+import numpy as np
+import pandas as pd
+from PIL import Image
 
 import torch
 from transformers import AutoImageProcessor, Dinov2Model
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+from huggingface_hub import hf_hub_download
+
+HF_DATASET_ID = "LevyJonas/sat_land_patches"
+CACHE_DIR = Path("hf_cache"); CACHE_DIR.mkdir(exist_ok=True, parents=True)
+
+EMB_DIR = Path("embeddings_part3")
+DB_E = np.load(EMB_DIR / "best_embeddings.npy").astype(np.float32)
+META = pd.read_csv(EMB_DIR / "best_metadata.csv")
+DB_FILES = META["filename"].values
+DB_LABELS = META["label"].values
 
-# Same model we selected in Part 3
 EMB_MODEL_ID = "facebook/dinov2-small"
+GEN_MODEL_ID = "stabilityai/sd-turbo"
+NEG = "cartoon, illustration, anime, text, watermark, logo, low quality, blurry, distorted, unrealistic"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print("Device:", device)
 
-processor = AutoImageProcessor.from_pretrained(EMB_MODEL_ID)
-embedder = Dinov2Model.from_pretrained(EMB_MODEL_ID).to(device)
-embedder.eval()
+# --- embedder ---
+proc = AutoImageProcessor.from_pretrained(EMB_MODEL_ID)
+emb = Dinov2Model.from_pretrained(EMB_MODEL_ID).to(device).eval()
 
-def embed_query_image(img: Image.Image) -> np.ndarray:
-    """Return L2-normalized embedding vector for a query PIL image."""
-    img = img.convert("RGB")
-    inputs = processor(images=[img], return_tensors="pt")
-    pixel_values = inputs["pixel_values"].to(device)
+# --- generators ---
+dtype = torch.float16 if device == "cuda" else torch.float32
+txt2img = StableDiffusionPipeline.from_pretrained(GEN_MODEL_ID, torch_dtype=dtype, variant="fp16" if device=="cuda" else None).to(device)
+img2img = StableDiffusionImg2ImgPipeline.from_pretrained(GEN_MODEL_ID, torch_dtype=dtype, variant="fp16" if device=="cuda" else None).to(device)
+txt2img.set_progress_bar_config(disable=True)
+img2img.set_progress_bar_config(disable=True)
+
+def _cap(x): return int(max(0, min(5, int(x))))
 
+def load_from_hf(rel_path: str) -> Image.Image:
+    p = hf_hub_download(repo_id=HF_DATASET_ID, repo_type="dataset", filename=rel_path,
+                        local_dir=str(CACHE_DIR), local_dir_use_symlinks=False)
+    return Image.open(p).convert("RGB")
+
+def embed_image(pil_img: Image.Image) -> np.ndarray:
+    x = proc(images=[pil_img.convert("RGB")], return_tensors="pt")["pixel_values"].to(device)
     with torch.inference_mode():
         if device == "cuda":
             with torch.autocast("cuda", dtype=torch.float16):
-                out = embedder(pixel_values=pixel_values)
+                out = emb(pixel_values=x)
         else:
-            out = embedder(pixel_values=pixel_values)
-
+            out = emb(pixel_values=x)
     v = out.last_hidden_state[:, 0, :].float().cpu().numpy()[0]
-    v = v / (np.linalg.norm(v) + 1e-12)
-    return v.astype(np.float32)
+    return (v / (np.linalg.norm(v) + 1e-12)).astype(np.float32)
 
-# ===============================================================================================
-
-from collections import Counter
-
-def retrieve_topk(query_vec: np.ndarray, k: int = 5):
-    """Cosine similarity = dot product because vectors are L2-normalized."""
-    k = int(max(0, min(5, k)))  # cap at 5
-    if k == 0:
-        return [], [], []
+def retrieve(query_vec: np.ndarray, k: int):
+    k = _cap(k)
+    if k == 0: return []
     sims = DB_E @ query_vec
     idx = np.argsort(-sims)[:k]
-    return idx, sims[idx], DB_labels[idx]
-
-def majority_label(labels):
-    if len(labels) == 0:
-        return None
-    return Counter(labels.tolist()).most_common(1)[0][0]
-
-def load_db_image(rel_path: str) -> Image.Image:
-    return Image.open(DATA_ROOT / rel_path).convert("RGB")
-
-# ===============================================================================================
-
-# Uses sd-turbo for both text-to-image and image-to-image.
-# Notes: keep steps low for speed (1-2). Works best on GPU.
-
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
-
-GEN_MODEL_ID = "stabilityai/sd-turbo"
-
-txt2img = StableDiffusionPipeline.from_pretrained(
-    GEN_MODEL_ID, torch_dtype=torch.float16, variant="fp16"
-).to("cuda")
-txt2img.set_progress_bar_config(disable=True)
-
-img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
-    GEN_MODEL_ID, torch_dtype=torch.float16, variant="fp16"
-).to("cuda")
-img2img.set_progress_bar_config(disable=True)
-
-# Optional speed-ups (safe to ignore if not available)
-try:
-    txt2img.enable_xformers_memory_efficient_attention()
-    img2img.enable_xformers_memory_efficient_attention()
-except Exception:
-    pass
-
-# ===============================================================================================
-
-# PROMPTS must exist: dict[label] -> prompt
-assert "PROMPTS" in globals(), "PROMPTS dict not found. Paste your PROMPTS (30 labels) before running Part 4."
-NEGATIVE = "cartoon, illustration, anime, text, watermark, logo, low quality, blurry, distorted, unrealistic"
-
-# ===============================================================================================
-
-import math
-from collections import Counter
-import torch
-from PIL import Image
-
-def _cap_0_5(x):
-    """Cap an integer to the range [0, 5]."""
-    return int(max(0, min(5, int(x))))
-
-def _majority_label(arr):
-    """Return majority label from an array of labels (or None)."""
-    if len(arr) == 0:
-        return None
-    return Counter(arr.tolist()).most_common(1)[0][0]
-
-def retrieve_topk(query_vec, k=5):
-    """
-    Retrieve top-k most similar items from DB using cosine similarity.
-    Cosine similarity = dot product because vectors are L2-normalized.
-    Returns: list of dicts with image, label, similarity, filename.
-    """
-    k = _cap_0_5(k)
-    if k == 0:
-        return []
-
-    sims = DB_E @ query_vec
-    idx = np.argsort(-sims)[:k]
-
-    results = []
-    for i in idx:
-        rel = DB_files[i]
-        results.append({
-            "img": load_db_image(rel),
-            "label": DB_labels[i],
-            "sim": float(sims[i]),
-            "filename": rel
-        })
-    return results
-
-def _safe_img2img_steps(strength, user_steps):
-    """
-    Diffusers img2img requires at least 1 effective denoising step:
-      effective = int(num_inference_steps * strength) >= 1
-    If not, tensors become empty and you get the reshape error.
-    This function chooses a safe num_inference_steps automatically.
-    """
-    strength = float(strength)
-    strength = max(1e-3, min(1.0, strength))  # keep in (0,1]
-
-    steps = int(user_steps)
-    steps = max(1, min(6, steps))  # keep small for turbo
-
-    # Ensure effective steps >= 1
-    if int(steps * strength) < 1:
-        steps = int(math.ceil(1.0 / strength))
-
-    # Clamp again to keep runtime bounded (still safe)
-    steps = max(2, min(6, steps))
-    return steps, strength
-
-def run_search_and_generate(
-    user_img: Image.Image,
-    k_retrieve: int = 2,
-    n_i2i: int = 1,
-    n_t2i: int = 1,
-    steps_t2i: int = 1,
-    strength_i2i: float = 0.35,
-    gen_size: int = 512,
-    seed: int = 123
-):
-    """
-    Pipeline:
-      1) Embed input image (DINOv2)
-      2) Retrieve top-k similar images from DB
-      3) Choose prompt based on majority retrieved label
-      4) Generate n_i2i images using img2img
-      5) Generate n_t2i images using txt2img
-
-    Returns:
-      retrieved: list[dict]  (each dict has img/label/sim/filename)
-      gen_i2i:   list[PIL.Image]
-      gen_t2i:   list[PIL.Image]
-      info:      dict (prompt/labels/params)
-    """
-    # --- Cap counts to [0,5] for app safety ---
-    k_retrieve = _cap_0_5(k_retrieve)
-    n_i2i = _cap_0_5(n_i2i)
-    n_t2i = _cap_0_5(n_t2i)
-
-    # --- Embed query image ---
-    q_vec = embed_query_image(user_img)
-
-    # --- Retrieve ---
-    retrieved = retrieve_topk(q_vec, k=k_retrieve)
-
-    # Decide label/prompt from retrieval results
-    retrieved_labels = np.array([r["label"] for r in retrieved]) if len(retrieved) else np.array([])
-    maj_label = _majority_label(retrieved_labels) if len(retrieved_labels) else None
-
-    prompt = PROMPTS.get(
-        maj_label,
-        "Satellite-like RGB patch, realistic remote sensing, top-down view"
-    )
-
-    # Prepare init image for img2img
-    init_img = user_img.convert("RGB").resize((gen_size, gen_size))
-
-    # --- Generate (img2img) ---
-    gen_i2i = []
-    if n_i2i > 0:
-        safe_steps_i2i, safe_strength = _safe_img2img_steps(strength_i2i, steps_t2i)
-
-        for i in range(n_i2i):
-            g = torch.Generator("cuda").manual_seed(seed + 10*i)
-            with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
-                im = img2img(
-                    prompt=prompt,
-                    negative_prompt=NEGATIVE,
-                    image=init_img,
-                    strength=safe_strength,
-                    num_inference_steps=safe_steps_i2i,
-                    guidance_scale=0.0,
-                    generator=g
-                ).images[0]
-            gen_i2i.append(im)
-
-    # --- Generate (txt2img) ---
-    gen_t2i = []
-    if n_t2i > 0:
-        # sd-turbo is designed for 1–2 steps
-        steps_txt = max(1, min(2, int(steps_t2i)))
-
-        for i in range(n_t2i):
-            g = torch.Generator("cuda").manual_seed(seed + 100 + 10*i)
-            with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
-                im = txt2img(
-                    prompt=prompt,
-                    negative_prompt=NEGATIVE,
-                    num_inference_steps=steps_txt,
-                    guidance_scale=0.0,
-                    height=gen_size,
-                    width=gen_size,
-                    generator=g
-                ).images[0]
-            gen_t2i.append(im)
+    return [{
+        "img": load_from_hf(DB_FILES[i]),
+        "label": DB_LABELS[i],
+        "sim": float(sims[i]),
+        "filename": DB_FILES[i],
+    } for i in idx]
+
+def _safe_i2i_steps(strength: float, user_steps: int):
+    strength = float(max(1e-3, min(1.0, strength)))
+    steps = int(max(1, min(2, user_steps)))      # user slider 1..2
+    # ensure int(steps_i2i * strength) >= 1
+    steps_i2i = max(2, int(math.ceil(1.0 / strength)))
+    steps_i2i = min(6, steps_i2i)                # keep fast
+    return steps, steps_i2i, strength
+
+def run_search_and_generate(user_img: Image.Image, user_prompt: str,
+                            k_retrieve=2, n_i2i=2, n_t2i=2,
+                            strength_i2i=0.35, steps=1, gen_size=512, seed=42):
+    k_retrieve, n_i2i, n_t2i = _cap(k_retrieve), _cap(n_i2i), _cap(n_t2i)
+    prompt = (user_prompt or "").strip()
+    if not prompt:
+        raise ValueError("Please enter a prompt (required for generation).")
+
+    q = embed_image(user_img)
+    retrieved = retrieve(q, k_retrieve)
+
+    steps_txt, steps_i2i, strength = _safe_i2i_steps(strength_i2i, steps)
+    init = user_img.convert("RGB").resize((gen_size, gen_size))
+
+    gen_i2i, gen_t2i = [], []
+    for i in range(n_i2i):
+        g = torch.Generator(device).manual_seed(seed + 10*i)
+        with torch.inference_mode():
+            if device == "cuda":
+                with torch.autocast("cuda", dtype=torch.float16):
+                    gen_i2i.append(img2img(prompt=prompt, negative_prompt=NEG, image=init,
+                                           strength=strength, num_inference_steps=steps_i2i,
+                                           guidance_scale=0.0, generator=g).images[0])
+            else:
+                gen_i2i.append(img2img(prompt=prompt, negative_prompt=NEG, image=init,
+                                       strength=strength, num_inference_steps=steps_i2i,
+                                       guidance_scale=0.0, generator=g).images[0])
+
+    for i in range(n_t2i):
+        g = torch.Generator(device).manual_seed(seed + 100 + 10*i)
+        with torch.inference_mode():
+            if device == "cuda":
+                with torch.autocast("cuda", dtype=torch.float16):
+                    gen_t2i.append(txt2img(prompt=prompt, negative_prompt=NEG,
+                                           num_inference_steps=steps_txt, guidance_scale=0.0,
+                                           height=gen_size, width=gen_size, generator=g).images[0])
+            else:
+                gen_t2i.append(txt2img(prompt=prompt, negative_prompt=NEG,
+                                       num_inference_steps=steps_txt, guidance_scale=0.0,
+                                       height=gen_size, width=gen_size, generator=g).images[0])
 
     info = {
-        "majority_label_from_retrieval": maj_label,
-        "used_prompt": prompt,
+        "prompt": prompt,
         "k_retrieve": k_retrieve,
         "n_img2img": n_i2i,
         "n_txt2img": n_t2i,
-        "steps_txt2img": max(1, min(2, int(steps_t2i))),
-        "requested_strength_img2img": float(strength_i2i),
-        "gen_size": gen_size,
-        "seed": seed
+        "strength_i2i": strength,
+        "steps_txt2img": steps_txt,
+        "steps_img2img": steps_i2i,
+        "dataset": HF_DATASET_ID
     }
-
-    return retrieved, gen_i2i, gen_t2i, info
-
-# ===============================================================================================
-
-import random
-import matplotlib.pyplot as plt
-from PIL import Image
-
-# --- Pick a demo input image from your dataset ---
-demo_rel = DB_files[random.randrange(len(DB_files))]
-user_img = load_db_image(demo_rel)
-
-# --- Run pipeline (you can change these 0-5 values) ---
-k_retrieve = 2   # 0..5 images from database
-n_i2i = 2        # 0..5 new images via image-to-image
-n_t2i = 2        # 0..5 new images via text-to-image
-
-retrieved, gen_i2i, gen_t2i, info = run_search_and_generate(
-    user_img=user_img,
-    k_retrieve=k_retrieve,
-    n_i2i=n_i2i,
-    n_t2i=n_t2i,
-    steps_t2i=1,          # txt2img steps (1-2 recommended for sd-turbo)
-    strength_i2i=0.35,    # img2img strength (0.25-0.60 is typical)
-    gen_size=512,
-    seed=42
-)
-
-print("=== PIPELINE INFO ===")
-for k, v in info.items():
-    print(f"{k}: {v}")
-
-# --- Helper to show a gallery in one row ---
-def show_row(images, titles, fig_w=16, fig_h=3, suptitle=None):
-    n = len(images)
-    if n == 0:
-        print(suptitle or "No images to show.")
-        return
-    plt.figure(figsize=(fig_w, fig_h))
-    for i, (im, t) in enumerate(zip(images, titles), 1):
-        ax = plt.subplot(1, n, i)
-        ax.imshow(im)
-        ax.set_title(t, fontsize=9)
-        ax.axis("off")
-    if suptitle:
-        plt.suptitle(suptitle, fontsize=12)
-    plt.tight_layout()
-    plt.show()
-
-# 1) Show input image
-show_row(
-    images=[user_img],
-    titles=[f"USER INPUT\n{demo_rel}"],
-    fig_w=6,
-    fig_h=4,
-    suptitle="User Input"
-)
-
-# 2) Show retrieved images
-if len(retrieved) > 0:
-    ret_imgs = [r["img"] for r in retrieved]
-    ret_titles = [f"{r['label']}\ncos={r['sim']:.3f}" for r in retrieved]
-    show_row(ret_imgs, ret_titles, fig_w=3.2*len(ret_imgs), fig_h=3, suptitle="Top-K Retrieved from Database")
-else:
-    print("No retrieval results (k_retrieve=0).")
-
-# 3) Show generated img2img images
-if len(gen_i2i) > 0:
-    titles = [f"img2img #{i+1}" for i in range(len(gen_i2i))]
-    show_row(gen_i2i, titles, fig_w=3.2*len(gen_i2i), fig_h=3, suptitle="Generated (Image-to-Image)")
-else:
-    print("No img2img generated (n_i2i=0).")
-
-# 4) Show generated txt2img images
-if len(gen_t2i) > 0:
-    titles = [f"txt2img #{i+1}" for i in range(len(gen_t2i))]
-    show_row(gen_t2i, titles, fig_w=3.2*len(gen_t2i), fig_h=3, suptitle="Generated (Text-to-Image)")
-else:
-    print("No txt2img generated (n_t2i=0).")
+    return retrieved, gen_i2i, gen_t2i, info