Update app.py

app.py

@@ -1,8 +1,8 @@
 import cv2
-import csv
 import tempfile
 import gradio as gr
 from ultralytics import YOLO
+import pandas as pd
 
 def process_video(video_file):
     # Define colors for each class (8 classes)
@@ -26,49 +26,62 @@ def process_video(video_file):
     # Open the video file
     cap = cv2.VideoCapture(video_file)
 
-    # Prepare CSV file for writing
-    csv_path = tempfile.mktemp(suffix=".csv")
-    with open(csv_path, mode='w', newline='') as csv_file:
-        writer = csv.writer(csv_file)
-        writer.writerow(["frame", "id", "class", "x", "y", "w", "h"])
+    # Prepare DataFrame for storing detection data
+    columns = ["frame", "id", "class", "x", "y", "w", "h"]
+    df = pd.DataFrame(columns=columns)
 
-        frame_id = 0
+    frame_id = 0
 
-        # Loop through the video frames
-        while cap.isOpened():
-            # Read a frame from the video
-            success, frame = cap.read()
+    # Loop through the video frames
+    while cap.isOpened():
+        # Read a frame from the video
+        success, frame = cap.read()
 
-            if success:
-                frame_id += 1
+        if success:
+            frame_id += 1
 
-                # Run YOLOv8 tracking on the frame, persisting tracks between frames
-                results = model.track(frame, persist=True, tracker="insect_tracker.yaml")
+            # Run YOLOv8 tracking on the frame, persisting tracks between frames
+            results = model.track(frame, persist=True, tracker="insect_tracker.yaml")
 
-                for result in results:
-                    boxes = result.boxes.cpu().numpy()
-                    confidences = boxes.conf
-                    class_ids = boxes.cls
+            for result in results:
+                boxes = result.boxes.cpu().numpy()
+                confidences = boxes.conf
+                class_ids = boxes.cls
 
-                    for i, box in enumerate(boxes):
-                        class_id = int(class_ids[i])
-                        confidence = confidences[i]
+                for i, box in enumerate(boxes):
+                    class_id = int(class_ids[i])
+                    confidence = confidences[i]
 
-                        # Write detection data to CSV
-                        writer.writerow([frame_id, box.id, int(box.cls[0]), box.xywh[0][0], box.xywh[0][1],
-                                        box.xywh[0][2], box.xywh[0][3]])
+                    # Append detection data to DataFrame
+                    new_row = pd.DataFrame({
+                        "frame": [frame_id],
+                        "id": [box.id],
+                        "class": [int(box.cls[0])],
+                        "x": [box.xywh[0][0]],
+                        "y": [box.xywh[0][1]],
+                        "w": [box.xywh[0][2]],
+                        "h": [box.xywh[0][3]]
+                    })
+                    df = pd.concat([df, new_row], ignore_index=True)
 
-            else:
-                break
+        else:
+            break
 
     # Release the video capture
     cap.release()
 
-    return csv_path
+    # Save DataFrame to CSV
+    csv_path = tempfile.mktemp(suffix=".csv")
+    df.to_csv(csv_path, index=False)
+
+    return df, csv_path
 
 # Create a Gradio interface
 inputs = gr.Video(label="Input Video")
-output = gr.File(label="CSV File")
+outputs = [
+    gr.DataFrame(label="Detection Data"),
+    gr.File(label="Download CSV")
+]
 
-gr.Interface(fn=process_video, inputs=inputs, outputs=output).launch()
+gr.Interface(fn=process_video, inputs=inputs, outputs=outputs).launch()