117 lines
3.7 KiB
Python
Executable file
117 lines
3.7 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""
|
|
Quick test for event-driven stream processing - runs for 20 seconds.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import logging
|
|
from dotenv import load_dotenv
|
|
|
|
from services import StreamConnectionManager, YOLOv8Utils, COCO_CLASSES
|
|
|
|
# Setup logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def main():
|
|
"""Quick test with callback pattern"""
|
|
logger.info("=== Quick Event-Driven Test (20 seconds) ===")
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
camera_url = os.getenv('CAMERA_URL_1')
|
|
if not camera_url:
|
|
logger.error("CAMERA_URL_1 not found in .env file")
|
|
return
|
|
|
|
# Create manager
|
|
manager = StreamConnectionManager(
|
|
gpu_id=0,
|
|
batch_size=16,
|
|
force_timeout=0.05, # 50ms
|
|
poll_interval=0.01, # 100 FPS
|
|
)
|
|
|
|
# Initialize with YOLOv8 model
|
|
model_path = "models/yolov8n.trt"
|
|
logger.info(f"Initializing with model: {model_path}")
|
|
|
|
await manager.initialize(
|
|
model_path=model_path,
|
|
model_id="yolo",
|
|
preprocess_fn=YOLOv8Utils.preprocess,
|
|
postprocess_fn=YOLOv8Utils.postprocess,
|
|
)
|
|
|
|
result_count = 0
|
|
|
|
# Define callback for tracking results
|
|
def on_tracking_result(result):
|
|
nonlocal result_count
|
|
result_count += 1
|
|
|
|
if result_count % 5 == 0: # Log every 5th result
|
|
logger.info(f"[{result.stream_id}] Frame {result.metadata.get('frame_number', 0)}")
|
|
logger.info(f" Tracked objects: {len(result.tracked_objects)}")
|
|
|
|
for obj in result.tracked_objects[:3]: # Show first 3
|
|
class_name = COCO_CLASSES.get(obj.class_id, f"Class {obj.class_id}")
|
|
logger.info(
|
|
f" Track ID {obj.track_id}: {class_name}, "
|
|
f"conf={obj.confidence:.2f}"
|
|
)
|
|
|
|
def on_error(error):
|
|
logger.error(f"Stream error: {error}")
|
|
|
|
# Connect to stream
|
|
logger.info(f"Connecting to stream...")
|
|
connection = await manager.connect_stream(
|
|
rtsp_url=camera_url,
|
|
stream_id="test_camera",
|
|
on_tracking_result=on_tracking_result,
|
|
on_error=on_error,
|
|
)
|
|
|
|
# Monitor for 20 seconds with stats updates
|
|
for i in range(4): # 4 x 5 seconds = 20 seconds
|
|
await asyncio.sleep(5)
|
|
|
|
stats = manager.get_stats()
|
|
model_stats = stats['model_controller']
|
|
|
|
logger.info(f"\n=== Stats Update {i+1}/4 ===")
|
|
logger.info(f"Results received: {result_count}")
|
|
logger.info(f"Buffer A: {model_stats['buffer_a_size']} ({model_stats['buffer_a_state']})")
|
|
logger.info(f"Buffer B: {model_stats['buffer_b_size']} ({model_stats['buffer_b_state']})")
|
|
logger.info(f"Active buffer: {model_stats['active_buffer']}")
|
|
logger.info(f"Total frames processed: {model_stats['total_frames_processed']}")
|
|
logger.info(f"Total batches: {model_stats['total_batches_processed']}")
|
|
logger.info(f"Avg batch size: {model_stats['avg_batch_size']:.2f}")
|
|
|
|
# Final statistics
|
|
stats = manager.get_stats()
|
|
logger.info("\n=== Final Statistics ===")
|
|
logger.info(f"Total results received: {result_count}")
|
|
logger.info(f"Manager: {stats['manager']}")
|
|
logger.info(f"Model Controller: {stats['model_controller']}")
|
|
logger.info(f"Connection: {stats['connections']['test_camera']}")
|
|
|
|
# Cleanup
|
|
logger.info("\nShutting down...")
|
|
await manager.shutdown()
|
|
logger.info("Test complete!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
logger.info("\nInterrupted by user")
|
|
except Exception as e:
|
|
logger.error(f"Error: {e}", exc_info=True)
|