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

# Function to update displayed quadrilateral based on selected index
def update_displayed_quadrilateral(valid_quadrilaterals, index):
    # Extract the four points of the current quadrilateral
    centroids = get_points_from_contours(valid_quadrilaterals[index])
    return centroids
    
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):
    

    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]

    # Ensure session_state variables are initialized
    if 'valid_quadrilaterals' not in st.session_state:
        st.session_state.valid_quadrilaterals = [perm for perm in itertools.permutations(significant_contours, 4) if is_valid_quadrilateral(get_points_from_contours(perm))]
    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'])

        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)

    # 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, 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

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)

            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()
            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')

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