Update app.py

app.py

@@ -1,130 +1,62 @@
-import os
-from pathlib import Path
-import random
-
+# app.py
 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:
+def run_app(user_img, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed):
+    try:
+        retrieved, gen_i2i, gen_t2i, info = run_search_and_generate(
+            user_img=user_img,
+            user_prompt=prompt,
+            k_retrieve=k_retrieve,
+            n_i2i=n_i2i,
+            n_t2i=n_t2i,
+            strength_i2i=strength_i2i,
+            steps=steps,
+            gen_size=gen_size,
+            seed=int(seed),
+        )
+        retr_gallery = [(r["img"], f"{r['label']} | cos={r['sim']:.3f}") for r in retrieved]
+        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 = "\n".join([f"{k}: {v}" for k, v in info.items()])
+        return retr_gallery, i2i_gallery, t2i_gallery, summary
+    except Exception as e:
+        return [], [], [], f"Error: {e}"
+
+with gr.Blocks(title="Satellite Patch: Retrieve + 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)."
+        "# Satellite Patch: Retrieve + Generate\n"
+        "Upload a satellite patch + write a prompt → get (1) similar images from the dataset and (2) newly generated images."
     )
 
     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)"
+            user_img = gr.Image(type="pil", label="Upload satellite patch")
+            prompt = gr.Textbox(
+                label="Prompt (required for generation)",
+                value="Satellite-like RGB patch, realistic remote sensing, top-down view",
+                lines=2
             )
-            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)")
+            k_retrieve = gr.Slider(0, 5, value=2, step=1, label="# Retrieved images (0–5)")
+            n_i2i = gr.Slider(0, 5, value=2, step=1, label="# img2img generated (0–5)")
+            n_t2i = gr.Slider(0, 5, value=2, step=1, label="# txt2img generated (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")
+            strength_i2i = gr.Slider(0.25, 0.80, value=0.35, step=0.01, label="img2img strength (lower = closer)")
+            steps = gr.Slider(1, 2, value=1, step=1, label="steps (1–2)")
+            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")
+            btn = 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],
-    )
+            out_retr = gr.Gallery(label="Retrieved from Dataset", columns=5, height=260)
+            out_i2i = gr.Gallery(label="Generated (img2img)", columns=5, height=260)
+            out_t2i = gr.Gallery(label="Generated (txt2img)", columns=5, height=260)
+            out_txt = gr.Textbox(label="Summary", lines=8)
+
+    btn.click(run_app,
+              inputs=[user_img, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed],
+              outputs=[out_retr, out_i2i, out_t2i, out_txt])
 
 demo.launch()