Merge ui

main.py

@@ -6,13 +6,10 @@ 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
-
-
-
+import plotly.graph_objects as go
+import numpy as np
+import ast
 
 def row_to_feature(row):
     feature = {
@@ -20,10 +17,9 @@ def row_to_feature(row):
         "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()]
 
@@ -40,35 +36,19 @@ def export_geojson(df, filename):
 
     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, progress=gr.Progress()):
+def process_image(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)
 
-    progress(0, desc="Running detectree2")
     run_detectree2(image_path, store_path=output_directory)
-    progress(0.1, desc="Postprocessing")
 
     processed_output_df = postprocess(output_directory + '/detectree2_delin.geojson', output_directory + '/processed_delin')
 
     processed_geojson = output_directory + '/processed_delin.geojson'
 
-    progress(0.2, desc="Generating tree images")
     generate_tree_images(processed_geojson, image_path)
 
     output_folder = './tree_images'
@@ -77,82 +57,82 @@ def greet(image_path: str, progress=gr.Progress()):
 
     for file_name in os.listdir(output_folder):
         file_path = os.path.join(output_folder, file_name)
+        print(file_path)
         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
+    # 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'
 
-    progress(0.3, desc="Exporting geojson")
     export_geojson(processed_output_df, final_output_path)
 
+    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(final_output_path)
-
-    # Set the interactive backend to Qt5Agg
-    # plt.switch_backend('Qt5Agg') # You have to install PyQt5
-
-    # Enable interactive mode
-    # plt.ion()
-
-    progress(0.4, desc="Plotting")
-
-    # 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
+    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):
+    geojson_path, image_path = process_image(image_path)
+    fig = plot_results(geojson_path, image_path)
+    
+    return fig
     
-    # 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
 
@@ -163,10 +143,15 @@ def greet(image_path: str, progress=gr.Progress()):
 # Read the TIF file
 
 
-demo = gr.Interface(
-    fn=greet,
-    inputs=gr.File(type='filepath'),
-    outputs=gr.Plot(label="Tree Crowns")
-)
+def main():
+    demo = gr.Interface(
+        fn=greet,
+        inputs=gr.File(type='filepath'),
+        outputs=gr.Plot(label="Tree Crowns")
+    )
+
+    demo.launch(server_name='0.0.0.0', share=True)
+
+if __name__ == "__main__":
+    main()
 
-demo.launch(server_name='0.0.0.0')