update

app.py

@@ -300,14 +300,17 @@ def main():
         # Set the local images to the uploaded images
         st.session_state['dir_input'] = st.session_state['dir_uploaded_images']
 
+        st.session_state['input_list'] = [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'))]
+
         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 {st.session_state['dir_input']}")
         
-    if not st.session_state['dir_input']:
+    if st.session_state['dir_input'] is None:
         st.session_state['dir_input'] = os.path.join(st.session_state['dir_home'],"demo")
+        st.session_state['input_list'] = [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'))]
     
     # 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"))
@@ -320,92 +323,88 @@ def main():
     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(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(st.session_state['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 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')