Initial commit

jetson/simple_rtsp_stream.py

@@ -0,0 +1,100 @@
+#!/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()