diff --git a/app.py b/app.py index 4e9be15..60beb27 100644 --- a/app.py +++ b/app.py @@ -29,6 +29,12 @@ app = FastAPI() # "models" now holds a nested dict: { camera_id: { modelId: model_tree } } models: Dict[str, Dict[str, Any]] = {} streams: Dict[str, Dict[str, Any]] = {} +# Store session IDs per display +session_ids: Dict[str, int] = {} +# Track shared camera streams by camera URL +camera_streams: Dict[str, Dict[str, Any]] = {} +# Map subscriptions to their camera URL +subscription_to_camera: Dict[str, str] = {} with open("config.json", "r") as f: config = json.load(f) @@ -184,9 +190,16 @@ async def detect(websocket: WebSocket): async def handle_detection(camera_id, stream, frame, websocket, model_tree, persistent_data): try: + # Apply crop if specified + cropped_frame = frame + if all(coord is not None for coord in [stream.get("cropX1"), stream.get("cropY1"), stream.get("cropX2"), stream.get("cropY2")]): + cropX1, cropY1, cropX2, cropY2 = stream["cropX1"], stream["cropY1"], stream["cropX2"], stream["cropY2"] + cropped_frame = frame[cropY1:cropY2, cropX1:cropX2] + logger.debug(f"Applied crop coordinates ({cropX1}, {cropY1}, {cropX2}, {cropY2}) to frame for camera {camera_id}") + logger.debug(f"Processing frame for camera {camera_id} with model {stream['modelId']}") start_time = time.time() - detection_result = run_pipeline(frame, model_tree) + detection_result = run_pipeline(cropped_frame, model_tree) process_time = (time.time() - start_time) * 1000 logger.debug(f"Detection for camera {camera_id} completed in {process_time:.2f}ms") @@ -235,22 +248,48 @@ async def detect(websocket: WebSocket): "box": [0, 0, 0, 0] } + # Convert detection format to match protocol - flatten detection attributes + detection_dict = {} + + # Handle different detection result formats + if isinstance(highest_confidence_detection, dict): + # Copy all fields from the detection result + for key, value in highest_confidence_detection.items(): + if key not in ["box", "id"]: # Skip internal fields + detection_dict[key] = value + + # Extract display identifier for session ID lookup + subscription_parts = stream["subscriptionIdentifier"].split(';') + display_identifier = subscription_parts[0] if subscription_parts else None + session_id = session_ids.get(display_identifier) if display_identifier else None + detection_data = { "type": "imageDetection", "subscriptionIdentifier": stream["subscriptionIdentifier"], "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S.%fZ", time.gmtime()), "data": { - "detection": highest_confidence_detection, # Send only the highest confidence detection + "detection": detection_dict, "modelId": stream["modelId"], "modelName": stream["modelName"] } } + # Add session ID if available + if session_id is not None: + detection_data["sessionId"] = session_id + if highest_confidence_detection["class"] != "none": logger.info(f"Camera {camera_id}: Detected {highest_confidence_detection['class']} with confidence {highest_confidence_detection['confidence']:.2f} using model {stream['modelName']}") + + # Log session ID if available + subscription_parts = stream["subscriptionIdentifier"].split(';') + display_identifier = subscription_parts[0] if subscription_parts else None + session_id = session_ids.get(display_identifier) if display_identifier else None + if session_id: + logger.debug(f"Detection associated with session ID: {session_id}") await websocket.send_json(detection_data) - logger.debug(f"Sent detection data to client for camera {camera_id}:\n{json.dumps(detection_data, indent=2)}") + logger.debug(f"Sent detection data to client for camera {camera_id}") return persistent_data except Exception as e: logger.error(f"Error in handle_detection for camera {camera_id}: {str(e)}", exc_info=True) @@ -521,50 +560,58 @@ async def detect(websocket: WebSocket): cropX2 = payload.get("cropX2") cropY2 = payload.get("cropY2") - camera_id = subscriptionIdentifier # Use subscriptionIdentifier as camera_id for mapping + # Extract camera_id from subscriptionIdentifier (format: displayIdentifier;cameraIdentifier) + parts = subscriptionIdentifier.split(';') + if len(parts) != 2: + logger.error(f"Invalid subscriptionIdentifier format: {subscriptionIdentifier}") + continue + + display_identifier, camera_identifier = parts + camera_id = subscriptionIdentifier # Use full subscriptionIdentifier as camera_id for mapping if model_url: with models_lock: if (camera_id not in models) or (modelId not in models[camera_id]): logger.info(f"Loading model from {model_url} for camera {camera_id}, modelId {modelId}") - extraction_dir = os.path.join("models", camera_id, str(modelId)) + extraction_dir = os.path.join("models", camera_identifier, 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"): - logger.info(f"Downloading remote model from {model_url}") - local_mpta = os.path.join(extraction_dir, os.path.basename(parsed.path)) + logger.info(f"Downloading remote .mpta file from {model_url}") + filename = os.path.basename(parsed.path) or f"model_{modelId}.mpta" + local_mpta = os.path.join(extraction_dir, filename) logger.debug(f"Download destination: {local_mpta}") local_path = download_mpta(model_url, local_mpta) if not local_path: - logger.error(f"Failed to download the remote mpta file from {model_url}") + logger.error(f"Failed to download the remote .mpta file from {model_url}") error_response = { "type": "error", - "cameraIdentifier": camera_id, + "subscriptionIdentifier": subscriptionIdentifier, "error": f"Failed to download model from {model_url}" } await websocket.send_json(error_response) continue model_tree = load_pipeline_from_zip(local_path, extraction_dir) else: - logger.info(f"Loading local model from {model_url}") + logger.info(f"Loading local .mpta file from {model_url}") # Check if file exists before attempting to load if not os.path.exists(model_url): - logger.error(f"Local model file not found: {model_url}") + logger.error(f"Local .mpta file not found: {model_url}") logger.debug(f"Current working directory: {os.getcwd()}") error_response = { "type": "error", - "cameraIdentifier": camera_id, + "subscriptionIdentifier": subscriptionIdentifier, "error": f"Model file not found: {model_url}" } await websocket.send_json(error_response) continue model_tree = load_pipeline_from_zip(model_url, extraction_dir) if model_tree is None: - logger.error(f"Failed to load model {modelId} from mpta file for camera {camera_id}") + logger.error(f"Failed to load model {modelId} from .mpta file for camera {camera_id}") error_response = { "type": "error", - "cameraIdentifier": camera_id, + "subscriptionIdentifier": subscriptionIdentifier, "error": f"Failed to load model {modelId}" } await websocket.send_json(error_response) @@ -573,20 +620,80 @@ async def detect(websocket: WebSocket): models[camera_id] = {} models[camera_id][modelId] = model_tree logger.info(f"Successfully loaded model {modelId} for camera {camera_id}") - success_response = { - "type": "modelLoaded", - "cameraIdentifier": camera_id, - "modelId": modelId - } - await websocket.send_json(success_response) + logger.debug(f"Model extraction directory: {extraction_dir}") if camera_id and (rtsp_url or snapshot_url): with streams_lock: + # Determine camera URL for shared stream management + camera_url = snapshot_url if snapshot_url else rtsp_url + if camera_id not in streams and len(streams) < max_streams: - buffer = queue.Queue(maxsize=1) - stop_event = threading.Event() + # Check if we already have a stream for this camera URL + shared_stream = camera_streams.get(camera_url) + + if shared_stream: + # Reuse existing stream + logger.info(f"Reusing existing stream for camera URL: {camera_url}") + buffer = shared_stream["buffer"] + stop_event = shared_stream["stop_event"] + thread = shared_stream["thread"] + mode = shared_stream["mode"] + + # Increment reference count + shared_stream["ref_count"] = shared_stream.get("ref_count", 0) + 1 + else: + # Create new stream + buffer = queue.Queue(maxsize=1) + stop_event = threading.Event() + + if snapshot_url and snapshot_interval: + logger.info(f"Creating new snapshot stream for camera {camera_id}: {snapshot_url}") + thread = threading.Thread(target=snapshot_reader, args=(camera_identifier, snapshot_url, snapshot_interval, buffer, stop_event)) + thread.daemon = True + thread.start() + mode = "snapshot" + + # Store shared stream info + shared_stream = { + "buffer": buffer, + "thread": thread, + "stop_event": stop_event, + "mode": mode, + "url": snapshot_url, + "snapshot_interval": snapshot_interval, + "ref_count": 1 + } + camera_streams[camera_url] = shared_stream + + elif rtsp_url: + logger.info(f"Creating new RTSP stream for camera {camera_id}: {rtsp_url}") + cap = cv2.VideoCapture(rtsp_url) + if not cap.isOpened(): + logger.error(f"Failed to open RTSP stream for camera {camera_id}") + continue + thread = threading.Thread(target=frame_reader, args=(camera_identifier, cap, buffer, stop_event)) + thread.daemon = True + thread.start() + mode = "rtsp" + + # Store shared stream info + shared_stream = { + "buffer": buffer, + "thread": thread, + "stop_event": stop_event, + "mode": mode, + "url": rtsp_url, + "cap": cap, + "ref_count": 1 + } + camera_streams[camera_url] = shared_stream + else: + logger.error(f"No valid URL provided for camera {camera_id}") + continue + + # Create stream info for this subscription stream_info = { "buffer": buffer, - "thread": None, + "thread": thread, "stop_event": stop_event, "modelId": modelId, "modelName": modelName, @@ -594,52 +701,25 @@ async def detect(websocket: WebSocket): "cropX1": cropX1, "cropY1": cropY1, "cropX2": cropX2, - "cropY2": cropY2 + "cropY2": cropY2, + "mode": mode, + "camera_url": camera_url } - if snapshot_url and snapshot_interval: - logger.info(f"Using snapshot mode for camera {camera_id}: {snapshot_url}") - thread = threading.Thread(target=snapshot_reader, args=(camera_id, snapshot_url, snapshot_interval, buffer, stop_event)) - thread.daemon = True - thread.start() - stream_info.update({ - "snapshot_url": snapshot_url, - "snapshot_interval": snapshot_interval, - "mode": "snapshot" - }) - stream_info["thread"] = thread - streams[camera_id] = stream_info - elif rtsp_url: - logger.info(f"Using RTSP mode for camera {camera_id}: {rtsp_url}") - cap = cv2.VideoCapture(rtsp_url) - if not cap.isOpened(): - logger.error(f"Failed to open RTSP stream for camera {camera_id}") - continue - thread = threading.Thread(target=frame_reader, args=(camera_id, cap, buffer, stop_event)) - thread.daemon = True - thread.start() - stream_info.update({ - "cap": cap, - "rtsp_url": rtsp_url, - "mode": "rtsp" - }) - stream_info["thread"] = thread - streams[camera_id] = stream_info - else: - logger.error(f"No valid URL provided for camera {camera_id}") - continue + + if mode == "snapshot": + stream_info["snapshot_url"] = snapshot_url + stream_info["snapshot_interval"] = snapshot_interval + elif mode == "rtsp": + stream_info["rtsp_url"] = rtsp_url + stream_info["cap"] = shared_stream["cap"] + + streams[camera_id] = stream_info + subscription_to_camera[camera_id] = camera_url + elif camera_id and camera_id in streams: # If already subscribed, unsubscribe first - stream = streams.pop(camera_id) - stream["stop_event"].set() - stream["thread"].join() - if "cap" in stream: - stream["cap"].release() - logger.info(f"Unsubscribed from camera {camera_id} for resubscription") - 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] + logger.info(f"Resubscribing to camera {camera_id}") + # Note: Keep models in memory for reuse across subscriptions elif msg_type == "unsubscribe": payload = data.get("payload", {}) subscriptionIdentifier = payload.get("subscriptionIdentifier") @@ -647,13 +727,25 @@ async def detect(websocket: WebSocket): with streams_lock: if camera_id and camera_id in streams: stream = streams.pop(camera_id) - stream["stop_event"].set() - stream["thread"].join() - if "cap" in stream: - stream["cap"].release() - with models_lock: - if camera_id in models: - del models[camera_id] + camera_url = subscription_to_camera.pop(camera_id, None) + + if camera_url and camera_url in camera_streams: + shared_stream = camera_streams[camera_url] + shared_stream["ref_count"] -= 1 + + # If no more references, stop the shared stream + if shared_stream["ref_count"] <= 0: + logger.info(f"Stopping shared stream for camera URL: {camera_url}") + shared_stream["stop_event"].set() + shared_stream["thread"].join() + if "cap" in shared_stream: + shared_stream["cap"].release() + del camera_streams[camera_url] + else: + logger.info(f"Shared stream for {camera_url} still has {shared_stream['ref_count']} references") + + logger.info(f"Unsubscribed from camera {camera_id}") + # Note: Keep models in memory for potential reuse elif msg_type == "requestState": cpu_usage = psutil.cpu_percent() memory_usage = psutil.virtual_memory().percent @@ -684,6 +776,37 @@ async def detect(websocket: WebSocket): "cameraConnections": camera_connections } await websocket.send_text(json.dumps(state_report)) + + elif msg_type == "setSessionId": + payload = data.get("payload", {}) + display_identifier = payload.get("displayIdentifier") + session_id = payload.get("sessionId") + + if display_identifier: + # Store session ID for this display + if session_id is None: + session_ids.pop(display_identifier, None) + logger.info(f"Cleared session ID for display {display_identifier}") + else: + session_ids[display_identifier] = session_id + logger.info(f"Set session ID {session_id} for display {display_identifier}") + + elif msg_type == "patchSession": + session_id = data.get("sessionId") + patch_data = data.get("data", {}) + + # For now, just acknowledge the patch - actual implementation depends on backend requirements + response = { + "type": "patchSessionResult", + "payload": { + "sessionId": session_id, + "success": True, + "message": "Session patch acknowledged" + } + } + await websocket.send_json(response) + logger.info(f"Acknowledged patch for session {session_id}") + else: logger.error(f"Unknown message type: {msg_type}") except json.JSONDecodeError: @@ -706,19 +829,23 @@ async def detect(websocket: WebSocket): stream_task.cancel() await stream_task with streams_lock: - for camera_id, stream in streams.items(): - stream["stop_event"].set() - stream["thread"].join() - # Only release cap if it exists (RTSP mode) - if "cap" in stream: - stream["cap"].release() - while not stream["buffer"].empty(): + # Clean up shared camera streams + for camera_url, shared_stream in camera_streams.items(): + shared_stream["stop_event"].set() + shared_stream["thread"].join() + if "cap" in shared_stream: + shared_stream["cap"].release() + while not shared_stream["buffer"].empty(): try: - stream["buffer"].get_nowait() + shared_stream["buffer"].get_nowait() except queue.Empty: pass - logger.info(f"Released camera {camera_id} and cleaned up resources") + logger.info(f"Released shared camera stream for {camera_url}") + streams.clear() + camera_streams.clear() + subscription_to_camera.clear() with models_lock: models.clear() + session_ids.clear() logger.info("WebSocket connection closed") diff --git a/test_protocol.py b/test_protocol.py new file mode 100644 index 0000000..74af7d8 --- /dev/null +++ b/test_protocol.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +""" +Test script to verify the worker implementation follows the protocol +""" +import json +import asyncio +import websockets +import time + +async def test_protocol(): + """Test the worker protocol implementation""" + uri = "ws://localhost:8000" + + try: + async with websockets.connect(uri) as websocket: + print("✓ Connected to worker") + + # Test 1: Check if we receive heartbeat (stateReport) + print("\n1. Testing heartbeat...") + try: + message = await asyncio.wait_for(websocket.recv(), timeout=5) + data = json.loads(message) + if data.get("type") == "stateReport": + print("✓ Received stateReport heartbeat") + print(f" - CPU Usage: {data.get('cpuUsage', 'N/A')}%") + print(f" - Memory Usage: {data.get('memoryUsage', 'N/A')}%") + print(f" - Camera Connections: {len(data.get('cameraConnections', []))}") + else: + print(f"✗ Expected stateReport, got {data.get('type')}") + except asyncio.TimeoutError: + print("✗ No heartbeat received within 5 seconds") + + # Test 2: Request state + print("\n2. Testing requestState...") + await websocket.send(json.dumps({"type": "requestState"})) + try: + message = await asyncio.wait_for(websocket.recv(), timeout=5) + data = json.loads(message) + if data.get("type") == "stateReport": + print("✓ Received stateReport response") + else: + print(f"✗ Expected stateReport, got {data.get('type')}") + except asyncio.TimeoutError: + print("✗ No response to requestState within 5 seconds") + + # Test 3: Set session ID + print("\n3. Testing setSessionId...") + session_message = { + "type": "setSessionId", + "payload": { + "displayIdentifier": "display-001", + "sessionId": 12345 + } + } + await websocket.send(json.dumps(session_message)) + print("✓ Sent setSessionId message") + + # Test 4: Test patchSession + print("\n4. Testing patchSession...") + patch_message = { + "type": "patchSession", + "sessionId": 12345, + "data": { + "currentCar": { + "carModel": "Civic", + "carBrand": "Honda" + } + } + } + await websocket.send(json.dumps(patch_message)) + + # Wait for patchSessionResult + try: + message = await asyncio.wait_for(websocket.recv(), timeout=5) + data = json.loads(message) + if data.get("type") == "patchSessionResult": + print("✓ Received patchSessionResult") + print(f" - Success: {data.get('payload', {}).get('success')}") + print(f" - Message: {data.get('payload', {}).get('message')}") + else: + print(f"✗ Expected patchSessionResult, got {data.get('type')}") + except asyncio.TimeoutError: + print("✗ No patchSessionResult received within 5 seconds") + + # Test 5: Test subscribe message format (without actual camera) + print("\n5. Testing subscribe message format...") + subscribe_message = { + "type": "subscribe", + "payload": { + "subscriptionIdentifier": "display-001;cam-001", + "snapshotUrl": "http://example.com/snapshot.jpg", + "snapshotInterval": 5000, + "modelUrl": "http://example.com/model.mpta", + "modelName": "Test Model", + "modelId": 101, + "cropX1": 100, + "cropY1": 200, + "cropX2": 300, + "cropY2": 400 + } + } + await websocket.send(json.dumps(subscribe_message)) + print("✓ Sent subscribe message (will fail without actual camera/model)") + + # Listen for a few more messages to catch any errors + print("\n6. Listening for additional messages...") + for i in range(3): + try: + message = await asyncio.wait_for(websocket.recv(), timeout=2) + data = json.loads(message) + msg_type = data.get("type") + print(f" - Received {msg_type}") + if msg_type == "error": + print(f" Error: {data.get('error')}") + except asyncio.TimeoutError: + break + + print("\n✓ Protocol test completed successfully!") + + except Exception as e: + print(f"✗ Connection failed: {e}") + print("Make sure the worker is running on localhost:8000") + +if __name__ == "__main__": + asyncio.run(test_protocol()) \ No newline at end of file