Update app.py

app.py

@@ -1,140 +1,89 @@
 import streamlit as st
 import pydicom
-import matplotlib.pyplot as plt
 import zipfile
 import os
 import subprocess
-from datetime import datetime
-import shutil
-import moviepy.video.io.ImageSequenceClip
-from io import BytesIO
-from tkinter import Tcl
-from PIL import Image
 
 subprocess_executed = False
 
-# logo_image_path = '1.png'  # Replace with your logo image path or URL
-
-# # Display the logo at the top of the page
-# st.image(logo_image_path, width=150)
-# st.title("Automated Abdominal Aortic Aneurysm Detection")
-
-# Upload a ZIP file containing DICOM slices
-st.write("Upload a ZIP file containing DICOM slices")
-uploaded_zip_file = st.file_uploader("Upload a .zip file", type=["zip"])
-
-@st.cache_resource
+# Function to install dependencies
+@st.cache(allow_output_mutation=True)
 def install_dependencies():
-        command = "chmod +x install.sh"
-        try:
-            subprocess.run(command, shell=True, check=True)
-            print("Script 'install.sh' has been made executable.")
-        except subprocess.CalledProcessError as e:
-            print(f"Error while making the script executable: {e}")
+    command = "chmod +x install.sh"
+    subprocess.run(command, shell=True, check=True)
 
-        command = "./install.sh"
-        try:
-            subprocess.run(command, shell=True, check=True)
-            print("Script 'install.sh' has started running.")
-            install_script_executed = True
-        except subprocess.CalledProcessError as e:
-            print(f"Error while running the script: {e}")
+    command = "./install.sh"
+    subprocess.run(command, shell=True, check=True)
 
-@st.cache_resource
+# Function to run inference
+@st.cache(allow_output_mutation=True)
 def run_inference():
-        command = "chmod +x inference.sh"
+    command = "chmod +x inference.sh"
+    subprocess.run(command, shell=True, check=True)
+
+    command = "./inference.sh"
+    subprocess.run(command, shell=True, check=True)
+
+# Function to zip and download files
+def zip_and_download(directory, zip_filename):
+    zipf = zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED)
+    for root, _, files in os.walk(directory):
+        for file in files:
+            zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), directory))
+    zipf.close()
+
+    with open(zip_filename, 'rb') as f:
+        bytes_data = f.read()
+    st.download_button(label="Download Outputs", data=bytes_data, file_name=zip_filename, mime="application/zip")
+
+# Main Streamlit code
+def main():
+    st.write("Upload a ZIP file containing DICOM slices")
+    uploaded_zip_file = st.file_uploader("Upload a .zip file", type=["zip"])
+
+    if uploaded_zip_file is not None:
         try:
-            subprocess.run(command, shell=True, check=True)
-            print("Script 'inference.sh' has been made executable.")
-        except subprocess.CalledProcessError as e:
-            print(f"Error while making the script executable: {e}")
+            install_dependencies()
 
-        command = "./inference.sh"
-        try:
-            subprocess.run(command, shell=True, check=True)
-            print("Script 'inference.sh' has started running.")
-        except subprocess.CalledProcessError as e:
-            print(f"Error while running the script: {e}")
-
-if uploaded_zip_file is not None:
-    try:    
-        install_dependencies()
-        
-        temp_dir = "/home/user/app/C2C/temp_dicom_dir"
-        
-        os.makedirs(temp_dir, exist_ok=True)
-        full_path = os.path.abspath(temp_dir)
-        st.success("Zip file uploaded successfully")
-
-        
-        with zipfile.ZipFile(uploaded_zip_file, "r") as zip_ref:
-            zip_ref.extractall(temp_dir)
-
-        # dicom_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir) if f.endswith(".dcm")]
-        # dicom_files.sort()  
-
-        # # Anonymize DICOM files
-        # for file_path in dicom_files:
-        #     ds = pydicom.dcmread(file_path)
-
-        #     # Anonymize personal information
-        #     if 'PatientName' in ds:
-        #         ds.PatientName = 'Anonymized'
-        #     if 'PatientID' in ds:
-        #         ds.PatientID = '00000000'
-
-        #     print("Anonymized")
-            
-        #     ds.save_as(file_path)
-            
-    except Exception as e:
-        st.error(f"Error: {str(e)}")
-
-if st.button("Analyze"):
-    try:
-        st.success("Analysis started (expected time: 5 mins)")
-        run_inference()
-
-        outputs_directory = "/home/user/app/C2C/outputs"
-        
-        if os.path.exists(outputs_directory):
-            subdirectories = [subdir for subdir in os.listdir(outputs_directory)
-                              if os.path.isdir(os.path.join(outputs_directory, subdir))]
-
-            first_subdirectory = subdirectories[0] if subdirectories else None
-
-            if first_subdirectory:
-                subdirectory_path = os.path.join(outputs_directory, first_subdirectory)
-                temp_dicom_dir_path = os.path.join(subdirectory_path, "temp_dicom_dir")
-
-                dicom_subdirectories = [subdir for subdir in os.listdir(temp_dicom_dir_path)
-                                        if os.path.isdir(os.path.join(temp_dicom_dir_path, subdir))]
-
-                first_dicom_subdirectory = dicom_subdirectories[0] if dicom_subdirectories else None
-
-                if first_dicom_subdirectory:
-                    video_path = os.path.join(temp_dicom_dir_path, first_dicom_subdirectory, "images/summary/aaa.mp4")
-                    image_path = os.path.join(temp_dicom_dir_path, first_dicom_subdirectory, "images/summary/diameter_graph.png")
-                    largest_slice = os.path.join(temp_dicom_dir_path, first_dicom_subdirectory, "images/summary/out.png")
-
-                    if os.path.exists(largest_slice):
-                        st.title("Largest Slice")
-                        st.image(largest_slice, use_column_width=True)
-
-                    # Display the video
-                    if os.path.exists(video_path):
-                        st.title("Video")
-                        st.video(video_path, format="video/mp4")
-
-                    # Display the image
-                    if os.path.exists(image_path):
-                        st.title("Diameter Graph")
-                        st.image(image_path, use_column_width=True)
+            temp_dir = "/home/user/app/C2C/temp_dicom_dir"
+            os.makedirs(temp_dir, exist_ok=True)
+
+            with zipfile.ZipFile(uploaded_zip_file, "r") as zip_ref:
+                zip_ref.extractall(temp_dir)
+
+            st.success("Zip file uploaded successfully")
+
+            if st.button("Analyze"):
+                st.success("Analysis started (expected time: 5 mins)")
+                run_inference()
+
+                outputs_directory = "/home/user/app/C2C/outputs"
+
+                if os.path.exists(outputs_directory):
+                    subdirectories = [subdir for subdir in os.listdir(outputs_directory)
+                                      if os.path.isdir(os.path.join(outputs_directory, subdir))]
+
+                    first_subdirectory = subdirectories[0] if subdirectories else None
+
+                    if first_subdirectory:
+                        subdirectory_path = os.path.join(outputs_directory, first_subdirectory)
+                        temp_dicom_dir_path = os.path.join(subdirectory_path, "temp_dicom_dir")
+
+                        video_path = os.path.join(temp_dicom_dir_path, "images/summary/aaa.mp4")
+                        image_path = os.path.join(temp_dicom_dir_path, "images/summary/diameter_graph.png")
+                        largest_slice = os.path.join(temp_dicom_dir_path, "images/summary/out.png")
+
+                        if os.path.exists(video_path) and os.path.exists(image_path) and os.path.exists(largest_slice):
+                            zip_filename = os.path.join(temp_dir, "outputs.zip")
+                            zip_and_download(temp_dicom_dir_path, zip_filename)
+                        else:
+                            st.error("Output files not found.")
+                    else:
+                        st.warning("No subdirectories found inside 'outputs'.")
                 else:
-                    st.warning("No subdirectories or folders found inside 'temp_dicom_dir'.")
-            else:
-                st.warning("No subdirectories or folders found inside the 'outputs' directory.")
-        else:
-            st.error("The 'outputs' directory does not exist.")
-    except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
+                    st.error("The 'outputs' directory does not exist.")
+        except Exception as e:
+            st.error(f"Error: {str(e)}")
+
+if __name__ == "__main__":
+    main()