101 lines
2.6 KiB
Python
101 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple RTSP stream with YOLOv8 detection using OpenCV + subprocess for RTSP.
|
|
"""
|
|
|
|
import cv2
|
|
import subprocess
|
|
import time
|
|
from ultralytics import YOLO
|
|
|
|
BALL_CLASS_ID = 32 # sports ball in COCO
|
|
WIDTH = 1280
|
|
HEIGHT = 720
|
|
FPS = 25
|
|
|
|
|
|
def main():
|
|
print("Loading YOLOv8...")
|
|
model = YOLO("yolov8n.pt")
|
|
model.to("cuda")
|
|
print("Model loaded on CUDA")
|
|
|
|
# Open video source
|
|
source = "/opt/nvidia/deepstream/deepstream/samples/streams/sample_1080p_h264.mp4"
|
|
cap = cv2.VideoCapture(source)
|
|
|
|
if not cap.isOpened():
|
|
print("ERROR: Cannot open video")
|
|
return
|
|
|
|
# Start ffmpeg to serve RTSP via tcp
|
|
ffmpeg_cmd = [
|
|
'ffmpeg',
|
|
'-y',
|
|
'-f', 'rawvideo',
|
|
'-vcodec', 'rawvideo',
|
|
'-pix_fmt', 'bgr24',
|
|
'-s', f'{WIDTH}x{HEIGHT}',
|
|
'-r', str(FPS),
|
|
'-i', '-',
|
|
'-c:v', 'libx264',
|
|
'-preset', 'ultrafast',
|
|
'-tune', 'zerolatency',
|
|
'-f', 'rtsp',
|
|
'-rtsp_transport', 'tcp',
|
|
'rtsp://localhost:8554/live'
|
|
]
|
|
|
|
print("Starting ffmpeg...")
|
|
proc = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE)
|
|
|
|
print(f"\nRTSP stream: rtsp://pickle:8554/live")
|
|
print("Press Ctrl+C to stop\n")
|
|
|
|
frame_count = 0
|
|
start = time.time()
|
|
|
|
try:
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
|
continue
|
|
|
|
frame = cv2.resize(frame, (WIDTH, HEIGHT))
|
|
|
|
# Detection
|
|
results = model(frame, verbose=False, classes=[BALL_CLASS_ID], conf=0.3)
|
|
for r in results:
|
|
for box in r.boxes:
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0])
|
|
conf = float(box.conf[0])
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
cv2.putText(frame, f"Ball {conf:.2f}", (x1, y1-10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
|
|
|
|
# FPS
|
|
frame_count += 1
|
|
fps = frame_count / (time.time() - start)
|
|
cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
|
|
|
|
# Write to ffmpeg
|
|
proc.stdin.write(frame.tobytes())
|
|
|
|
if frame_count % 100 == 0:
|
|
print(f"Frame {frame_count}, FPS: {fps:.1f}")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nStopping...")
|
|
except BrokenPipeError:
|
|
print("ffmpeg pipe broken")
|
|
finally:
|
|
cap.release()
|
|
proc.stdin.close()
|
|
proc.wait()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|