Create app.py

app.py

@@ -0,0 +1,130 @@
+import os
+from pathlib import Path
+import random
+
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+# Import your pipeline (you already created pipeline.py)
+# Must expose: run_search_and_generate(user_img, k_retrieve, n_i2i, n_t2i, steps_t2i, strength_i2i, gen_size, seed)
+from pipeline import run_search_and_generate
+
+# ---- HF dataset repo (read images directly from the dataset repo) ----
+HF_DATASET_ID = "LevyJonas/sat_land_patches"
+CACHE_DIR = Path("hf_cache")
+CACHE_DIR.mkdir(exist_ok=True, parents=True)
+
+# ---- Utility: download an image file from HF dataset repo (cached) ----
+def load_image_from_hf(rel_path: str) -> Image.Image:
+    local_path = 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(local_path).convert("RGB")
+
+# ---- Quick starters: choose 3 representative images by known labels ----
+# These paths must exist in your dataset repo.
+# If a file does not exist, just replace with another filename from your metadata.csv.
+QUICK_STARTERS = [
+    ("LakeWater example", "images/LakeWater/LakeWater_000000.jpg"),
+    ("DenseForest example", "images/DenseForest/DenseForest_000000.jpg"),
+    ("ResidentialDense example", "images/ResidentialDense/ResidentialDense_000000.jpg"),
+]
+
+def get_quickstarter_image(choice_name: str):
+    for name, rel in QUICK_STARTERS:
+        if name == choice_name:
+            return load_image_from_hf(rel)
+    # fallback
+    return load_image_from_hf(QUICK_STARTERS[0][1])
+
+# ---- App core function ----
+def app_run(user_img, k_retrieve, n_i2i, n_t2i, strength_i2i, steps_t2i, gen_size, seed):
+    if user_img is None:
+        return [], [], [], "Please upload an image."
+
+    # run pipeline
+    retrieved, gen_i2i, gen_t2i, info = run_search_and_generate(
+        user_img=user_img,
+        k_retrieve=int(k_retrieve),
+        n_i2i=int(n_i2i),
+        n_t2i=int(n_t2i),
+        steps_t2i=int(steps_t2i),
+        strength_i2i=float(strength_i2i),
+        gen_size=int(gen_size),
+        seed=int(seed),
+    )
+
+    # Format outputs for Gradio Gallery (image, caption)
+    retrieved_gallery = []
+    for r in retrieved:
+        cap = f"{r.get('label','')} | cos={r.get('sim',0):.3f}"
+        retrieved_gallery.append((r["img"], cap))
+
+    i2i_gallery = [(im, f"img2img #{i+1}") for i, im in enumerate(gen_i2i)]
+    t2i_gallery = [(im, f"txt2img #{i+1}") for i, im in enumerate(gen_t2i)]
+
+    # Summary text
+    summary_lines = []
+    if isinstance(info, dict):
+        summary_lines.append(f"majority_label_from_retrieval: {info.get('majority_label_from_retrieval')}")
+        summary_lines.append(f"used_prompt: {info.get('used_prompt')}")
+        summary_lines.append(f"k_retrieve: {info.get('k_retrieve')}, n_img2img: {info.get('n_img2img')}, n_txt2img: {info.get('n_t2i', info.get('n_txt2img'))}")
+        summary_lines.append(f"strength_i2i: {info.get('requested_strength_img2img', info.get('strength_i2i'))}, steps_txt2img: {info.get('steps_txt2img', info.get('steps_t2i'))}")
+        summary_lines.append(f"gen_size: {info.get('gen_size')}, seed: {info.get('seed')}")
+
+    return retrieved_gallery, i2i_gallery, t2i_gallery, "\n".join(summary_lines)
+
+# ---- Quick starter click: set input image ----
+def load_quickstarter(choice_name):
+    return get_quickstarter_image(choice_name)
+
+# ---- Build UI ----
+with gr.Blocks(title="Satellite Patch Search + Generate") as demo:
+    gr.Markdown(
+        "# Satellite Patch Search + Generate\n"
+        "Upload a satellite-like patch → get **Top-K similar patches** from the dataset + **new generated variants** "
+        "(Image-to-Image + Text-to-Image)."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            inp = gr.Image(type="pil", label="Upload satellite patch (input)")
+
+            qs = gr.Dropdown(
+                choices=[x[0] for x in QUICK_STARTERS],
+                value=QUICK_STARTERS[0][0],
+                label="Quick Starters (1-click examples)"
+            )
+            btn_qs = gr.Button("Load Quick Starter")
+
+            k_retrieve = gr.Slider(0, 5, value=2, step=1, label="How many images to retrieve from DB (0–5)")
+            n_i2i = gr.Slider(0, 5, value=2, step=1, label="How many img2img generated images (0–5)")
+            n_t2i = gr.Slider(0, 5, value=2, step=1, label="How many txt2img generated images (0–5)")
+
+            strength_i2i = gr.Slider(0.25, 0.80, value=0.35, step=0.01, label="img2img strength (lower = closer to input)")
+            steps_t2i = gr.Slider(1, 2, value=1, step=1, label="Generation steps (1–2 for sd-turbo)")
+            gen_size = gr.Radio([384, 512], value=512, label="Generation size")
+            seed = gr.Number(value=42, precision=0, label="Seed")
+
+            btn_run = gr.Button("Run")
+
+        with gr.Column(scale=2):
+            out_retr = gr.Gallery(label="Retrieved from Dataset (Top-K)", columns=5, height=260)
+            out_i2i = gr.Gallery(label="Generated (Image-to-Image)", columns=5, height=260)
+            out_t2i = gr.Gallery(label="Generated (Text-to-Image)", columns=5, height=260)
+            out_txt = gr.Textbox(label="Summary", lines=6)
+
+    # Wire buttons
+    btn_qs.click(load_quickstarter, inputs=[qs], outputs=[inp])
+    btn_run.click(
+        app_run,
+        inputs=[inp, k_retrieve, n_i2i, n_t2i, strength_i2i, steps_t2i, gen_size, seed],
+        outputs=[out_retr, out_i2i, out_t2i, out_txt],
+    )
+
+demo.launch()