Update app.py

app.py

@@ -3,7 +3,8 @@ import tempfile
 import gradio as gr
 from ultralytics import YOLO
 import pandas as pd
-import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+import numpy as np
 
 # Define the label mapping
 label_mapping = {
@@ -17,18 +18,6 @@ label_mapping = {
 }
 
 def process_video(video_file):
-    # Define colors for each class (8 classes)
-    colors = [
-        (255, 0, 0),   # Class 0 - Blue
-        (50, 205, 50),   # Class 1 - Green
-        (0, 0, 255),   # Class 2 - Red
-        (255, 255, 0), # Class 3 - Cyan
-        (255, 0, 255), # Class 4 - Magenta
-        (255, 140, 0), # Class 5 - Orange
-        (128, 0, 128), # Class 6 - Purple
-        (0, 128, 128)  # Class 7 - Teal
-    ]
-
     # Load the YOLOv8 model
     model = YOLO("insect_detection4.pt")
 
@@ -100,37 +89,41 @@ def process_video(video_file):
     # Read the DataFrame from the CSV file
     df_from_csv = pd.read_csv(csv_path)
 
-    # Create the plot from the CSV data
-    plt.figure(figsize=(10, 6))
-    for insect_id in df_from_csv['insect_id'].unique():
-        frames = df_from_csv[df_from_csv['insect_id'] == insect_id]['frame']
-        insect_class = label_mapping[df_from_csv[df_from_csv['insect_id'] == insect_id]['class'].values[0]]
-        plt.plot(frames, [insect_id] * len(frames), 'o-', label=f'{insect_class} {insect_id}')
+    # Create the interactive plot from the CSV data
+    fig = go.Figure()
 
-    plt.xlabel('Frame')
-    plt.ylabel('Insect ID')
-    plt.title('Temporal distribution of insects')
-    plt.legend()
-    plt.grid(True)
+    for insect_id, group in df_from_csv.groupby('insect_id'):
+        class_name = label_mapping[group.iloc[0]['class']]
+        color = 'rgb({}, {}, {})'.format(*np.random.randint(0, 256, 3))
+        hover_text = group.apply(lambda row: f'Insect ID: {int(row["insect_id"])}, Class: {class_name}, Frame: {int(row["frame"])}', axis=1)
+        fig.add_trace(go.Scatter(x=group['frame'], y=group['insect_id'], mode='markers', marker=dict(color=color), name=f'{class_name} {insect_id}', 
+                                 hoverinfo='text', hovertext=hover_text))
 
-    # Save the plot to a temporary file
-    plot_path = tempfile.mktemp(suffix=".png")
-    plt.savefig(plot_path)
-    plt.close()
+    fig.update_layout(title='Temporal distribution of insects',
+                      xaxis_title='Frame',
+                      yaxis_title='Insect ID',
+                      hovermode='closest')
 
     gallery_items = [(crop_path, f'{label} {insect_id}') for insect_id, (crop_path, label) in unique_insect_crops.items()]
 
-    return plot_path, gallery_items, csv_path
+    return fig, gallery_items, csv_path
 
 # Create a Gradio interface
-example_video = "insect_trap_video_example.mp4"  # Replace with the actual path to your example video
+example_video = "insect_trap_video_exemple.mp4"  # Replace with the actual path to your example video
 
 inputs = gr.Video(label="Input Insect Trap Video", value=example_video)
 outputs = [
-    gr.Image(label="Insect Detection Plot"),
-    gr.Gallery(label="Unique Insect Crops"),  # Added a gallery to display insect crops with labels
+    gr.Plot(label="Insect Detection Plot"),
+    gr.Gallery(label="Insect Gallery"),  # Added a gallery to display insect crops with labels
     gr.File(label="Download CSV")
 ]
 
-gr.Interface(fn=process_video, inputs=inputs, outputs=outputs, title="InsectSpy Automated Insect Detection", examples=[example_video]).launch()
+description = """
+Uncover the Secret Lives of Insects! 🐝🦋🕷️
+
+Upload your video now to track, visualize, and explore insect activity with our cutting-edge detection tool. You can get started with the example video.
+"""
+
+gr.Interface(fn=process_video, inputs=inputs, outputs=outputs, title= 'InsectSpy 🕵️‍♂️🦗', description=description, examples=[example_video]).launch()
+