python-rtsp-worker/stream_decoder_test.py
2025-11-11 02:28:33 +07:00

60 lines
1.3 KiB
Python

#!/usr/bin/env python3
"""
Simple example: Decode 4 RTSP streams and display with OpenCV using callbacks
Usage:
python stream_decoder_test.py
"""
import os
import cv2
from dotenv import load_dotenv
from services import StreamDecoderFactory
load_dotenv()
# Frame storage for each camera
frames = {1: None, 2: None, 3: None, 4: None}
def make_callback(cam_id):
"""Create callback for specific camera"""
def callback(frame_ref):
# Transfer to CPU and convert RGB to BGR
frame = frame_ref.rgb_tensor.cpu().permute(1, 2, 0).numpy()
frames[cam_id] = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
frame_ref.free()
return callback
# Initialize factory and decoders
factory = StreamDecoderFactory(gpu_id=0)
decoders = []
for i in range(1, 5):
url = os.getenv(f"CAMERA_URL_{i}")
decoder = factory.create_decoder(url, buffer_size=5)
decoder.register_frame_callback(make_callback(i))
decoder.start()
decoders.append(decoder)
print(f"Camera {i} started")
# Display loop
print("Press 'q' to quit")
while True:
# Show each camera in separate window
for cam_id, frame in frames.items():
if frame is not None:
cv2.imshow(f"Camera {cam_id}", frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
# Cleanup
for decoder in decoders:
decoder.stop()
cv2.destroyAllWindows()