Update app.py for new hub version and download to local

app.py

@@ -5,6 +5,8 @@ import os
 import posixpath
 from huggingface_hub import hf_hub_download
 from huggingface_hub import list_repo_files
+import io
+import zipfile
 
 # Replace this with your actual Hugging Face repo ID
 REPO_ID = "PortPy-Project/PortPy_Dataset"
@@ -32,7 +34,7 @@ def load_all_metadata(disease_site):
     filtered_patients = patient_df[patient_df["disease_site"] == disease_site]
 
     metadata = {}
-    for patient_id in filtered_patients["patient_id"]:
+    for patient_id in filtered_patients["patient_id"]: # TODO: limit for testing
         # Load structure metadata for the patient
         structs = load_structure_metadata(patient_id)
         # Load beam metadata for the patient
@@ -144,8 +146,7 @@ def download_data(repo_id, patient_ids, beam_ids=None, planner_beam_ids=True, ma
             "CT_Data.h5", "CT_MetaData.json",
             "StructureSet_Data.h5", "StructureSet_MetaData.json",
             "OptimizationVoxels_Data.h5", "OptimizationVoxels_MetaData.json",
-            "PlannerBeams.json",
-            "rt_dose_echo_imrt.dcm", "rt_plan_echo_imrt.dcm"
+            "PlannerBeams.json"
         ]
         for filename in static_files:
             hf_path = posixpath.join("data", patient_id, filename)
@@ -164,6 +165,32 @@ def download_data(repo_id, patient_ids, beam_ids=None, planner_beam_ids=True, ma
                     if attempt == max_retries - 1:
                         st.error(f"Failed to download {hf_path}: {e}")
 
+        # ---------------------------------------------------------------
+        # 2. Download all DICOM files under data/<patient_id>/DicomFiles/
+        # ---------------------------------------------------------------
+        try:
+            all_files = list_repo_files(repo_id, repo_type="dataset")
+            dicom_prefix = f"data/{patient_id}/DicomFiles/"
+            dicom_files = [f for f in all_files if f.startswith(dicom_prefix)]
+
+            for hf_path in dicom_files:
+                for attempt in range(max_retries):
+                    try:
+                        local_path = hf_hub_download(
+                            repo_id=repo_id,
+                            repo_type="dataset",
+                            filename=hf_path,
+                            local_dir=local_dir,
+                            token=token
+                        )
+                        downloaded_files.append(local_path)
+                        break
+                    except Exception as e:
+                        if attempt == max_retries - 1:
+                            st.error(f"Failed to download {hf_path}: {e}")
+
+        except Exception as e:
+            st.error(f"Error listing DICOM files for {patient_id}: {e}")
         if planner_beam_ids:
             planner_file = os.path.join(local_dir, 'data', patient_id, "PlannerBeams.json")
             try:
@@ -239,7 +266,10 @@ def main():
     query_ptv_vol = st.sidebar.number_input("Minimum PTV volume (cc):", value=0)
 
     # Checkbox: Only planner beams
-    only_planner = st.sidebar.checkbox("Show only planner beams", value=True)
+    only_planner = st.sidebar.checkbox(
+        "Show only planner beams (if selected it will download only planner beams)",
+        value=True,
+    )
 
     results_df = filter_matched_data(
         filtered_patients, query_ptv_vol, beam_gantry_filter,
@@ -271,21 +301,71 @@ def main():
     #     st.dataframe(pd.DataFrame(structs), use_container_width=True)
     #     st.subheader(f"📡 Beams for {selected_patient}")
     #     st.dataframe(pd.DataFrame(beams), use_container_width=True)
-
-    with st.expander("Download matched patients"):
+    if "open_download_expander" not in st.session_state:
+        st.session_state["open_download_expander"] = False
+    with st.expander("Download matched patients", expanded=st.session_state["open_download_expander"]):
         # Multi-select and download
         to_download = st.sidebar.multiselect("Select Patients to Download", results_df["patient_id"].tolist())
-        local_dir = st.sidebar.text_input("Enter local directory to download data:", value="./downloaded")
+        # local_dir = st.sidebar.text_input("Enter local directory to download data:", value="./downloaded")
+        # if st.sidebar.button("Download Selected Patients"):
+        #     if to_download:
+        #         patient_to_beams = {
+        #             row["patient_id"]: row["beam_ids"] for ind, row in results_df.iterrows() if ind in to_download
+        #         }
+        #         for pid, beam_ids in patient_to_beams.items():
+        #             download_data(REPO_ID, [pid], beam_ids=beam_ids, planner_beam_ids=False, local_dir=local_dir)
+        #         st.success("Download complete!")
+        #     else:
+        #         st.warning("No patients selected.")
+
         if st.sidebar.button("Download Selected Patients"):
-            if to_download:
+            st.session_state["open_download_expander"] = True  # Force open expander
+            if not to_download:
+                st.warning("No patients selected.")
+            else:
+                progress = st.progress(0)
+                status = st.empty()
+
+                local_dir = "./downloaded"
+                os.makedirs(local_dir, exist_ok=True)
+
                 patient_to_beams = {
-                    row["patient_id"]: row["beam_ids"] for ind, row in results_df.iterrows() if ind in to_download
+                    row["patient_id"]: row["selected_beam_ids"]
+                    for _, row in results_df.iterrows()
+                    if row["patient_id"] in to_download
                 }
-                for pid, beam_ids in patient_to_beams.items():
-                    download_data(REPO_ID, [pid], beam_ids=beam_ids, planner_beam_ids=False, local_dir=local_dir)
-                st.success("Download complete!")
-            else:
-                st.warning("No patients selected.")
+
+                total = len(patient_to_beams)
+                for i, (pid, beam_ids) in enumerate(patient_to_beams.items(), start=1):
+                    status.write(f"Downloading {pid} ({i}/{total})…")
+
+                    download_data(REPO_ID, [pid], beam_ids=beam_ids,
+                                  planner_beam_ids=only_planner,
+                                  local_dir=local_dir)
+
+                    progress.progress(i / total)
+
+
+                status.success("All downloads complete. Preparing zip…")
+
+
+                # Create zip in memory
+                buf = io.BytesIO()
+                with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+                    for root, _, files in os.walk(local_dir):
+                        for f in files:
+                            full_path = os.path.join(root, f)
+                            rel_path = os.path.relpath(full_path, local_dir)
+                            zf.write(full_path, rel_path)
+                buf.seek(0)
+
+                # Trigger file download automatically from the SAME BUTTON CLICK
+                st.download_button(
+                    label="Your download is ready! Click to save.",
+                    data=buf,
+                    file_name="portpy_patients.zip",
+                    mime="application/zip",
+                )
 
         # if st.button("Download Data"):
         #     patients_to_download = results_df["patient_id"].tolist()
@@ -293,4 +373,5 @@ def main():
         #     st.success("Download complete!")
 
 if __name__ == "__main__":
-    main()
+    main()
+    #to run: streamlit run app.py