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

app.py

@@ -8,6 +8,7 @@ from huggingface_hub import list_repo_files
 import io
 import zipfile
 import shutil
+import tempfile, uuid
 
 # Replace this with your actual Hugging Face repo ID
 REPO_ID = "PortPy-Project/PortPy_Dataset"
@@ -140,7 +141,7 @@ def filter_matched_data(filtered_patients, query_ptv_vol, beam_gantry_filter,
 
     return pd.DataFrame(matched)
 
-def download_data(repo_id, patient_ids, beam_ids=None, planner_beam_ids=True, max_retries=2, local_dir='./'):
+def download_data(repo_id, patient_ids, beam_ids=None, planner_beam_ids=True, max_retries=2, local_dir='./', download_dicom=True):
     from huggingface_hub import hf_hub_download
 
     downloaded_files = []
@@ -171,29 +172,30 @@ def download_data(repo_id, patient_ids, beam_ids=None, planner_beam_ids=True, ma
         # ---------------------------------------------------------------
         # 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 download_dicom:
+            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:
@@ -346,31 +348,60 @@ def main():
 
                     download_data(REPO_ID, [pid], beam_ids=beam_ids,
                                   planner_beam_ids=only_planner,
-                                  local_dir=local_dir)
+                                  local_dir=local_dir, download_dicom=False)
 
                     progress.progress(i / total)
 
 
                 status.success("All downloads complete. Preparing zip…")
 
+                zip_path = os.path.join(tempfile.gettempdir(), f"portpy_patients_{uuid.uuid4().hex}.zip")
+
+                # optional: guard size to avoid crashes
+                total_bytes = 0
+                for root, _, files in os.walk(local_dir):
+                    for f in files:
+                        total_bytes += os.path.getsize(os.path.join(root, f))
+                total_gb = total_bytes / (1024 ** 3)
+                status.write(f"Preparing zip (~{total_gb:.2f} GB)…")
+                if total_gb > 40.0:
+                    st.error("Selection too large for a single zip. Please download fewer patients.")
+                    st.stop()
+
+                if os.path.exists(zip_path):
+                    os.remove(zip_path)
 
-                # Create zip in memory
-                buf = io.BytesIO()
-                with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) as zf:
+                with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_STORED, allowZip64=True) 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",
-                )
+
+                with open(zip_path, "rb") as fp:
+                    st.download_button(
+                        label="Your download is ready! Click to save.",
+                        data=fp,
+                        file_name="portpy_patients.zip",
+                        mime="application/zip",
+                    )
+                # # Create zip in memory
+                # buf = io.BytesIO()
+                # with zipfile.ZipFile(buf, "w", zipfile.ZIP_STORED) 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()