Update app.py

app.py

@@ -1,15 +1,37 @@
+# app.py (clean, no duplicates)
 import gradio as gr
-from pipeline import run_search_and_generate, load_from_hf  # reuse pipeline loader
+from PIL import Image
+
+# Pipeline (must support: run_search_and_generate(user_imgs=[...], user_prompt=...))
+from pipeline import run_search_and_generate, load_from_hf
 
 # --- Quick Starter file paths (must exist in your HF dataset repo) ---
 QS_1_PATH = "images/LakeWater/LakeWater_000550.jpg"
 QS_2_PATH = "images/DenseForest/DenseForest_000000.jpg"
 QS_3_PATH = "images/ResidentialDense/ResidentialDense_001050.jpg"
 
-def run_app(img1, img2, img3, img4, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed):
+
+def _files_to_pil_list(files, n_use: int):
+    """Convert uploaded files to a list of PIL images (use first n_use, capped to 1..4)."""
+    if not files:
+        return []
+    n = max(1, min(4, int(n_use), len(files)))
+    imgs = []
+    for f in files[:n]:
+        path = f.name if hasattr(f, "name") else str(f)
+        imgs.append(Image.open(path).convert("RGB"))
+    return imgs
+
+
+def run_app(files, n_user_imgs, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed):
+    """Main app handler."""
     try:
+        user_imgs = _files_to_pil_list(files, n_user_imgs)
+        if len(user_imgs) == 0:
+            return [], [], [], "Error: Please upload at least 1 image."
+
         retrieved, gen_i2i, gen_t2i, info = run_search_and_generate(
-            user_imgs=[img1, img2, img3, img4],
+            user_imgs=user_imgs,
             user_prompt=prompt,
             k_retrieve=k_retrieve,
             n_i2i=n_i2i,
@@ -19,37 +41,79 @@ def run_app(img1, img2, img3, img4, prompt, k_retrieve, n_i2i, n_t2i, strength_i
             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}"
 
+
 def quickstarter_fill(rel_path):
-    # Fill only Image 1 by default (user can add more)
-    return load_from_hf(rel_path), None, None, None
+    """
+    1-click quick starter:
+    loads a dataset image from HF and returns it as the ONLY uploaded file.
+    The user can then adjust the slider (#images to use) and/or upload more.
+    """
+    img = load_from_hf(rel_path)
+    return [img]  # This will populate the gr.Files-like input (we will use a hidden state)
+
+
+# NOTE:
+# Gradio's gr.Files expects actual uploaded files, not PIL objects.
+# So for Quick Starters, we will route through a separate hidden Image input
+# and run the pipeline directly (no file upload required).
+
+def run_quickstarter(rel_path, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed):
+    try:
+        img = load_from_hf(rel_path)
+        retrieved, gen_i2i, gen_t2i, info = run_search_and_generate(
+            user_imgs=[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: Retrieve + Generate\n"
-        "Upload **up to 4** satellite patches + write a prompt → retrieve similar images and generate new variants.\n\n"
-        "**Goal:** help predict / hallucinate a missing satellite patch from context."
+        "Upload **up to 4 satellite patches** + write a prompt → retrieve similar images and generate new variants.\n\n"
+        "**Quick Start:** click one of the buttons below to instantly see an example (no upload required)."
     )
 
+    # --- Quick Starters (1-click) ---
     gr.Markdown("### Quick Starters (1-click examples)")
     with gr.Row():
-        qs1 = gr.Button("Quick Starter: LakeWater")
-        qs2 = gr.Button("Quick Starter: DenseForest")
-        qs3 = gr.Button("Quick Starter: ResidentialDense")
+        qs1 = gr.Button("LakeWater")
+        qs2 = gr.Button("DenseForest")
+        qs3 = gr.Button("ResidentialDense")
 
     with gr.Row():
+        # --- Inputs ---
         with gr.Column(scale=1):
-            img1 = gr.Image(type="pil", label="Input Image 1 (required)")
-            img2 = gr.Image(type="pil", label="Input Image 2 (optional)")
-            img3 = gr.Image(type="pil", label="Input Image 3 (optional)")
-            img4 = gr.Image(type="pil", label="Input Image 4 (optional)")
+            files = gr.Files(
+                label="Upload up to 4 satellite patch images",
+                file_types=["image"],
+                file_count="multiple"
+            )
+            n_user_imgs = gr.Slider(1, 4, value=1, step=1, label="How many uploaded images to use (1–4)")
 
             prompt = gr.Textbox(
                 label="Prompt (required for generation)",
@@ -68,21 +132,35 @@ with gr.Blocks(title="Satellite Patch: Retrieve + Generate") as demo:
 
             btn = gr.Button("Run")
 
+        # --- Outputs ---
         with gr.Column(scale=2):
             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)
 
+    # Normal run (user-uploaded files)
     btn.click(
         run_app,
-        inputs=[img1, img2, img3, img4, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed],
+        inputs=[files, n_user_imgs, prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed],
         outputs=[out_retr, out_i2i, out_t2i, out_txt],
     )
 
-    # Quick starter: fill Image 1, then user can click Run
-    qs1.click(quickstarter_fill, inputs=[gr.State(QS_1_PATH)], outputs=[img1, img2, img3, img4])
-    qs2.click(quickstarter_fill, inputs=[gr.State(QS_2_PATH)], outputs=[img1, img2, img3, img4])
-    qs3.click(quickstarter_fill, inputs=[gr.State(QS_3_PATH)], outputs=[img1, img2, img3, img4])
+    # Quick Starter runs (1-click)
+    qs1.click(
+        run_quickstarter,
+        inputs=[gr.State(QS_1_PATH), prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed],
+        outputs=[out_retr, out_i2i, out_t2i, out_txt],
+    )
+    qs2.click(
+        run_quickstarter,
+        inputs=[gr.State(QS_2_PATH), prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed],
+        outputs=[out_retr, out_i2i, out_t2i, out_txt],
+    )
+    qs3.click(
+        run_quickstarter,
+        inputs=[gr.State(QS_3_PATH), prompt, k_retrieve, n_i2i, n_t2i, strength_i2i, steps, gen_size, seed],
+        outputs=[out_retr, out_i2i, out_t2i, out_txt],
+    )
 
 demo.launch()