update

app.py

@@ -306,104 +306,106 @@ def main():
         st.session_state['uploader_idk'] += 1
         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()
+    if not st.session_state['dir_input']:
+        st.session_state['dir_input'] = os.path.join(st.session_state['dir_home'],"demo")
+    
+    # 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(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 '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 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 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 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')