Create pipeline.py

pipeline.py

@@ -0,0 +1,336 @@
+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)))
+
+ # ===============================================================================================
+
+import torch
+from transformers import AutoImageProcessor, Dinov2Model
+
+# Same model we selected in Part 3
+EMB_MODEL_ID = "facebook/dinov2-small"
+
+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()
+
+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)
+
+    with torch.inference_mode():
+        if device == "cuda":
+            with torch.autocast("cuda", dtype=torch.float16):
+                out = embedder(pixel_values=pixel_values)
+        else:
+            out = embedder(pixel_values=pixel_values)
+
+    v = out.last_hidden_state[:, 0, :].float().cpu().numpy()[0]
+    v = v / (np.linalg.norm(v) + 1e-12)
+    return v.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 [], [], []
+    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)
+
+    info = {
+        "majority_label_from_retrieval": maj_label,
+        "used_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
+    }
+
+    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).")