Update app.py

app.py

@@ -18,7 +18,6 @@ logo_image_path = '1.png'
 st.image(logo_image_path, width=150)
 st.title("Automated Abdominal Aortic Aneurysm Detection")
 
-# Function to install dependencies
 @st.cache_resource
 def install_dependencies():
     command = "chmod +x install.sh"
@@ -27,7 +26,6 @@ def install_dependencies():
     command = "./install.sh"
     subprocess.run(command, shell=True, check=True)
 
-# Function to run inference
 @st.cache_resource
 def run_inference():
     command = "chmod +x inference.sh"
@@ -53,7 +51,6 @@ def search_file(start_path, target_file):
             return os.path.join(root, target_file)
     return None
 
-# 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"])
@@ -86,26 +83,35 @@ def main():
                         subdirectory_path = os.path.join(outputs_directory, first_subdirectory)
                         temp_dicom_dir_path = os.path.join(subdirectory_path, "temp_dicom_dir")
 
-                        video_path = search_file(temp_dicom_dir_path, "aaa.mp4")
-                        image_path = search_file(temp_dicom_dir_path, "diameter_graph.png")
-                        largest_slice = search_file(temp_dicom_dir_path, "out.png")
+                        dicom_subdirectories = [subdir for subdir in os.listdir(temp_dicom_dir_path)
+                                                if os.path.isdir(os.path.join(temp_dicom_dir_path, subdir))]
 
-                        if video_path and image_path and largest_slice:
-                            zip_filename = os.path.join(temp_dir, "outputs.zip")
-                            zip_and_download(temp_dicom_dir_path, zip_filename)
+                        dicom_folder = dicom_subdirectories[0] if dicom_subdirectories else None
 
-                            st.title("Largest Slice")
-                            st.image(largest_slice, use_column_width=True)
+                        if dicom_folder:
+                            dicom_folder_path = os.path.join(temp_dicom_dir_path, dicom_folder)
+                            video_path = search_file(dicom_folder_path, "aaa.mp4")
+                            image_path = search_file(dicom_folder_path, "diameter_graph.png")
+                            largest_slice = search_file(dicom_folder_path, "out.png")
 
-                            st.title("Video")
-                            st.video(video_path, format="video/mp4")
+                            if video_path and image_path and largest_slice:
+                                zip_filename = os.path.join(temp_dir, "outputs.zip")
+                                zip_and_download(dicom_folder_path, zip_filename)
 
-                            st.title("Diameter Graph")
-                            st.image(image_path, use_column_width=True)
+                                st.title("Largest Slice")
+                                st.image(largest_slice, use_column_width=True)
+
+                                st.title("Video")
+                                st.video(video_path, format="video/mp4")
+
+                                st.title("Diameter Graph")
+                                st.image(image_path, use_column_width=True)
+                            else:
+                                st.error("Output files not found")
                         else:
-                            st.error("Output files not found")
+                            st.error("DICOM folder not found")
                     else:
-                        st.warning("Ouput files not found")
+                        st.warning("Output files not found")
                 else:
                     st.error("Output files not found")
         except Exception as e: