Upload 3 files

app.py

@@ -1,6 +1,5 @@
-import os
 import cv2
-from fastapi import FastAPI, Request, UploadFile, File
+from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse, HTMLResponse
 from fastapi.templating import Jinja2Templates
 from typing import Generator
@@ -8,63 +7,67 @@ from ultralytics import YOLO
 import numpy as np
 
 app = FastAPI()
-templates = Jinja2Templates(directory="templates")
 
 # Load the YOLOv8 model
-model = YOLO("yolov8l.pt")
+model = YOLO("yolov8n.pt")
+
+# Open the video file
+video_path = "demo.mp4"
+cap = cv2.VideoCapture(video_path)
 
-video_path = None
-cap = None
 bird_count = 0
 tracker_initialized = False
-trackers = None
-
-@app.post("/upload_video/")
-async def upload_video(file: UploadFile = File(...)):
-    global cap, tracker_initialized, trackers
-
-    # Save uploaded video file
-    file_location = f"uploads/{file.filename}"
-    with open(file_location, "wb") as f:
-        f.write(file.file.read())
 
-    # Open the uploaded video file
-    cap = cv2.VideoCapture(file_location)
-
-    # Reset tracker and counter
-    tracker_initialized = False
+# Initialize trackers based on OpenCV version
+try:
+    if hasattr(cv2, 'legacy'):
+        trackers = cv2.legacy.MultiTracker_create()
+    else:
+        trackers = cv2.TrackerCSRT_create()
+except AttributeError:
     trackers = None
-
-    return {"info": f"file '{file.filename}' saved at '{file_location}'"}
+    tracker_initialized = False
 
 def process_video() -> Generator[bytes, None, None]:
-    global bird_count, tracker_initialized, trackers, cap
+    global bird_count, tracker_initialized, trackers
     while cap.isOpened():
+        # Read a frame from the video
         success, frame = cap.read()
 
         if success:
             frame_height, frame_width = frame.shape[:2]
             if not tracker_initialized:
+                # Run YOLOv8 inference on the frame
                 results = model(frame)
+
+                # Extract the detected objects
                 detections = results[0].boxes.data.cpu().numpy()
+
+                # Filter results to include only the "bird" class (class id 14 in COCO)
                 bird_results = [detection for detection in detections if int(detection[5]) == 14]
 
+                # Initialize trackers for bird results
                 try:
-                    trackers = cv2.legacy.MultiTracker_create()
+                    if hasattr(cv2, 'legacy'):
+                        trackers = cv2.legacy.MultiTracker_create()
+                    else:
+                        trackers = cv2.MultiTracker_create()
+
+                    for res in bird_results:
+                        x1, y1, x2, y2, confidence, class_id = res
+                        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                        if 0 <= x1 < frame_width and 0 <= y1 < frame_height and x2 <= frame_width and y2 <= frame_height:
+                            bbox = (x1, y1, x2 - x1, y2 - y1)
+                            tracker = cv2.legacy.TrackerCSRT_create() if hasattr(cv2, 'legacy') else cv2.TrackerCSRT_create()
+                            trackers.add(tracker, frame, bbox)
+
+                    bird_count = len(bird_results)
+                    tracker_initialized = True
                 except AttributeError:
-                    trackers = cv2.MultiTracker_create()
-
-                for res in bird_results:
-                    x1, y1, x2, y2, confidence, class_id = res
-                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
-                    if 0 <= x1 < frame_width and 0 <= y1 < frame_height and x2 <= frame_width and y2 <= frame_height:
-                        bbox = (x1, y1, x2 - x1, y2 - y1)
-                        tracker = cv2.legacy.TrackerCSRT_create() if hasattr(cv2, 'legacy') else cv2.TrackerCSRT_create()
-                        trackers.add(tracker, frame, bbox)
-
-                bird_count = len(bird_results)
-                tracker_initialized = True
+                    trackers = None
+                    tracker_initialized = False
             else:
+                # Update trackers and get updated positions
                 success, boxes = trackers.update(frame)
 
                 if success:
@@ -76,14 +79,19 @@ def process_video() -> Generator[bytes, None, None]:
                 else:
                     tracker_initialized = False
 
+            # Encode the frame in JPEG format
             ret, buffer = cv2.imencode('.jpg', frame)
             frame = buffer.tobytes()
+
+            # Use generator to yield the frame
             yield (b'--frame\r\n'
                    b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
         else:
             break
     cap.release()
 
+templates = Jinja2Templates(directory="templates")
+
 @app.get("/", response_class=HTMLResponse)
 async def index(request: Request):
     return templates.TemplateResponse("index.html", {"request": request, "bird_count": bird_count})