Update app.py

app.py

@@ -710,7 +710,7 @@ def _same_tab_download_button(label: str, data_bytes: bytes, file_name: str, mim
 # =========================
 #    Per-file processing
 # =========================
-def process_nifti_case(nifti_path: str, model, rows_acc: List[Dict], per_frame_rows_acc: List[pd.DataFrame]):
+def process_nifti_case(nifti_path: str, model, rows_acc: List[Dict], per_frame_rows_acc: List[pd.DataFrame]) -> Optional[str]:
     pid = Path(nifti_path).stem
     log(f"Using NIfTI input: {nifti_path}")
 
@@ -785,11 +785,13 @@ def process_nifti_case(nifti_path: str, model, rows_acc: List[Dict], per_frame_r
     # GIF: ED vs ES (slices animate) 
     gif_path = gif_animation_for_patient_pred_only(imgs_4d, preds_4d, pid, ed_idx, es_idx, GIFS_DIR)
     try:
-        st.image(gif_path, caption="Generated GIF", use_column_width=True)
+        st.image(gif_path, caption=f"Generated GIF — {pid}", use_column_width=True)
     except TypeError:
-        st.image(gif_path, caption="Generated GIF")
+        st.image(gif_path, caption=f"Generated GIF — {pid}")
     log(f"GIF saved: {gif_path}")
-
+    
+    return gif_path
+    
 
 # =========================
 #           UI
@@ -917,15 +919,22 @@ def main():
     
         st.markdown('</div></div>', unsafe_allow_html=True)
 
-        # ---- Clear stale CSV when a new ZIP is picked ----
+
+
+        # ---- Clear stale results when a new ZIP is picked ----
         if uploaded_zip is not None:
             if st.session_state.get("_last_zip_name") != uploaded_zip.name:
                 st.session_state.pop("csv_bytes", None)
                 st.session_state.pop("csv_name", None)
                 st.session_state.pop("rows_count", None)
                 st.session_state.pop("_dl_token", None)  # reset download-button identity
+        
+                # NEW: clear any previously displayed GIFs
+                st.session_state.pop("gif_paths", None)
+        
                 st.session_state["_last_zip_name"] = uploaded_zip.name
 
+
         # ---- Extract helper ----
         def extract_zip(zip_path, extract_to):
             os.makedirs(extract_to, exist_ok=True)
@@ -974,13 +983,21 @@ def main():
 
                     rows: List[Dict] = []
                     per_frame_rows: List[pd.DataFrame] = []
+                    
+                    # NEW: collect GIFs generated during this run
+                    gif_paths: List[str] = []
+
 
                     for fp in sorted(nii_files):
                         try:
-                            process_nifti_case(fp, model, rows, per_frame_rows)
+                            gif_path = process_nifti_case(fp, model, rows, per_frame_rows)
+                            if gif_path:
+                                gif_paths.append(gif_path)
                         except Exception as e:
                             st.warning(f"Failed: {Path(fp).name} — {e}")
 
+
+
                     # ---- BELOW the GIF(s): write CSV & persist bytes/name ----
                     csv_path = write_all_in_one_csv(rows, per_frame_rows, CSV_DIR)
                     csv_download_name = f"{Path(zip_label).stem}_Results.csv"
@@ -990,12 +1007,25 @@ def main():
                     st.session_state["csv_bytes"] = csv_bytes
                     st.session_state["csv_name"]  = csv_download_name
                     st.session_state["rows_count"] = len(rows)
+                    # NEW: persist GIFs so they can be shown after reruns
+                    st.session_state["gif_paths"] = gif_paths
+
 
-        # ---- Re-render success + robust download on EVERY run (same tab, no 404) ----
+        # ---- Re-render GIFs + success + download on EVERY run ----
         if "csv_bytes" in st.session_state and "csv_name" in st.session_state:
+            # NEW: show GIFs again (persisted from the run that created the CSV)
+            if "gif_paths" in st.session_state and st.session_state["gif_paths"]:
+                st.markdown("### Generated GIFs")
+                for gp in st.session_state["gif_paths"]:
+                    # Optional: show filename/patient id
+                    cap = f"{Path(gp).stem}"
+                    try:
+                        st.image(gp, caption=cap, use_column_width=True)
+                    except TypeError:
+                        st.image(gp, caption=cap)
+        
             st.success(f"Processed {st.session_state.get('rows_count', 0)} NIfTI file(s).")
-
-            # Use our same-tab data:URI button that looks like Streamlit's and turns red on hover
+        
             _same_tab_download_button(
                 label="Download CSV",
                 data_bytes=st.session_state["csv_bytes"],
@@ -1005,7 +1035,6 @@ def main():
             )
 
 
-
     # ===== Tab 2: Dataset =====
     with tab2:
         st.markdown(