!/usr/bin/env python3

import cv2
import torch
import yt_dlp
import subprocess
import time
from flask import Flask, Response
import numpy as np
import warnings
import os

warnings.simplefilter("ignore", category=FutureWarning)

app = Flask(__name__)

YOUTUBE_URL = "https://www.youtube.com/watch?v=t45_gP7I82I"  # Stream URL
CONFIDENCE_THRESHOLD = 0.4  # Confidence threshold for object detection
MODEL = "yolov5s"  # YOLO model version (yolov5s, yolov5m, etc.)

# Load YOLOv5 model 
print("\U0001F504 Loading YOLOv5 model...")
model = torch.hub.load("ultralytics/yolov5", "custom", path=MODEL, force_reload=True)
print("\u2705 YOLOv5 loaded successfully!")

# Load object exclusion list
EXCLUSION_FILE = "objectnotlist.txt"
detected_objects = set()
if os.path.exists(EXCLUSION_FILE):
    with open(EXCLUSION_FILE, "r") as f:
        object_exclusion_list = set(line.strip().lower() for line in f)
else:
    object_exclusion_list = set()

def get_stream_url():
    """Fetch fresh 720p YouTube stream URL using yt-dlp."""
    ydl_opts = {'quiet': True, 'format': 'bestvideo[height=720]'}
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(YOUTUBE_URL, download=False)
        return info_dict.get("url", None)

def generate_frames():
    """Capture video frames, apply object detection, and stream as MJPEG."""
    stream_url = get_stream_url()
    if not stream_url:
        print("\u274C Failed to fetch stream URL!")
        return

    print("\U0001F3A5 Starting FFmpeg stream...")
    ffmpeg_process = subprocess.Popen([
        "ffmpeg", "-re", "-i", stream_url, "-r", "10",  # frame rate to 10 FPS
        "-fflags", "nobuffer", "-flags", "low_delay",  # buffering delay
        "-f", "rawvideo", "-pix_fmt", "bgr24", "pipe:1"
    ], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=10**8)

    while True:
        raw_frame = ffmpeg_process.stdout.read(1280 * 720 * 3)  # Read raw BGR frame for 720p
        if not raw_frame:
            print("\u274C No frame received!")
            break

        frame = np.frombuffer(raw_frame, np.uint8).reshape((720, 1280, 3))  # Convert to NumPy array

        # Run YOLO object detection
        results = model(frame)
        detections = results.pandas().xyxy[0]  # Convert detections to Pandas DataFrame

        if detections.empty:
            print("\u274C No objects detected in this frame.")
        else:
            for _, row in detections.iterrows():
                detected_object = row["name"].lower()
                if detected_object not in object_exclusion_list and detected_object not in detected_objects:
                    detected_objects.add(detected_object)
                    with open("detections.txt", "a") as f:
                        f.write(detected_object + "\n")
                    print(f"\u2705 >: {detected_object}")

        # Encode and yield the frame as JPEG
        _, buffer = cv2.imencode('.jpg', frame)
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')

@app.route('/video')
def video_feed():
    """Stream processed video frames."""
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    print("\U0001F680 Server running at http://localhost:5000/video")
    app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)