app.py

import os, math, csv, shutil, itertools
import streamlit as st
from streamlit_image_select import image_select
import cv2
import numpy as np
from PIL import Image
import matplotlib.colors as mcolors
from io import BytesIO


MAX_GALLERY_IMAGES = 50
GALLERY_IMAGE_SIZE = 128
MIN_AREA = 10

class DirectoryManager:
    def __init__(self, output_dir):
        self.dir_output = output_dir
        self.mask_flag = os.path.join(output_dir, "mask_flag")
        self.mask_plant = os.path.join(output_dir, "mask_plant")
        self.mask_plant_plot = os.path.join(output_dir, "mask_plant_plot")
        self.plant_rgb = os.path.join(output_dir, "plant_rgb")
        self.plot_rgb = os.path.join(output_dir, "plot_rgb")
        self.plant_rgb_warp = os.path.join(output_dir, "plant_rgb_warp")
        self.plant_mask_warp = os.path.join(output_dir, "plant_mask_warp")
        self.data = os.path.join(output_dir, "data")

    def create_directories(self):
        os.makedirs(self.dir_output, exist_ok=True)
        os.makedirs(self.mask_flag, exist_ok=True)
        os.makedirs(self.mask_plant, exist_ok=True)
        os.makedirs(self.mask_plant_plot, exist_ok=True)
        os.makedirs(self.plant_rgb, exist_ok=True)
        os.makedirs(self.plot_rgb, exist_ok=True)
        os.makedirs(self.plant_rgb_warp, exist_ok=True)
        os.makedirs(self.plant_mask_warp, exist_ok=True)
        os.makedirs(self.data, exist_ok=True)



def hex_to_hsv_bounds(hex_color, sat_value, val_value):
    # Convert RGB hex to color
    rgb_color = mcolors.hex2color(hex_color)
    hsv_color = mcolors.rgb_to_hsv(np.array(rgb_color).reshape(1, 1, 3))
    
    # Adjust the saturation and value components based on user's input
    hsv_color[0][0][1] = sat_value / 255.0  # Saturation
    hsv_color[0][0][2] = val_value / 255.0  # Value

    hsv_bound = tuple((hsv_color * np.array([179, 255, 255])).astype(int)[0][0])
    
    return hsv_bound

def warp_image(img, vertices):
    # Compute distances between the vertices to determine the size of the target square
    distances = [np.linalg.norm(np.array(vertices[i]) - np.array(vertices[i+1])) for i in range(len(vertices)-1)]
    distances.append(np.linalg.norm(np.array(vertices[-1]) - np.array(vertices[0])))  # Add the distance between the last and first point
    max_distance = max(distances)

    # Define target vertices for the square
    dst_vertices = np.array([
        [max_distance - 1, 0],
        [0, 0],
        [0, max_distance - 1],
        [max_distance - 1, max_distance - 1]
    ], dtype="float32")

    # Compute the perspective transform matrix using the provided vertices
    matrix = cv2.getPerspectiveTransform(np.array(vertices, dtype="float32"), dst_vertices)
    
    # Warp the image to the square
    warped_img = cv2.warpPerspective(img, matrix, (int(max_distance), int(max_distance)))

    return warped_img

# Assuming get_points_from_contours is a function that takes a tuple of four contours
# and returns their respective centroid points as a list of tuples [(x1,y1), (x2,y2), (x3,y3), (x4,y4)]
def get_points_from_contours(contours):
    centroids = []
    for contour in contours:
        # Compute the centroid for the contour
        M = cv2.moments(contour)
        if M["m00"] != 0:
            cX = int(M["m10"] / M["m00"])
            cY = int(M["m01"] / M["m00"])
            centroids.append((cX, cY))
        else:
            # If the contour is a single point or line (which should not happen with flags), handle it here
            pass
    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()
    
    # 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 quadrilateral_area(centroids):
    # Assuming centroids are in correct order (A, B, C, D) to form a quadrilateral
    def distance(p1, p2):
        return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
    
    A, B, C, D = centroids
    # Using Bretschneider's formula to calculate area of a quadrilateral
    a = distance(A, B)
    b = distance(B, C)
    c = distance(C, D)
    d = distance(D, A)
    p = (a + b + c + d) / 2  # semi-perimeter
    return math.sqrt((p - a) * (p - b) * (p - c) * (p - d))

def sort_permutations_by_area(valid_permutations):
    # Calculate area for each permutation and return sorted list
    perm_areas = [(perm, quadrilateral_area(get_points_from_contours(perm))) for perm in valid_permutations]
    # Sort by area in descending order (largest first)
    perm_areas.sort(key=lambda x: x[1], reverse=True)
    # Return only the sorted permutations, not the areas
    sorted_permutations = [perm for perm, area in perm_areas]
    return sorted_permutations

def is_valid_quadrilateral(centroids):
    if len(centroids) != 4:
        return False

    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 = centroids
    return not (intersect(A, B, C, D) or intersect(A, D, B, C))

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])

    img = cv2.imread(image_path)
    
    # Check if image is valid
    if img is None:
        print(f"Error reading image from path: {image_path}")
        return None, None, None, None, None, None, None, None, None, None

    hsv_img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)  # Convert image to HSV
    
    # Explicitly ensure bounds are integer tuples
    flag_lower = tuple(int(x) for x in flag_lower)
    flag_upper = tuple(int(x) for x in flag_upper)
    plant_lower = tuple(int(x) for x in plant_lower)
    plant_upper = tuple(int(x) for x in plant_upper)

    flag_mask = cv2.inRange(hsv_img, flag_lower, flag_upper)
    plant_mask = cv2.inRange(hsv_img, plant_lower, plant_upper)

    # # Find contours
    # contours, _ = cv2.findContours(flag_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # # Sort contours by area and keep only the largest 4
    # sorted_contours = sorted(contours, key=cv2.contourArea, reverse=True)[:4]
    
    # # If there are not 4 largest contours, return
    # if len(sorted_contours) != 4:
    #     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]

    # 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:
        st.session_state['keep_quad'] = False
        # 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
        if len(significant_contours) >= 4:
            # Generate all permutations of four points from the significant contours
            permutations_of_four = list(itertools.permutations(significant_contours, 4))

            # Filter out invalid quadrilaterals
            valid_permutations0 = [perm for perm in permutations_of_four if is_valid_quadrilateral(get_points_from_contours(perm))]

            valid_permutations = sort_permutations_by_area(valid_permutations0)

            if not valid_permutations:
                st.error("No valid quadrilaterals found.")
                return None, None, None, None, None, None, None, None, None, None

        # 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(valid_permutations[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'):
                st.session_state.quad_index = (st.session_state.quad_index - 1) % len(valid_permutations)
                centroids = update_displayed_quadrilateral(st.session_state.quad_index)

        with btn_next:
            # Button to go to the next quadrilateral
            if st.button('Next'):
                st.session_state.quad_index = (st.session_state.quad_index + 1) % len(valid_permutations)
                centroids = update_displayed_quadrilateral(st.session_state.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

        # Sort the centroids
        centroids.sort(key=lambda point: (-math.atan2(point[1] - centroid_y, point[0] - centroid_x)) % (2 * np.pi))
    
    if len(centroids) == 4:
        # 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)

        # 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, centroids[i], 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, 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

def calculate_coverage(mask_plant_plot, plant_mask_warp, black_pixels_in_quad):
    # Calculate the percentage of white pixels for mask_plant_plot
    white_pixels_plot = np.sum(mask_plant_plot > 0)
    total_pixels_plot = mask_plant_plot.size
    plot_coverage = (white_pixels_plot / black_pixels_in_quad) * 100

    # Convert plant_mask_warp to grayscale
    plant_mask_warp_gray = cv2.cvtColor(plant_mask_warp, cv2.COLOR_BGR2GRAY)

    # Calculate the percentage of white pixels for plant_mask_warp
    white_pixels_warp = np.sum(plant_mask_warp_gray > 0)
    total_pixels_warp = plant_mask_warp_gray.size
    warp_coverage = (white_pixels_warp / total_pixels_warp) * 100

    # Calculate the area in cm^2 of the mask_plant_plot
    # Given that the real-life size of the square is 2 square meters or 20000 cm^2
    plot_area_cm2 = (white_pixels_warp / total_pixels_warp) * 20000

    return round(plot_coverage,2), round(warp_coverage,2), round(plot_area_cm2,2)

def get_color_parameters():
    # Color pickers for hue component
    FL, FL_S, FL_SS = st.columns([2,4,4])
    with FL:
        flag_lower_hex = st.color_picker("Flag Color Lower Bound Hue", "#33211f")
    with FL_S:
        flag_lower_sat = st.slider("Flag Lower Bound Saturation", 0, 255, 120)
    with FL_SS:
        flag_lower_val = st.slider("Flag Lower Bound Value", 0, 255, 150)

    FU, FU_S, FU_SS = st.columns([2,4,4])
    with FU:
        flag_upper_hex = st.color_picker("Flag Color Upper Bound Hue", "#ff7700")
    with FU_S:
        flag_upper_sat = st.slider("Flag Upper Bound Saturation", 0, 255, 255)
    with FU_SS:
        flag_upper_val = st.slider("Flag Upper Bound Value", 0, 255, 255)

    PL, PL_S, PL_SS = st.columns([2,4,4])
    with PL:
        plant_lower_hex = st.color_picker("Plant Color Lower Bound Hue", "#504F49")
    with PL_S:
        plant_lower_sat = st.slider("Plant Lower Bound Saturation", 0, 255, 30)
    with PL_SS:
        plant_lower_val = st.slider("Plant Lower Bound Value", 0, 255, 30)

    PU, PU_S, PU_SS = st.columns([2,4,4])
    with PU:
        plant_upper_hex = st.color_picker("Plant Color Upper Bound Hue", "#00CFFF")
    with PU_S:
        plant_upper_sat = st.slider("Plant Upper Bound Saturation", 0, 255, 255)
    with PU_SS:
        plant_upper_val = st.slider("Plant Upper Bound Value", 0, 255, 255)  

    # Get HSV bounds using the modified function
    flag_lower_bound = hex_to_hsv_bounds(flag_lower_hex, flag_lower_sat, flag_lower_val)
    flag_upper_bound = hex_to_hsv_bounds(flag_upper_hex, flag_upper_sat, flag_upper_val)
    plant_lower_bound = hex_to_hsv_bounds(plant_lower_hex, plant_lower_sat, plant_lower_val)
    plant_upper_bound = hex_to_hsv_bounds(plant_upper_hex, plant_upper_sat, plant_upper_val)

    return flag_lower_bound, flag_upper_bound, plant_lower_bound, plant_upper_bound

def save_img(directory, base_name, mask):
    mask_name = os.path.join(directory, os.path.basename(base_name))
    cv2.imwrite(mask_name, mask)

def validate_dir(dir):
    if not os.path.exists(dir):
        os.makedirs(dir, exist_ok=True)

def make_zipfile(source_dir, output_filename):
    shutil.make_archive(output_filename, 'zip', source_dir)
    return output_filename + '.zip'

def save_uploaded_file(directory, img_file, image=None):
    if not os.path.exists(directory):
        os.makedirs(directory)
    # Assuming the uploaded file is an image
    if image is None:
        with Image.open(img_file) as image:
            full_path = os.path.join(directory, img_file.name)
            image.save(full_path, "JPEG")
        # Return the full path of the saved image
        return full_path
    else:
        full_path = os.path.join(directory, img_file.name)
        image.save(full_path, "JPEG")
        return full_path

def create_download_button(dir_to_zip, zip_filename):
    zip_filepath = make_zipfile(dir_to_zip, zip_filename)
    with open(zip_filepath, 'rb') as f:
        bytes_io = BytesIO(f.read())
    st.download_button(
        label=f"Download Results for{st.session_state['processing_add_on']}",type='primary',
        data=bytes_io,
        file_name=os.path.basename(zip_filepath),
        mime='application/zip'
    )

def delete_directory(dir_path):
    try:
        shutil.rmtree(dir_path)
        st.session_state['input_list'] = []
        st.session_state['input_list_small'] = []
        # st.success(f"Deleted previously uploaded images, making room for new images: {dir_path}")
    except OSError as e:
        st.error(f"Error: {dir_path} : {e.strerror}")

def clear_image_gallery():
    delete_directory(st.session_state['dir_uploaded_images'])
    delete_directory(st.session_state['dir_uploaded_images_small'])
    validate_dir(st.session_state['dir_uploaded_images'])
    validate_dir(st.session_state['dir_uploaded_images_small'])

def reset_demo_images():
    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'))]
    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"
    st.session_state['uploader_idk'] += 1

def main():
    _, R_coverage, R_plot_area_cm2, R_save = st.columns([5,2,2,2])
    img_gallery, img_main, img_seg, img_green, img_warp = st.columns([1,4,2,2,2])

    st.session_state['dir_uploaded_images'] = os.path.join(st.session_state['dir_home'],'uploads')
    st.session_state['dir_uploaded_images_small'] = os.path.join(st.session_state['dir_home'],'uploads_small')
    uploaded_files = st.file_uploader("Upload Images", type=['jpg', 'jpeg'], accept_multiple_files=True, key=st.session_state['uploader_idk'])
    if uploaded_files:
        # Clear input image gallery and input list
        clear_image_gallery()

        # Process the new iamges
        for uploaded_file in uploaded_files:
            file_path = save_uploaded_file(st.session_state['dir_uploaded_images'], uploaded_file)
            st.session_state['input_list'].append(file_path)

            img = Image.open(file_path)
            img.thumbnail((GALLERY_IMAGE_SIZE, GALLERY_IMAGE_SIZE), Image.Resampling.LANCZOS)  
            file_path_small = save_uploaded_file(st.session_state['dir_uploaded_images_small'], uploaded_file, img)
            st.session_state['input_list_small'].append(file_path_small)
            print(uploaded_file.name)

        # 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 st.session_state['dir_input'] is None:
        reset_demo_images()
    
    # 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)
    st.button("Reset Demo Images", on_click=reset_demo_images)

        
    if len(st.session_state['input_list']) == 0 or st.session_state['input_list'] is None:
        st.balloons()
        create_download_button(dir_output, run_name)

    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)
            create_download_button(dir_output, run_name)
        
        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, 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)

                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:
                    st.write(f"Showing plot outline #{st.session_state.quad_index}")
                    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.session_state['quad_index'] = 0
                        st.rerun()
            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.session_state['quad_index'] = 0
                        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'

if 'uploader_idk' not in st.session_state:
    st.session_state['uploader_idk'] = 1

if 'input_list' not in st.session_state:
    st.session_state['input_list'] = []

if 'input_list_small' not in st.session_state:
    st.session_state['input_list_small'] = []

if 'dir_uploaded_images' not in st.session_state:
    st.session_state['dir_uploaded_images'] = os.path.join(st.session_state['dir_home'],'uploads')
    validate_dir(os.path.join(st.session_state['dir_home'],'uploads'))

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 'keep_quad' not in st.session_state:
    st.session_state['keep_quad'] = False

if 'quad_index' not in st.session_state:
    st.session_state['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.")
st.write("If you exit the session before completing the segmentation of all images, all progress will be lost!")
main()