Update app.py

app.py

@@ -1,58 +1,109 @@
-import gradio as gr
 import cv2
+import csv
 import tempfile
-import numpy as np
+import gradio as gr
 from ultralytics import YOLO
 
-# Load the YOLOv8 model
-model = YOLO('yolov8m.pt')  # Ensure you have the correct model path
-
 def process_video(video_file):
-    # Create a temporary directory to store processed frames
-    temp_dir = tempfile.TemporaryDirectory()
-    
+    # 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
+    ]
+
+    # 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")
+
     # Open the video file
-    cap = cv2.VideoCapture(video_file.name)
-    
-    # Get video properties
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    codec = cv2.VideoWriter_fourcc(*'mp4v')
-    
-    # Create a VideoWriter object to save the processed video
-    output_path = f"{temp_dir.name}/output.mp4"
-    out = cv2.VideoWriter(output_path, codec, fps, (width, height))
+    cap = cv2.VideoCapture(video_file)
 
+    # Prepare CSV file for writing
+    csv_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
+    writer = csv.writer(csv_file)
+    writer.writerow(["frame", "id", "class", "x", "y", "w", "h"])
+
+    frame_id = 0
+
+    # Initialize a list to store annotated frames
+    annotated_frames = []
+
+    # Loop through the video frames
     while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+        # Read a frame from the video
+        success, frame = cap.read()
+
+        if success:
+            frame_id += 1
+
+            # Run YOLOv8 tracking on the frame, persisting tracks between frames
+            results = model.track(frame, persist=True)
 
-        # Use YOLO model to detect objects in the frame
-        results = model(frame)
+            for result in results:
+                boxes = result.boxes.cpu().numpy()
+                confidences = boxes.conf
+                class_ids = boxes.cls
 
-        # Draw bounding boxes and labels on the frame
-        annotated_frame = results[0].plot()
+                for i, box in enumerate(boxes):
+                    class_id = int(class_ids[i])
+                    confidence = confidences[i]
 
-        # Write the frame to the output video
-        out.write(annotated_frame)
-    
-    # Release the VideoCapture and VideoWriter objects
+                    color = colors[class_id % len(colors)]  # Use the color corresponding to the class
+                    label = f'{class_names[class_id]}: {confidence:.2f}'
+
+                    # Draw the rectangle
+                    cv2.rectangle(frame, (int(box.xyxy[0][0]), int(box.xyxy[0][1])),
+                                (int(box.xyxy[0][2]), int(box.xyxy[0][3])), color, 2)
+
+                    # Display the label above the rectangle
+                    label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 2.0, 2)
+                    label_y = max(int(box.xyxy[0][1]) - label_size[1], 0)
+                    cv2.rectangle(frame, (int(box.xyxy[0][0]), label_y - label_size[1]),
+                                (int(box.xyxy[0][0]) + label_size[0], label_y + label_size[1]), color, -1)
+                    cv2.putText(frame, label, (int(box.xyxy[0][0]), label_y), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (255, 255, 255), 2)
+
+                    # 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]])
+
+            # Add annotated frame to the list
+            annotated_frames.append(frame)
+
+        else:
+            break
+
+    # Release the video capture
     cap.release()
+
+    # Compile annotated frames into a video
+    output_video_path = tempfile.NamedTemporaryFile(suffix='.mp4').name
+    height, width, _ = annotated_frames[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_video_path, fourcc, 30.0, (width, height))
+    for frame in annotated_frames:
+        out.write(frame)
     out.release()
 
-    return output_path
+    # Close CSV file
+    csv_file.close()
+
+    return output_video_path, csv_file.name
+
+# Create a Gradio interface
+inputs = gr.Video(label="Input Video")
+outputs = [gr.Video(label="Annotated Video"), gr.File(label="CSV File")]
+
+gradio_app=gr.Interface(fn=process_video, inputs=inputs, outputs=outputs)
+
 
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=process_video,
-    inputs=gr.Video(),
-    outputs=gr.Video(),
-    title="YOLOv8 Video Object Detection",
-    description="Upload a video and apply YOLOv8 object detection."
-)
 
-# Launch the Gradio app
 if __name__ == "__main__":
-    iface.launch()
+    gradio_app.launch()