from typing import Any, Dict import os import json import time import queue import torch import cv2 import base64 import logging import threading import requests import asyncio import psutil import zipfile from urllib.parse import urlparse from fastapi import FastAPI, WebSocket from fastapi.websockets import WebSocketDisconnect from websockets.exceptions import ConnectionClosedError from ultralytics import YOLO # Import shared pipeline functions from siwatsystem.pympta import load_pipeline_from_zip, run_pipeline app = FastAPI() # Global dictionaries to keep track of models and streams # "models" now holds a nested dict: { camera_id: { modelId: model_tree } } models: Dict[str, Dict[str, Any]] = {} streams: Dict[str, Dict[str, Any]] = {} with open("config.json", "r") as f: config = json.load(f) poll_interval = config.get("poll_interval_ms", 100) reconnect_interval = config.get("reconnect_interval_sec", 5) TARGET_FPS = config.get("target_fps", 10) poll_interval = 1000 / TARGET_FPS logging.info(f"Poll interval: {poll_interval}ms") max_streams = config.get("max_streams", 5) max_retries = config.get("max_retries", 3) # Configure logging logging.basicConfig( level=logging.DEBUG, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.FileHandler("app.log"), logging.StreamHandler() ] ) # Ensure the models directory exists os.makedirs("models", exist_ok=True) # Constants for heartbeat and timeouts HEARTBEAT_INTERVAL = 2 # seconds WORKER_TIMEOUT_MS = 10000 # Locks for thread-safe operations streams_lock = threading.Lock() models_lock = threading.Lock() # Add helper to download mpta ZIP file from a remote URL def download_mpta(url: str, dest_path: str) -> str: try: os.makedirs(os.path.dirname(dest_path), exist_ok=True) response = requests.get(url, stream=True) if response.status_code == 200: with open(dest_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) logging.info(f"Downloaded mpta file from {url} to {dest_path}") return dest_path else: logging.error(f"Failed to download mpta file (status code {response.status_code})") return None except Exception as e: logging.error(f"Exception downloading mpta file from {url}: {e}") return None #################################################### # Detection and frame processing functions #################################################### @app.websocket("/") async def detect(websocket: WebSocket): logging.info("WebSocket connection accepted") persistent_data_dict = {} async def handle_detection(camera_id, stream, frame, websocket, model_tree, persistent_data): try: detection_result = run_pipeline(frame, model_tree) detection_data = { "type": "imageDetection", "cameraIdentifier": camera_id, "timestamp": time.time(), "data": { "detection": detection_result if detection_result else None, "modelId": stream["modelId"], "modelName": stream["modelName"] } } logging.debug(f"Sending detection data for camera {camera_id}: {detection_data}") await websocket.send_json(detection_data) return persistent_data except Exception as e: logging.error(f"Error in handle_detection for camera {camera_id}: {e}") return persistent_data def frame_reader(camera_id, cap, buffer, stop_event): retries = 0 try: while not stop_event.is_set(): try: ret, frame = cap.read() if not ret: logging.warning(f"Connection lost for camera: {camera_id}, retry {retries+1}/{max_retries}") cap.release() time.sleep(reconnect_interval) retries += 1 if retries > max_retries and max_retries != -1: logging.error(f"Max retries reached for camera: {camera_id}") break # Re-open cap = cv2.VideoCapture(streams[camera_id]["rtsp_url"]) if not cap.isOpened(): logging.error(f"Failed to reopen RTSP stream for camera: {camera_id}") continue continue retries = 0 # Overwrite old frame if buffer is full if not buffer.empty(): try: buffer.get_nowait() except queue.Empty: pass buffer.put(frame) except cv2.error as e: logging.error(f"OpenCV error for camera {camera_id}: {e}") cap.release() time.sleep(reconnect_interval) retries += 1 if retries > max_retries and max_retries != -1: logging.error(f"Max retries reached after OpenCV error for camera {camera_id}") break cap = cv2.VideoCapture(streams[camera_id]["rtsp_url"]) if not cap.isOpened(): logging.error(f"Failed to reopen RTSP stream for camera {camera_id} after OpenCV error") continue except Exception as e: logging.error(f"Unexpected error for camera {camera_id}: {e}") cap.release() break except Exception as e: logging.error(f"Error in frame_reader thread for camera {camera_id}: {e}") async def process_streams(): logging.info("Started processing streams") try: while True: start_time = time.time() with streams_lock: current_streams = list(streams.items()) for camera_id, stream in current_streams: buffer = stream["buffer"] if not buffer.empty(): frame = buffer.get() with models_lock: model_tree = models.get(camera_id, {}).get(stream["modelId"]) key = (camera_id, stream["modelId"]) persistent_data = persistent_data_dict.get(key, {}) updated_persistent_data = await handle_detection( camera_id, stream, frame, websocket, model_tree, persistent_data ) persistent_data_dict[key] = updated_persistent_data elapsed_time = (time.time() - start_time) * 1000 # ms sleep_time = max(poll_interval - elapsed_time, 0) logging.debug(f"Elapsed time: {elapsed_time}ms, sleeping for: {sleep_time}ms") await asyncio.sleep(sleep_time / 1000.0) except asyncio.CancelledError: logging.info("Stream processing task cancelled") except Exception as e: logging.error(f"Error in process_streams: {e}") async def send_heartbeat(): while True: try: cpu_usage = psutil.cpu_percent() memory_usage = psutil.virtual_memory().percent if torch.cuda.is_available(): gpu_usage = torch.cuda.memory_allocated() / (1024 ** 2) # MB gpu_memory_usage = torch.cuda.memory_reserved() / (1024 ** 2) # MB else: gpu_usage = None gpu_memory_usage = None camera_connections = [ { "cameraIdentifier": camera_id, "modelId": stream["modelId"], "modelName": stream["modelName"], "online": True } for camera_id, stream in streams.items() ] state_report = { "type": "stateReport", "cpuUsage": cpu_usage, "memoryUsage": memory_usage, "gpuUsage": gpu_usage, "gpuMemoryUsage": gpu_memory_usage, "cameraConnections": camera_connections } await websocket.send_text(json.dumps(state_report)) logging.debug("Sent stateReport as heartbeat") await asyncio.sleep(HEARTBEAT_INTERVAL) except Exception as e: logging.error(f"Error sending stateReport heartbeat: {e}") break async def on_message(): while True: try: msg = await websocket.receive_text() logging.debug(f"Received message: {msg}") data = json.loads(msg) msg_type = data.get("type") if msg_type == "subscribe": payload = data.get("payload", {}) camera_id = payload.get("cameraIdentifier") rtsp_url = payload.get("rtspUrl") model_url = payload.get("modelUrl") # may be remote or local modelId = payload.get("modelId") modelName = payload.get("modelName") if model_url: with models_lock: if camera_id not in models: models[camera_id] = {} if modelId not in models[camera_id]: logging.info(f"Loading model from {model_url}") extraction_dir = os.path.join("models", camera_id, str(modelId)) os.makedirs(extraction_dir, exist_ok=True) # If model_url is remote, download it first. parsed = urlparse(model_url) if parsed.scheme in ("http", "https"): local_mpta = os.path.join(extraction_dir, os.path.basename(parsed.path)) local_path = download_mpta(model_url, local_mpta) if not local_path: logging.error("Failed to download the remote mpta file.") continue model_tree = load_pipeline_from_zip(local_path, extraction_dir) else: model_tree = load_pipeline_from_zip(model_url, extraction_dir) if model_tree is None: logging.error("Failed to load model from mpta file.") continue models[camera_id][modelId] = model_tree logging.info(f"Loaded model {modelId} for camera {camera_id}") if camera_id and rtsp_url: with streams_lock: if camera_id not in streams and len(streams) < max_streams: cap = cv2.VideoCapture(rtsp_url) if not cap.isOpened(): logging.error(f"Failed to open RTSP stream for camera {camera_id}") continue buffer = queue.Queue(maxsize=1) stop_event = threading.Event() thread = threading.Thread(target=frame_reader, args=(camera_id, cap, buffer, stop_event)) thread.daemon = True thread.start() streams[camera_id] = { "cap": cap, "buffer": buffer, "thread": thread, "rtsp_url": rtsp_url, "stop_event": stop_event, "modelId": modelId, "modelName": modelName } logging.info(f"Subscribed to camera {camera_id} with modelId {modelId}, modelName {modelName}, URL {rtsp_url}") elif camera_id and camera_id in streams: # If already subscribed, unsubscribe stream = streams.pop(camera_id) stream["cap"].release() logging.info(f"Unsubscribed from camera {camera_id}") with models_lock: if camera_id in models and modelId in models[camera_id]: del models[camera_id][modelId] if not models[camera_id]: del models[camera_id] elif msg_type == "unsubscribe": payload = data.get("payload", {}) camera_id = payload.get("cameraIdentifier") logging.debug(f"Unsubscribing from camera {camera_id}") with streams_lock: if camera_id and camera_id in streams: stream = streams.pop(camera_id) stream["stop_event"].set() stream["thread"].join() stream["cap"].release() logging.info(f"Unsubscribed from camera {camera_id}") with models_lock: if camera_id in models: del models[camera_id] elif msg_type == "requestState": cpu_usage = psutil.cpu_percent() memory_usage = psutil.virtual_memory().percent if torch.cuda.is_available(): gpu_usage = torch.cuda.memory_allocated() / (1024 ** 2) gpu_memory_usage = torch.cuda.memory_reserved() / (1024 ** 2) else: gpu_usage = None gpu_memory_usage = None camera_connections = [ { "cameraIdentifier": camera_id, "modelId": stream["modelId"], "modelName": stream["modelName"], "online": True } for camera_id, stream in streams.items() ] state_report = { "type": "stateReport", "cpuUsage": cpu_usage, "memoryUsage": memory_usage, "gpuUsage": gpu_usage, "gpuMemoryUsage": gpu_memory_usage, "cameraConnections": camera_connections } await websocket.send_text(json.dumps(state_report)) else: logging.error(f"Unknown message type: {msg_type}") except json.JSONDecodeError: logging.error("Received invalid JSON message") except (WebSocketDisconnect, ConnectionClosedError) as e: logging.warning(f"WebSocket disconnected: {e}") break except Exception as e: logging.error(f"Error handling message: {e}") break try: await websocket.accept() stream_task = asyncio.create_task(process_streams()) heartbeat_task = asyncio.create_task(send_heartbeat()) message_task = asyncio.create_task(on_message()) await asyncio.gather(heartbeat_task, message_task) except Exception as e: logging.error(f"Error in detect websocket: {e}") finally: stream_task.cancel() await stream_task with streams_lock: for camera_id, stream in streams.items(): stream["stop_event"].set() stream["thread"].join() stream["cap"].release() while not stream["buffer"].empty(): try: stream["buffer"].get_nowait() except queue.Empty: pass logging.info(f"Released camera {camera_id} and cleaned up resources") streams.clear() with models_lock: models.clear() logging.info("WebSocket connection closed")