Update app.py

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 import requests
 import os
 from PIL import Image
+import shutil
+import zipfile
+from io import BytesIO
 
 # Streamlit app title
 st.title("Background Remover App")
@@ -23,7 +26,12 @@ if st.button("Start Processing"):
         st.error("Please upload at least one image.")
     else:
         progress_bar = st.progress(0)
-        results = []
+        processed_files = []
+        
+        # Create a temporary directory to store processed files
+        temp_dir = "processed_images"
+        os.makedirs(temp_dir, exist_ok=True)
+        
         for i, uploaded_file in enumerate(uploaded_files):
             image = uploaded_file.read()
             try:
@@ -38,19 +46,35 @@ if st.button("Start Processing"):
                 )
                 
                 if response.status_code == 200:
-                    output_file_path = f"{uploaded_file.name.split('.')[0]}_processed.{output_format}"
+                    output_file_path = os.path.join(temp_dir, f"{uploaded_file.name.split('.')[0]}_processed.{output_format}")
                     with open(output_file_path, "wb") as f:
                         f.write(response.content)
                     
-                    results.append((uploaded_file.name, output_file_path))
+                    processed_files.append(output_file_path)
                     st.success(f"Processed: {uploaded_file.name}")
-                    st.image(Image.open(output_file_path), caption=f"Processed: {output_file_path}")
+                    st.image(Image.open(output_file_path), caption=f"Processed: {uploaded_file.name}")
                 else:
                     st.error(f"Failed to process {uploaded_file.name}: {response.json()}")
             except Exception as e:
                 st.error(f"Error processing {uploaded_file.name}: {e}")
             progress_bar.progress((i + 1) / len(uploaded_files))
         
-        if results:
-            st.success("Processing completed!")
+        if processed_files:
+            # Create a zip archive of all processed images
+            zip_buffer = BytesIO()
+            with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+                for file_path in processed_files:
+                    zip_file.write(file_path, os.path.basename(file_path))
+            zip_buffer.seek(0)
 
+            # Provide download link for the zip file
+            st.success("Processing completed!")
+            st.download_button(
+                label="Download All Processed Images as ZIP",
+                data=zip_buffer,
+                file_name="processed_images.zip",
+                mime="application/zip"
+            )
+        
+        # Cleanup temporary directory
+        shutil.rmtree(temp_dir)