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
from rasterio.plot import show
import geopandas as gpd
import matplotlib.pyplot as plt
from shapely.geometry import Point




def row_to_feature(row):
    feature = {
        "id": row["id"],
        "type": "Feature",
        "properties": {"Confidence_score": row["Confidence_score"], "species": row['species']},
        "geometry": {"type": "Polygon", "coordinates": [row["coordinates"]]},
        
    }
    return feature
# [[first guess, prob],[second guess, prob],[third guess, prob]] 
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.")

"""
tif_input: the file containing a tif that we are analyzing

tif_file_name: the file name of the tif input. tif_input is the folder in which the tif file lies 
(detectree2 works with that) but generate_tree_images requires path including the file hence the file name is needed

output_directory: the directory were all in-between and final files are stored

generate_tree_images stores the cutout tree images in a separate folder

"""


def greet(image_path: str):
    current_directory = os.getcwd()

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

    run_detectree2(image_path, store_path=output_directory)

    processed_output_df = postprocess(output_directory + '/detectree2_delin.geojson', output_directory + '/processed_delin')

    processed_geojson = output_directory + '/processed_delin.geojson'

    generate_tree_images(processed_geojson, image_path)

    output_folder = './tree_images'

    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

    final_output_path = 'result'
    export_geojson(processed_output_df, final_output_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

    # Read the GeoJSON file
    geojson_data = gpd.read_file(final_output_path)

    # Set the interactive backend to Qt5Agg
    plt.switch_backend('Qt5Agg') # You have to install PyQt5

    # Enable interactive mode
    plt.ion()

    # Plotting
    fig, ax = plt.subplots(figsize=(10, 10))

    # Plot the RGB TIF image
    show(tif_image, transform=tif_transform, ax=ax)

    # Plot the GeoJSON polygons
    geojson_data.plot(ax=ax, facecolor='none', edgecolor='red')

    # Set plot title
    ax.set_title('TIF Image with Tree Crowns Overlay')

    # Create an annotation box
    annot = ax.annotate("", xy=(0, 0), xytext=(20, 20),
                    textcoords="offset points",
                    bbox=dict(boxstyle="round", fc="w"),
                    arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)

    # Create a function to handle mouse clicks
    def on_click(event):
        if event.inaxes is not None:
            # Get the coordinates of the click
            click_point = Point(event.xdata, event.ydata)
            # Check if the click is within any of the polygons
            for idx, row in geojson_data.iterrows():
                if row['geometry'].contains(click_point):
                    # Access the properties dictionarㅋ
                    # Extract species and confidence score
                    species_info = row['species']
                    confidence_score = row['Confidence_score']
                    # Display information about the clicked polygon
                    annot.xy = (event.xdata, event.ydata)
                    text = f"Polygon {idx}\n\nConfidence:\n{confidence_score}\n\nSpecies and their probability:\n{species_info}"
                    annot.set_text(text)
                    annot.set_visible(True)
                    fig.canvas.draw()
                    break
    
    # Connect the click event to the handler function
    cid = fig.canvas.mpl_connect('button_press_event', on_click)

    figure = plt.figure()

    return figure

#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


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

demo.launch()