update

app.py

@@ -212,86 +212,66 @@ def process_image(image_path, flag_lower, flag_upper, plant_lower, plant_upper,
     #     return None, None, None, None, None, None, None, None, None, None
 
 
-    # Find contours
     contours, _ = cv2.findContours(flag_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-    # Sort contours by area and keep a significant number, assuming noise has much smaller area
     sorted_contours = sorted(contours, key=cv2.contourArea, reverse=True)
-
-    # Filter out noise based on a predefined area threshold
     significant_contours = [cnt for cnt in sorted_contours if cv2.contourArea(cnt) > MIN_AREA]
 
-    # Ensure session_state variables are initialized
-    if 'valid_quadrilaterals' not in st.session_state:
+    # Generate all permutations of four points and filter valid quadrilaterals
+    if 'valid_quadrilaterals' not in st.session_state or st.session_state.valid_quadrilaterals is None:
         st.session_state.valid_quadrilaterals = [perm for perm in itertools.permutations(significant_contours, 4) if is_valid_quadrilateral(get_points_from_contours(perm))]
+    
+    if not st.session_state.valid_quadrilaterals:
+        st.warning("No valid quadrilaterals found.")
+        return None, None, None, None, None, None, None, None, None, None
+
+    # Initialize or update the selected_quad_index and centroids
     if 'selected_quad_index' not in st.session_state:
         st.session_state.selected_quad_index = 0
     if 'centroids' not in st.session_state:
-        st.session_state.centroids = get_points_from_contours(st.session_state.valid_quadrilaterals[0])
-
-    # Logic to handle cases where there are more than 4 significant contours
-    if len(significant_contours) < 4:
-        with loc:
-            st.info("Not enough points to form a quadrilateral.")
-        return None, None, None, None, None, None, None, None, None, None
-    else:
-        # Generate all permutations of four points
-        permutations_of_four = list(itertools.permutations(significant_contours, 4))
-        valid_quadrilaterals = [perm for perm in permutations_of_four if is_valid_quadrilateral(get_points_from_contours(perm))]
-        
-        if not valid_quadrilaterals:
-            with loc:
-                st.info("No valid quadrilaterals found.")
-            return None, None, None, None, None, None, None, None, None, None
-        
-        # Initial quadrilateral index
-        st.session_state['selected_quad_index'] = 0
-        st.session_state.centroids = get_points_from_contours(valid_quadrilaterals[st.session_state['selected_quad_index']])
-        
-
-        # Show initial quadrilateral
-        st.session_state.centroids = update_displayed_quadrilateral(st.session_state['selected_quad_index'])
+        st.session_state.centroids = get_points_from_contours(st.session_state.valid_quadrilaterals[st.session_state.selected_quad_index])
 
+    # Logic to handle cases where there are 4 or more significant contours
+    if len(significant_contours) >= 4:
         with btn_back:
             if st.button('Previous'):
                 decrement_index()  # Call function to decrement index
 
         with btn_next:
             if st.button('Next'):
-                increment_index()
-
-    poly_mask = create_polygon_mask(st.session_state.centroids, flag_mask.shape)
+                increment_index()  # Call function to increment index
 
-    # Mask the plant_mask with poly_mask
-    mask_plant_plot = cv2.bitwise_and(plant_mask, plant_mask, mask=poly_mask)
+        poly_mask = create_polygon_mask(st.session_state.centroids, flag_mask.shape)
 
-    # Count the number of black pixels inside the quadrilateral
-    total_pixels_in_quad = np.prod(poly_mask.shape)
-    white_pixels_in_quad = np.sum(poly_mask == 255)
-    black_pixels_in_quad = total_pixels_in_quad - white_pixels_in_quad
-    
-    # Extract the RGB pixels from the original image using the mask_plant_plot
-    plant_rgb = cv2.bitwise_and(img, img, mask=mask_plant_plot)
+        # Mask the plant_mask with poly_mask
+        mask_plant_plot = cv2.bitwise_and(plant_mask, plant_mask, mask=poly_mask)
 
-    # Draw the bounding quadrilateral
-    plot_rgb = plant_rgb.copy()
-    for i in range(4):
-        cv2.line(plot_rgb, st.session_state.centroids[i], st.session_state.centroids[(i+1)%4], (0, 0, 255), 3)
-
-    # Convert the masks to RGB for visualization
-    flag_mask_rgb = cv2.cvtColor(flag_mask, cv2.COLOR_GRAY2RGB)
-    orange_color = [255, 165, 0]  # RGB value for orange
-    flag_mask_rgb[np.any(flag_mask_rgb != [0, 0, 0], axis=-1)] = orange_color
-
-    plant_mask_rgb = cv2.cvtColor(plant_mask, cv2.COLOR_GRAY2RGB)
-    mask_plant_plot_rgb = cv2.cvtColor(mask_plant_plot, cv2.COLOR_GRAY2RGB)
-    bright_green_color = [0, 255, 0]
-    plant_mask_rgb[np.any(plant_mask_rgb != [0, 0, 0], axis=-1)] = bright_green_color
-    mask_plant_plot_rgb[np.any(mask_plant_plot_rgb != [0, 0, 0], axis=-1)] = bright_green_color
-    
-    # Warp the images
-    plant_rgb_warp = warp_image(plant_rgb, st.session_state.centroids)
-    plant_mask_warp = warp_image(mask_plant_plot_rgb, st.session_state.centroids)
+        # Count the number of black pixels inside the quadrilateral
+        total_pixels_in_quad = np.prod(poly_mask.shape)
+        white_pixels_in_quad = np.sum(poly_mask == 255)
+        black_pixels_in_quad = total_pixels_in_quad - white_pixels_in_quad
+        
+        # Extract the RGB pixels from the original image using the mask_plant_plot
+        plant_rgb = cv2.bitwise_and(img, img, mask=mask_plant_plot)
+
+        # Draw the bounding quadrilateral
+        plot_rgb = plant_rgb.copy()
+        for i in range(4):
+            cv2.line(plot_rgb, st.session_state.centroids[i], st.session_state.centroids[(i+1)%4], (0, 0, 255), 3)
+
+        # Convert the masks to RGB for visualization
+        flag_mask_rgb = cv2.cvtColor(flag_mask, cv2.COLOR_GRAY2RGB)
+        orange_color = [255, 165, 0]  # RGB value for orange
+        flag_mask_rgb[np.any(flag_mask_rgb != [0, 0, 0], axis=-1)] = orange_color
+
+        plant_mask_rgb = cv2.cvtColor(plant_mask, cv2.COLOR_GRAY2RGB)
+        mask_plant_plot_rgb = cv2.cvtColor(mask_plant_plot, cv2.COLOR_GRAY2RGB)
+        bright_green_color = [0, 255, 0]
+        plant_mask_rgb[np.any(plant_mask_rgb != [0, 0, 0], axis=-1)] = bright_green_color
+        mask_plant_plot_rgb[np.any(mask_plant_plot_rgb != [0, 0, 0], axis=-1)] = bright_green_color
+        
+        # Warp the images
+        plant_rgb_warp = warp_image(plant_rgb, st.session_state.centroids)
+        plant_mask_warp = warp_image(mask_plant_plot_rgb, st.session_state.centroids)
 
     return flag_mask_rgb, plant_mask_rgb, mask_plant_plot_rgb, plant_rgb, plot_rgb, plant_rgb_warp, plant_mask_warp, plant_mask, mask_plant_plot, black_pixels_in_quad