Added in slider after episode is completed

app.py

@@ -131,8 +131,18 @@ def process_search_tta(
 ):
     """Run both TTA and non-TTA search episodes concurrently and stream both heat-maps."""
 
-    # Disable Run button and clear image/status outputs
-    yield gr.update(interactive=False), gr.update(value=None), gr.update(value=None), gr.update(value="Initializing model…"), gr.update(value="Initializing model…")
+    # Disable Run button and clear image/status outputs, hide sliders, clear frame states
+    yield (
+        gr.update(interactive=False),
+        gr.update(value=None),
+        gr.update(value=None),
+        gr.update(value="Initializing model…"),
+        gr.update(value="Initializing model…"),
+        gr.update(visible=False),
+        gr.update(visible=False),
+        [],
+        [],
+    )
 
     # Bail early if satellite image missing
     if sat_path is None:
@@ -246,7 +256,17 @@ def process_search_tta(
 
             # Emit update if we have a new frame OR if the status label has changed
             if updated or status_tta != prev_status_tta or status_no != prev_status_no:
-                yield gr.update(interactive=False), last_tta, last_no, gr.update(value=status_tta), gr.update(value=status_no)
+                yield (
+                    gr.update(interactive=False),
+                    last_tta,
+                    last_no,
+                    gr.update(value=status_tta, visible=True),
+                    gr.update(value=status_no, visible=True),
+                    gr.update(visible=False),
+                    gr.update(visible=False),
+                    gr.update(),
+                    gr.update(),
+                )
                 prev_status_tta = status_tta
                 prev_status_no  = status_no
 
@@ -270,8 +290,24 @@ def process_search_tta(
             sent_no.add(fp)
             last_no = fp
 
-    # Final emit after both finish (and after final scan)
-    yield gr.update(interactive=True), last_tta, last_no, gr.update(value="Done."), gr.update(value="Done.")
+    # Prepare frames list and slider configs
+    frames_tta = sorted(glob.glob(os.path.join(gifs_dir_tta, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+    frames_no = sorted(glob.glob(os.path.join(gifs_dir_no, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+    n_tta = len(frames_tta) or 1  # prevent zero-range slider
+    n_no  = len(frames_no) or 1
+
+    # Final emit: re-enable button, hide statuses, show sliders set to last frame
+    yield (
+        gr.update(interactive=True),
+        last_tta,
+        last_no,
+        gr.update(visible=False),
+        gr.update(visible=False),
+        gr.update(visible=True, minimum=1, maximum=n_tta, value=n_tta),
+        gr.update(visible=True, minimum=1, maximum=n_no,  value=n_no),
+        frames_tta,
+        frames_no,
+    )
 
 
 # ────────────────────────── Gradio UI ─────────────────────────────────
@@ -333,9 +369,24 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
             # gr.Markdown("### Live Heatmap (with TTA)")
             display_img_tta = gr.Image(label="Heatmap (TTA per 20 steps)", type="filepath", height=400)  # 512
             status_tta = gr.Markdown("")
-            # gr.Markdown("### Live Heatmap (without TTA)")
+            slider_tta = gr.Slider(label="TTA Frame", minimum=1, maximum=1, step=1, value=1, visible=False)
+
             display_img_no_tta = gr.Image(label="Heatmap (no TTA)", type="filepath", height=400)  # 512
             status_no_tta = gr.Markdown("")
+            slider_no = gr.Slider(label="No-TTA Frame", minimum=1, maximum=1, step=1, value=1, visible=False)
+
+            frames_state_tta = gr.State([])
+            frames_state_no = gr.State([])
+
+    # Slider callbacks (updates image when user drags slider)
+    def _show_frame(idx: int, frames: list[str]):
+        # Slider is 1-indexed; convert to 0-indexed list access
+        if 1 <= idx <= len(frames):
+            return frames[idx - 1]
+        return gr.update()
+
+    slider_tta.change(_show_frame, inputs=[slider_tta, frames_state_tta], outputs=display_img_tta)
+    slider_no.change(_show_frame, inputs=[slider_no, frames_state_no], outputs=display_img_no_tta)
 
     # Bind callback
 
@@ -362,7 +413,7 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
                 ],
             ],
             inputs=[sat_input, ground_input, taxonomy_input],
-            outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta],
+            outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta, slider_tta, slider_no, frames_state_tta, frames_state_no],
             fn=process_search_tta,
             cache_examples=False,
         )
@@ -371,7 +422,7 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
     run_btn.click(
         fn=process_search_tta,
         inputs=[sat_input, ground_input, taxonomy_input],
-        outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta],
+        outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta, slider_tta, slider_no, frames_state_tta, frames_state_no],
     )
 
     # Footer to point out to model and data from app page.