main.py

from detectree2model.predictions.predict import run_detectree2
from polygons_processing.postpprocess_detectree2 import postprocess
from generate_tree_images.generate_tree_images import generate_tree_images
from classification.classification_predict import classify
import os
import json
import gradio as gr
import rasterio
import geopandas as gpd
import plotly.graph_objects as go
import numpy as np
import ast

def row_to_feature(row):
    if (row['geometry'].geom_type == 'Polygon'):
        coordinates = row['geometry'].exterior.coords.xy
        coordinate_list = [[x, y] for x, y in zip(coordinates[0], coordinates[1])]
        feature = {
        "id": row["id"],
        "type": "Feature",
        "properties": {"Confidence_score": row["Confidence_score"], "species": row['species']},
        "geometry": {"type": "Polygon", "coordinates": [coordinate_list]},
        }
        return feature

 
   

def export_geojson(df, filename):
    
    features = [row_to_feature(row) for idx, row in df.iterrows()]

    feature_collection = {
        "type": "FeatureCollection",
        "crs": {"type": "name", "properties": {"name": "urn:ogc:def:crs:EPSG::32720"}},
        "features": features,
    } 

    output_geojson = json.dumps(feature_collection)

    with open(f"{filename}.geojson", "w") as f:
        f.write(output_geojson)

    print(f"GeoJSON data exported to '{filename}.geojson' file.")
    
def delete_files_in_directory(directory_path: str) -> None:
    """
    Delete all files in the given directory.

    :param directory_path: Path to the directory
    """
    if not os.path.exists(directory_path):
        return

    if not os.path.isdir(directory_path):
        raise ValueError(f"The provided path '{directory_path}' is not a valid directory.")

    for filename in os.listdir(directory_path):
        file_path = os.path.join(directory_path, filename)
        if os.path.isfile(file_path):
            os.remove(file_path)
            print(f"Deleted file: {file_path}")
        else:
            print(f"Skipped non-file: {file_path}")

def process_image(image_path: str, progress: gr.Progress):
    current_directory = os.getcwd()

    output_directory = os.path.join(current_directory, "outputs")
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    progress(0, desc="Deleting files")
    delete_files_in_directory("tiles")
    delete_files_in_directory("tiles/predictions")
    delete_files_in_directory("tiles/predictions_geo")
    delete_files_in_directory("tree_images")
    delete_files_in_directory("detected_trees")
    progress(0.01, desc="Running detectree2")



    run_detectree2(image_path, store_path=output_directory)
    progress(0.2, desc="Ran detectree2")


    processed_output_df = postprocess(output_directory + '/detectree2_delin.geojson', output_directory + '/processed_delin.geojson')
    progress(0.4, desc="Post processed")

    processed_geojson = output_directory + '/processed_delin.geojson'

    progress(0.6, desc="Generating tree images")
    generate_tree_images(processed_geojson, image_path)
    progress(0.7, desc="Generated tree images")

    output_folder = './tree_images'

    processed_output_df = gpd.read_file(processed_geojson)

    all_top_3_list = []  # Initialize an empty list to accumulate all top_3 lists

    for file_name in os.listdir(output_folder):
        file_path = os.path.join(output_folder, file_name)

        probs = classify(file_path)
        top_3 = probs.head(3)
        top_3_list = [[cls, prob] for cls, prob in top_3.items()]

        # Accumulate the top_3_list for each file
        all_top_3_list.append(top_3_list)

    # Assign the accumulated top_3_list to the 'species' column of the dataframe
    processed_output_df['species'] = all_top_3_list

    print(processed_output_df.head())
    print(processed_output_df.columns)

    final_output_path = 'result'

    progress(0.8, desc="Exporting geojson")
    export_geojson(processed_output_df, final_output_path)
    progress(0.9, desc="Exported geojson")

    return final_output_path, image_path

def plot_results(geojson_path, image_path):
    with rasterio.open(image_path) as src:
        tif_image = src.read([1, 2, 3])  # Read the first three bands (RGB)
        tif_transform = src.transform
        height, width = tif_image.shape[1], tif_image.shape[2]

    # Read the GeoJSON file
    geojson_data = gpd.read_file(geojson_path + '.geojson')

    # Create Plotly figure
    fig = go.Figure()

    # Add image to the figure
    fig.add_trace(go.Image(z=tif_image.transpose((1, 2, 0)), hoverinfo = 'none'))

    # Add polygons to the plot
    for idx, row in geojson_data.iterrows():
        coordinates = row['geometry'].exterior.coords.xy
        x, y = list(coordinates[0]), list(coordinates[1])  # Convert to list

        # Transform coordinates to match image pixel space
        x_transformed = [(xi - tif_transform.c) / tif_transform.a for xi in x]
        y_transformed = [(yi - tif_transform.f) / tif_transform.e for yi in y]

        species_info_str = row['species']
        species_info = ast.literal_eval(species_info_str)
        first_array = species_info[0]
        second_array = species_info[1]
        third_array = species_info[2]
        confidence_score = row['Confidence_score']
        hovertemplate = f"Polygon:<br>{idx}<br><br>Species and Probability:<br>{first_array}<br>{second_array}<br>{third_array}<br><br>Confidence:<br>{confidence_score}"

        fig.add_trace(go.Scatter(
            x=x_transformed, 
            y=y_transformed, 
            mode='lines',  
            name = '',
            line=dict(color='red'),
            hovertemplate=hovertemplate,
            hoverinfo='text'
        ))

    fig.update_layout(
        title='TIF Image with Tree Crowns Overlay', 
        xaxis_title='X',
        yaxis_title='Y',
        showlegend=False, # Hide the legend
    )

    return fig



def greet(image_path: str, progress=gr.Progress()):
    geojson_path, image_path = process_image(image_path, progress)
    progress(0.0, desc="Plotting results")
    fig = plot_results (geojson_path, image_path)
    progress(1, desc="Plotted results")
    
    return fig
    
#tif_file_name = "TreeCrownVectorDataset_761588_9673769_20_20_32720.tif"
#tif_input = "/Users/jonathanseele/ETH/Hackathons/EcoHackathon/WeCanopy/test/" + tif_file_name

# File paths
#tif_file_path = '/Users/taekim/ecohackathon/WeCanopy/test/TreeCrownVectorDataset_761588_9673769_20_20_32720.tif'
#geojson_file_path = '/Users/taekim/ecohackathon/WeCanopy/test/result.geojson'

# Read the TIF file


def main():
    demo = gr.Interface(
        fn=greet,
        inputs=gr.File(type='filepath'),
        outputs=gr.Plot(label="Tree Crowns"),
    )

    demo.queue()
    demo.launch(server_name='0.0.0.0', debug=True, show_error=True)

if __name__ == "__main__":
    main()