80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import cv2
|
|
import yt_dlp
|
|
import subprocess
|
|
import numpy as np
|
|
from flask import Flask, Response
|
|
from ultralytics import YOLO #changed
|
|
|
|
|
|
app = Flask(__name__)
|
|
|
|
YOUTUBE_URL = "https://www.youtube.com/watch?v=i3w7qZVSAsY"
|
|
CONFIDENCE_THRESHOLD = 0.25
|
|
|
|
# Load YOLOv8 model
|
|
print("Loading YOLOv8 model...")
|
|
model = YOLO("yolov8s.pt") # Using YOLOv8 small model
|
|
print("YOLOv8 loaded successfully!")
|
|
|
|
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("Failed to fetch stream URL!")
|
|
return
|
|
|
|
print("Starting FFmpeg stream...")
|
|
ffmpeg_process = subprocess.Popen([
|
|
"ffmpeg", "-re", "-i", stream_url, "-r", "10", # Frame rate to 10 FPS
|
|
"-fflags", "nobuffer", "-flags", "low_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)
|
|
if not raw_frame:
|
|
print("No frame received!")
|
|
break
|
|
|
|
frame = np.frombuffer(raw_frame, np.uint8).reshape((720, 1280, 3))
|
|
|
|
# Run YOLOv8 object detection
|
|
results = model(frame, stream=True)
|
|
|
|
for result in results:
|
|
boxes = result.boxes.xyxy.cpu().numpy()
|
|
scores = result.boxes.conf.cpu().numpy()
|
|
class_ids = result.boxes.cls.cpu().numpy().astype(int)
|
|
names = result.names
|
|
|
|
for i in range(len(boxes)):
|
|
if scores[i] > CONFIDENCE_THRESHOLD:
|
|
x1, y1, x2, y2 = map(int, boxes[i])
|
|
label = f"{names[class_ids[i]]} ({scores[i]:.2f})"
|
|
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
|
|
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
|
# 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("Running at http://localhost:5000/video")
|
|
app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)
|
|
|