update

app.py

@@ -1,10 +1,16 @@
-import os, math, csv
+import os, math, csv, shutil
 import streamlit as st
 from streamlit_image_select import image_select
 import cv2
 import numpy as np
 from PIL import Image
 import matplotlib.colors as mcolors
+from io import BytesIO
+
+
+MAX_GALLERY_IMAGES = 50
+GALLERY_IMAGE_SIZE = 128
+
 
 class DirectoryManager:
     def __init__(self, output_dir):
@@ -220,13 +226,89 @@ def save_img(directory, base_name, mask):
     mask_name = os.path.join(directory, os.path.basename(base_name))
     cv2.imwrite(mask_name, mask)
 
-def main():
+def validate_dir(dir):
+    if not os.path.exists(dir):
+        os.makedirs(dir, exist_ok=True)
+
+def make_zipfile(source_dir, output_filename):
+    shutil.make_archive(output_filename, 'zip', source_dir)
+    return output_filename + '.zip'
+
+def save_uploaded_file(directory, img_file, image=None):
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+    # Assuming the uploaded file is an image
+    if image is None:
+        with Image.open(img_file) as image:
+            full_path = os.path.join(directory, img_file.name)
+            image.save(full_path, "JPEG")
+        # Return the full path of the saved image
+        return full_path
+    else:
+        full_path = os.path.join(directory, img_file.name)
+        image.save(full_path, "JPEG")
+        return full_path
+
+def create_download_button(dir_to_zip, zip_filename):
+    zip_filepath = make_zipfile(dir_to_zip, zip_filename)
+    with open(zip_filepath, 'rb') as f:
+        bytes_io = BytesIO(f.read())
+    st.download_button(
+        label=f"Download Results for{st.session_state['processing_add_on']}",type='primary',
+        data=bytes_io,
+        file_name=os.path.basename(zip_filepath),
+        mime='application/zip'
+    )
+
+def delete_directory(dir_path):
+    try:
+        shutil.rmtree(dir_path)
+        st.session_state['input_list'] = []
+        st.session_state['input_list_small'] = []
+        # st.success(f"Deleted previously uploaded images, making room for new images: {dir_path}")
+    except OSError as e:
+        st.error(f"Error: {dir_path} : {e.strerror}")
+
+def clear_image_gallery():
+    delete_directory(st.session_state['dir_uploaded_images'])
+    delete_directory(st.session_state['dir_uploaded_images_small'])
+    validate_dir(st.session_state['dir_uploaded_images'])
+    validate_dir(st.session_state['dir_uploaded_images_small'])
 
+def main():
     _, R_coverage, R_plot_area_cm2, R_save = st.columns([5,2,2,2])
     img_gallery, img_main, img_seg, img_green, img_warp = st.columns([1,4,2,2,2])
 
-    dir_input = st.text_input("Input directory for images:", value="D:\Dropbox\GreenSight\demo")
-    dir_output = st.text_input("Output directory:", value="D:\Dropbox\GreenSight\demo_out")
+    st.session_state['dir_uploaded_images'] = os.path.join(st.session_state.dir_home,'uploads')
+    st.session_state['dir_uploaded_images_small'] = os.path.join(st.session_state.dir_home,'uploads_small')
+    uploaded_files = st.file_uploader("Upload Images", type=['jpg', 'jpeg'], accept_multiple_files=True, key=st.session_state['uploader_idk'])
+    if uploaded_files:
+        # Clear input image gallery and input list
+        clear_image_gallery()
+
+        # Process the new iamges
+        for uploaded_file in uploaded_files:
+            file_path = save_uploaded_file(st.session_state['dir_uploaded_images'], uploaded_file)
+            st.session_state['input_list'].append(file_path)
+
+            img = Image.open(file_path)
+            img.thumbnail((GALLERY_IMAGE_SIZE, GALLERY_IMAGE_SIZE), Image.Resampling.LANCZOS)  
+            file_path_small = save_uploaded_file(st.session_state['dir_uploaded_images_small'], uploaded_file, img)
+            st.session_state['input_list_small'].append(file_path_small)
+            print(uploaded_file.name)
+
+        # Set the local images to the uploaded images
+        dir_input = st.session_state['dir_uploaded_images']
+
+        n_images = len([f for f in os.listdir(dir_input) if os.path.isfile(os.path.join(dir_input, f))])
+        st.session_state['processing_add_on'] = f" {n_images} Images"
+        uploaded_files = None
+        st.session_state['uploader_idk'] += 1
+        st.info(f"Processing **{n_images}** images from {dir_input}")
+        
+
+    # dir_input = st.text_input("Input directory for images:", value=os.path.join(st.session_state['dir_home'],"demo"))
+    dir_output = os.path.join(st.session_state['dir_home'],"demo_out") # st.text_input("Output directory:", value=os.path.join(st.session_state['dir_home'],"demo_out"))
     
     directory_manager = DirectoryManager(dir_output)
     directory_manager.create_directories()
@@ -304,6 +386,7 @@ def main():
                             # Remove processed image from the list
                             st.session_state.input_list.remove(selected_img)
                             st.rerun()
+                        create_download_button(dir_output, run_name)
                 else:
                     with R_save:
                         if st.button('Save as Failure'):
@@ -323,5 +406,9 @@ def main():
                             st.rerun()
 
 st.set_page_config(layout="wide", page_title='GreenSight')
+
+if 'dir_home' not in st.session_state:
+    st.session_state['dir_home'] = os.path.dirname(__file__)
+
 st.title("GreenSight")
 main()