update

app.py

@@ -298,118 +298,121 @@ def main():
             print(uploaded_file.name)
 
         # Set the local images to the uploaded images
-        dir_input = st.session_state['dir_uploaded_images']
+        st.session_state['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))])
+        n_images = len([f for f in os.listdir(st.session_state['dir_input']) if os.path.isfile(os.path.join(st.session_state['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}")
+        st.info(f"Processing **{n_images}** images from {st.session_state['dir_input']}")
         
+    if st.session_state['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()
 
-    # 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()
-
-    run_name = st.text_input("Run name:", value="test")
-    file_name = os.path.join(directory_manager.data, f"{run_name}.csv")
-    headers = ['image',"plant_coverage_uncorrected_percen", "plant_coverage_corrected_percent", "plant_area_corrected_cm2"]
-    file_exists = os.path.isfile(file_name)
+        run_name = st.text_input("Run name:", value="test")
+        file_name = os.path.join(directory_manager.data, f"{run_name}.csv")
+        headers = ['image',"plant_coverage_uncorrected_percen", "plant_coverage_corrected_percent", "plant_area_corrected_cm2"]
+        file_exists = os.path.isfile(file_name)
 
-    if 'input_list' not in st.session_state:
-        input_images = [os.path.join(dir_input, fname) for fname in os.listdir(dir_input) if fname.endswith(('.jpg', '.jpeg', '.png'))]
-        st.session_state.input_list = input_images
+        if 'input_list' not in st.session_state:
+            input_images = [os.path.join(st.session_state['dir_input'], fname) for fname in os.listdir(st.session_state['dir_input']) if fname.endswith(('.jpg', '.jpeg', '.png'))]
+            st.session_state.input_list = input_images
 
-    if os.path.exists(dir_input):
-        
-        if len(st.session_state.input_list) == 0 or st.session_state.input_list is None:
-            st.balloons()
-        else:
-            with img_gallery:
-                selected_img = image_select("Select an image", st.session_state.input_list, use_container_width=False)
-                base_name = os.path.basename(selected_img)
+        if os.path.exists(st.session_state['dir_input']):
             
-            if selected_img:
-
-                selected_img_view = Image.open(selected_img)
-                with img_main:
-                    st.image(selected_img_view, caption="Selected Image", use_column_width='auto')
-
-                    flag_lower_bound, flag_upper_bound, plant_lower_bound, plant_upper_bound = get_color_parameters()
-
-                flag_mask, plant_mask, mask_plant_plot, plant_rgb, plot_rgb, plant_rgb_warp, plant_mask_warp, plant_mask_bi, mask_plant_plot_bi, black_pixels_in_quad = process_image(selected_img, flag_lower_bound, flag_upper_bound, plant_lower_bound, plant_upper_bound)
-
-                if plant_mask_warp is not None:
-                    plot_coverage, warp_coverage, plot_area_cm2 = calculate_coverage(mask_plant_plot_bi, plant_mask_warp, black_pixels_in_quad)
-
-                    with R_coverage:
-                        st.markdown(f"Uncorrected Plant Coverage: {plot_coverage}%")
-                    with R_plot_area_cm2:
-                        st.markdown(f"Corrected Plant Coverage: {warp_coverage}%")
-                        st.markdown(f"Corrected Plant Area: {plot_area_cm2}cm2")
-
-                    # Display masks in galleries
-                    with img_seg:
-                        st.image(plant_mask, caption="Plant Mask", use_column_width=True)
-                        st.image(flag_mask, caption="Flag Mask", use_column_width=True)
-                    with img_green:
-                        st.image(mask_plant_plot, caption="Plant Mask Inside Plot", use_column_width=True)
-                        st.image(plant_rgb, caption="Plant Material", use_column_width=True)
-                    with img_warp:
-                        st.image(plot_rgb, caption="Plant Material Inside Plot", use_column_width=True)
-                        st.image(plant_rgb_warp, caption="Plant Mask Inside Plot Warped to Square", use_column_width=True)
-                        # st.image(plot_rgb_warp, caption="Flag Mask", use_column_width=True)
-                    with R_save:
-                        if st.button('Save'):
-                            # Save the masks to their respective folders
-                            save_img(directory_manager.mask_flag, base_name, flag_mask)
-                            save_img(directory_manager.mask_plant, base_name, plant_mask)
-                            save_img(directory_manager.mask_plant_plot, base_name, mask_plant_plot)
-                            save_img(directory_manager.plant_rgb, base_name, plant_rgb)
-                            save_img(directory_manager.plot_rgb, base_name, plot_rgb)
-                            save_img(directory_manager.plant_rgb_warp, base_name, plant_rgb_warp)
-                            save_img(directory_manager.plant_mask_warp, base_name, plant_mask_warp)
-
-                            # Append the data to the CSV file
-                            with open(file_name, mode='a', newline='') as file:
-                                writer = csv.writer(file)
-                                
-                                # If the file doesn't exist, write the headers
-                                if not file_exists:
-                                    writer.writerow(headers)
-                                
-                                # Write the data
-                                writer.writerow([f"{base_name}",f"{plot_coverage}", f"{warp_coverage}", f"{plot_area_cm2}"])
-
-                            # 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'):
-                            # Append the data to the CSV file
-                            with open(file_name, mode='a', newline='') as file:
-                                writer = csv.writer(file)
-                                
-                                # If the file doesn't exist, write the headers
-                                if not file_exists:
-                                    writer.writerow(headers)
-                                
-                                # Write the data
-                                writer.writerow([f"{base_name}",f"NA", f"NA", f"NA"])
-
-                            # Remove processed image from the list
-                            st.session_state.input_list.remove(selected_img)
-                            st.rerun()
+            if len(st.session_state.input_list) == 0 or st.session_state.input_list is None:
+                st.balloons()
+            else:
+                with img_gallery:
+                    selected_img = image_select("Select an image", st.session_state.input_list, use_container_width=False)
+                    base_name = os.path.basename(selected_img)
+                
+                if selected_img:
+
+                    selected_img_view = Image.open(selected_img)
+                    with img_main:
+                        st.image(selected_img_view, caption="Selected Image", use_column_width='auto')
+
+                        flag_lower_bound, flag_upper_bound, plant_lower_bound, plant_upper_bound = get_color_parameters()
+
+                    flag_mask, plant_mask, mask_plant_plot, plant_rgb, plot_rgb, plant_rgb_warp, plant_mask_warp, plant_mask_bi, mask_plant_plot_bi, black_pixels_in_quad = process_image(selected_img, flag_lower_bound, flag_upper_bound, plant_lower_bound, plant_upper_bound)
+
+                    if plant_mask_warp is not None:
+                        plot_coverage, warp_coverage, plot_area_cm2 = calculate_coverage(mask_plant_plot_bi, plant_mask_warp, black_pixels_in_quad)
+
+                        with R_coverage:
+                            st.markdown(f"Uncorrected Plant Coverage: {plot_coverage}%")
+                        with R_plot_area_cm2:
+                            st.markdown(f"Corrected Plant Coverage: {warp_coverage}%")
+                            st.markdown(f"Corrected Plant Area: {plot_area_cm2}cm2")
+
+                        # Display masks in galleries
+                        with img_seg:
+                            st.image(plant_mask, caption="Plant Mask", use_column_width=True)
+                            st.image(flag_mask, caption="Flag Mask", use_column_width=True)
+                        with img_green:
+                            st.image(mask_plant_plot, caption="Plant Mask Inside Plot", use_column_width=True)
+                            st.image(plant_rgb, caption="Plant Material", use_column_width=True)
+                        with img_warp:
+                            st.image(plot_rgb, caption="Plant Material Inside Plot", use_column_width=True)
+                            st.image(plant_rgb_warp, caption="Plant Mask Inside Plot Warped to Square", use_column_width=True)
+                            # st.image(plot_rgb_warp, caption="Flag Mask", use_column_width=True)
+                        with R_save:
+                            if st.button('Save'):
+                                # Save the masks to their respective folders
+                                save_img(directory_manager.mask_flag, base_name, flag_mask)
+                                save_img(directory_manager.mask_plant, base_name, plant_mask)
+                                save_img(directory_manager.mask_plant_plot, base_name, mask_plant_plot)
+                                save_img(directory_manager.plant_rgb, base_name, plant_rgb)
+                                save_img(directory_manager.plot_rgb, base_name, plot_rgb)
+                                save_img(directory_manager.plant_rgb_warp, base_name, plant_rgb_warp)
+                                save_img(directory_manager.plant_mask_warp, base_name, plant_mask_warp)
+
+                                # Append the data to the CSV file
+                                with open(file_name, mode='a', newline='') as file:
+                                    writer = csv.writer(file)
+                                    
+                                    # If the file doesn't exist, write the headers
+                                    if not file_exists:
+                                        writer.writerow(headers)
+                                    
+                                    # Write the data
+                                    writer.writerow([f"{base_name}",f"{plot_coverage}", f"{warp_coverage}", f"{plot_area_cm2}"])
+
+                                # 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'):
+                                # Append the data to the CSV file
+                                with open(file_name, mode='a', newline='') as file:
+                                    writer = csv.writer(file)
+                                    
+                                    # If the file doesn't exist, write the headers
+                                    if not file_exists:
+                                        writer.writerow(headers)
+                                    
+                                    # Write the data
+                                    writer.writerow([f"{base_name}",f"NA", f"NA", f"NA"])
+
+                                # Remove processed image from the list
+                                st.session_state.input_list.remove(selected_img)
+                                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__)
 
+if 'dir_input' not in st.session_state:
+    st.session_state['dir_input'] = None
+
 if 'processing_add_on' not in st.session_state:
     st.session_state['processing_add_on'] = ' 1 Image'