Update app.py

app.py

@@ -5,6 +5,17 @@ from ultralytics import YOLO
 import pandas as pd
 import matplotlib.pyplot as plt
 
+# Define the label mapping
+label_mapping = {
+    0: 'Hymenoptera',
+    1: 'Mantodea',
+    2: 'Odonata',
+    3: 'Orthoptera',
+    4: 'Coleoptera',
+    5: 'Lepidoptera',
+    6: 'Hemiptera'
+}
+
 def process_video(video_file):
     # Define colors for each class (8 classes)
     colors = [
@@ -18,9 +29,6 @@ def process_video(video_file):
         (0, 128, 128)  # Class 7 - Teal
     ]
 
-    # Define class names (example names, replace with actual class names if available)
-    class_names = ['Hymenoptera', 'Mantodea', 'Odonata', 'Orthoptera', 'Coleoptera', 'Lepidoptera', 'Hemiptera']
-
     # Load the YOLOv8 model
     model = YOLO("insect_detection4.pt")
 
@@ -28,7 +36,7 @@ def process_video(video_file):
     cap = cv2.VideoCapture(video_file)
 
     # Prepare DataFrame for storing detection data
-    columns = ["frame", "id", "class", "x", "y", "w", "h"]
+    columns = ["frame", "insect_id", "class", "x", "y", "w", "h"]
     df = pd.DataFrame(columns=columns)
 
     frame_id = 0
@@ -53,13 +61,13 @@ def process_video(video_file):
                 for i, box in enumerate(boxes):
                     class_id = int(class_ids[i])
                     confidence = confidences[i]
-                    unique_id = int(box.id[0])  # Ensure the ID is an integer
+                    insect_id = int(box.id[0])  # Ensure the ID is an integer
 
                     # Append detection data to DataFrame
                     new_row = pd.DataFrame({
                         "frame": [frame_id],
-                        "id": [unique_id],
-                        "class": [int(box.cls[0])],
+                        "insect_id": [insect_id],
+                        "class": [class_id],
                         "x": [box.xywh[0][0]],
                         "y": [box.xywh[0][1]],
                         "w": [box.xywh[0][2]],
@@ -68,7 +76,7 @@ def process_video(video_file):
                     df = pd.concat([df, new_row], ignore_index=True)
 
                     # Crop and save the image of the insect
-                    if unique_id not in unique_insect_crops:
+                    if insect_id not in unique_insect_crops:
                         x_center, y_center, width, height = box.xywh[0]
                         x1 = int(x_center - width / 2)
                         y1 = int(y_center - height / 2)
@@ -77,7 +85,7 @@ def process_video(video_file):
                         insect_crop = frame[y1:y2, x1:x2]
                         crop_path = tempfile.mktemp(suffix=".png")
                         cv2.imwrite(crop_path, insect_crop)
-                        unique_insect_crops[unique_id] = crop_path
+                        unique_insect_crops[insect_id] = (crop_path, label_mapping[class_id])
 
         else:
             break
@@ -94,9 +102,10 @@ def process_video(video_file):
 
     # Create the plot from the CSV data
     plt.figure(figsize=(10, 6))
-    for unique_id in df_from_csv['id'].unique():
-        frames = df_from_csv[df_from_csv['id'] == unique_id]['frame']
-        plt.plot(frames, [unique_id] * len(frames), 'o-', label=f'Insect {unique_id}')
+    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}')
 
     plt.xlabel('Frame')
     plt.ylabel('Insect ID')
@@ -109,17 +118,19 @@ def process_video(video_file):
     plt.savefig(plot_path)
     plt.close()
 
-    gallery_items = [(crop_path, f'Insect {unique_id}') for unique_id, crop_path in unique_insect_crops.items()]
+    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
 
 # Create a Gradio interface
-inputs = gr.Video(label="Input Insect Trap Video")
+example_video = "path_to_example_video.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.File(label="Download CSV")
 ]
 
-gr.Interface(fn=process_video, inputs=inputs, outputs=outputs).launch()
+gr.Interface(fn=process_video, inputs=inputs, outputs=outputs, examples=[example_video]).launch()