update

app.py

@@ -89,97 +89,33 @@ def get_points_from_contours(contours):
     return centroids
 
 # Function to display the image with the selected quadrilateral superimposed
-# def display_image_with_quadrilateral(image, points):
-#     # Make a copy of the image to draw on
-#     overlay_image = image.copy()
+def display_image_with_quadrilateral(image, points):
+    # Make a copy of the image to draw on
+    overlay_image = image.copy()
     
-#     # Draw the quadrilateral
-#     cv2.polylines(overlay_image, [np.array(points)], isClosed=True, color=(0, 255, 0), thickness=3)
+    # Draw the quadrilateral
+    cv2.polylines(overlay_image, [np.array(points)], isClosed=True, color=(0, 255, 0), thickness=3)
     
-#     # Display the image with the quadrilateral
-#     st.image(overlay_image, caption="Quadrilateral on Image", use_column_width='auto')
-
-# # Function to update displayed quadrilateral based on selected index
-# def update_displayed_quadrilateral(index, point_combinations, base_image_path):
-#     # Extract the four points of the current quadrilateral
-#     quad_points = get_points_from_contours(point_combinations[index])
-    
-#     # Read the base image
-#     base_image = cv2.imread(base_image_path)
-    
-#     # If the image is not found, handle the error appropriately
-#     if base_image is None:
-#         st.error("Failed to load image.")
-#         return
-    
-#     # Display the image with the selected quadrilateral
-#     display_image_with_quadrilateral(base_image, quad_points)
-
-def increment_index():
-    st.session_state.selected_quad_index = (st.session_state.selected_quad_index + 1) % len(st.session_state.valid_quadrilaterals)
-    st.session_state.centroids = update_displayed_quadrilateral(st.session_state.selected_quad_index)
-
-def decrement_index():
-    st.session_state.selected_quad_index = (st.session_state.selected_quad_index - 1) % len(st.session_state.valid_quadrilaterals)
-    st.session_state.centroids = update_displayed_quadrilateral(st.session_state.selected_quad_index)
+    # Display the image with the quadrilateral
+    st.image(overlay_image, caption="Quadrilateral on Image", use_column_width='auto')
 
 # Function to update displayed quadrilateral based on selected index
-def update_displayed_quadrilateral(valid_quadrilaterals, index):
+def update_displayed_quadrilateral(index, point_combinations, base_image_path):
     # Extract the four points of the current quadrilateral
-    centroids = get_points_from_contours(valid_quadrilaterals[index])
-    return centroids
+    quad_points = get_points_from_contours(point_combinations[index])
     
-def get_centroid(contour):
-    # Compute the centroid for the contour
-    M = cv2.moments(contour)
-    if M["m00"] != 0:
-        return (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
-    return None
-
-def get_points_from_contours(contours):
-    centroids = [get_centroid(contour) for contour in contours if get_centroid(contour) is not None]
-    return centroids
-
-# Function to check if a quadrilateral is valid (i.e., no sides intersect)
-def is_valid_quadrilateral(pts):
-    def ccw(A, B, C):
-        return (C[1] - A[1]) * (B[0] - A[0]) > (B[1] - A[1]) * (C[0] - A[0])
-
-    def intersect(A, B, C, D):
-        return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D)
-
-    A, B, C, D = pts
-    return not (intersect(A, B, C, D) or intersect(A, D, B, C))
-
-def create_polygon_mask(centroids, flag_mask_shape):
-    # Create a polygon mask using the sorted centroids
-    poly_mask = np.zeros(flag_mask_shape, dtype=np.uint8)
-    cv2.fillPoly(poly_mask, [np.array(centroids)], 255)
-    return poly_mask
-
-def warp_and_display_images(img, centroids, base_name, flag_mask_rgb, plant_mask_rgb, mask_plant_plot_rgb):
-    # Warp the images
-    plant_rgb_warp = warp_image(img, centroids)
-    plant_mask_warp = warp_image(mask_plant_plot_rgb, centroids)
-
-    # Extract the RGB pixels from the original image using the mask_plant_plot
-    plant_rgb = cv2.bitwise_and(img, img, mask=create_polygon_mask(centroids, img.shape[:2]))
-
-    # Draw the bounding quadrilateral
-    plot_rgb = plant_rgb.copy()
-    for i in range(4):
-        cv2.line(plot_rgb, centroids[i], centroids[(i+1)%4], (0, 0, 255), 3)
-
-    # Display the images
-    st.image([flag_mask_rgb, plant_mask_rgb, mask_plant_plot_rgb, plot_rgb, plant_rgb_warp, plant_mask_warp], 
-             caption=["Flag Mask", "Plant Mask", "Mask Plant Plot", "Plot RGB", "Plant RGB Warp", "Plant Mask Warp"],
-             use_column_width=True)
-
-    return plant_rgb, plot_rgb, plant_rgb_warp, plant_mask_warp
-
-def process_image(image_path, flag_lower, flag_upper, plant_lower, plant_upper, loc):
+    # Read the base image
+    base_image = cv2.imread(base_image_path)
+    
+    # If the image is not found, handle the error appropriately
+    if base_image is None:
+        st.error("Failed to load image.")
+        return
     
+    # Display the image with the selected quadrilateral
+    display_image_with_quadrilateral(base_image, quad_points)
 
+def process_image(image_path, flag_lower, flag_upper, plant_lower, plant_upper, loc, file_name, file_exists, selected_img, headers, base_name):
     with loc:
         btn_back, btn_next = st.columns([2,2])
 
@@ -212,66 +148,127 @@ 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]
 
-    # 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.")
+    # Logic to handle cases where there are more than 4 significant contours
+    centroids = []
+    if len(significant_contours) < 4:
         return None, None, None, None, None, None, None, None, None, None
+    elif len(significant_contours) > 4:
+        while not st.session_state['keep_quad']:
+            with loc:
+                st.warning("Cycle until correct plot bounds are found")
+            # Create all possible combinations of four points
+            point_combinations = list(itertools.combinations(significant_contours, 4))
+            
+            # Placeholder for quadrilateral indices
+            selected_quad_index = 0
+
+            # Function to update displayed quadrilateral based on selected index
+            def update_displayed_quadrilateral(index):
+                # Extract the four points of the current quadrilateral
+                centroids = get_points_from_contours(point_combinations[index])
+                return centroids
+
+            # Show initial quadrilateral
+            centroids = update_displayed_quadrilateral(selected_quad_index)
+
+            with btn_back:
+                # Button to go to the previous quadrilateral
+                if st.button('Previous'):
+                    selected_quad_index = max(selected_quad_index - 1, 0)
+                    centroids = update_displayed_quadrilateral(selected_quad_index)
+
+            with btn_next:
+                # Button to go to the next quadrilateral
+                if st.button('Next'):
+                    selected_quad_index = min(selected_quad_index + 1, len(point_combinations) - 1)
+                    centroids = update_displayed_quadrilateral(selected_quad_index)
+            
+            with loc:
+                if st.button('Keep Plot Bounds'):
+                    st.session_state['keep_quad'] = True
+                if st.button('Save as Failure'):
+                    st.session_state['keep_quad'] = True
+                    # 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 there are exactly 4 largest contours, proceed with existing logic
+    elif len(significant_contours) == 4:
+        # Create a new mask with only the largest 4 contours
+        largest_4_flag_mask = np.zeros_like(flag_mask)
+        cv2.drawContours(largest_4_flag_mask, sorted_contours, -1, (255), thickness=cv2.FILLED)
+        
+        # Compute the centroid for each contour
+        for contour in sorted_contours:
+            M = cv2.moments(contour)
+            if M["m00"] != 0:
+                cx = int(M["m10"] / M["m00"])
+                cy = int(M["m01"] / M["m00"])
+            else:
+                cx, cy = 0, 0
+            centroids.append((cx, cy))
+    
+        # Compute the centroid of the centroids
+        centroid_x = sum(x for x, y in centroids) / 4
+        centroid_y = sum(y for x, y in centroids) / 4
 
-    # 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[st.session_state.selected_quad_index])
+        # Sort the centroids
+        centroids.sort(key=lambda point: (-math.atan2(point[1] - centroid_y, point[0] - centroid_x)) % (2 * np.pi))
 
-    # 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
+    # Create a polygon mask using the sorted centroids
+    poly_mask = np.zeros_like(flag_mask)
+    cv2.fillPoly(poly_mask, [np.array(centroids)], 255)
+    
+    # Mask the plant_mask with poly_mask
+    mask_plant_plot = cv2.bitwise_and(plant_mask, plant_mask, mask=poly_mask)
 
-        with btn_next:
-            if st.button('Next'):
-                increment_index()  # Call function to increment index
+    # 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)
 
-        poly_mask = create_polygon_mask(st.session_state.centroids, flag_mask.shape)
+    # Draw the bounding quadrilateral
+    plot_rgb = plant_rgb.copy()
+    for i in range(4):
+        cv2.line(plot_rgb, centroids[i], centroids[(i+1)%4], (0, 0, 255), 3)
 
-        # Mask the plant_mask with poly_mask
-        mask_plant_plot = cv2.bitwise_and(plant_mask, plant_mask, mask=poly_mask)
+    # 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
 
-        # 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)
+    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, centroids)
+    plant_mask_warp = warp_image(mask_plant_plot_rgb, 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
 
@@ -464,7 +461,7 @@ def main():
 
                 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, R_save)
+            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, R_save, file_name, file_exists, selected_img, headers, base_name)
 
             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)
@@ -558,9 +555,6 @@ if 'dir_uploaded_images_small' not in st.session_state:
     st.session_state['dir_uploaded_images_small'] = os.path.join(st.session_state['dir_home'],'uploads_small')
     validate_dir(os.path.join(st.session_state['dir_home'],'uploads_small'))
 
-if 'selected_quad_index' not in st.session_state:
-    st.session_state['selected_quad_index'] = 0
-
 
 st.title("GreenSight")
 st.write("Simple color segmentation app to estimate the vegetation coverage in a plot. Corners of the plot need to be marked with solid, uniforly colored flags.")