adsist-cms/python-detector-worker
3
0
Fork 0
Code Issues Pull requests Projects Releases Packages 1 Wiki Activity Actions

gracefully stop WS connection

Browse source
This commit is contained in:
Siwat Sirichai 2025-01-11 01:29:07 +07:00
parent d5e0ebbe9d
commit af26c1477c
2 changed files with 130 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

192
app.py
View file

@ -1,4 +1,6 @@
from fastapi import FastAPI, WebSocket from fastapi import FastAPI, WebSocket
from fastapi.websockets import WebSocketDisconnect
from websockets.exceptions import ConnectionClosedError
from ultralytics import YOLO from ultralytics import YOLO
import torch import torch
import cv2 import cv2
@ -22,10 +24,12 @@ with open("config.json", "r") as f:
config = json.load(f) config = json.load(f)
poll_interval = config.get("poll_interval_ms", 100) poll_interval = config.get("poll_interval_ms", 100)
reconnect_interval = config.get("reconnect_interval_sec", 5) # New setting
TARGET_FPS = config.get("target_fps", 10) # Add TARGET_FPS TARGET_FPS = config.get("target_fps", 10) # Add TARGET_FPS
poll_interval = 1000 / TARGET_FPS # Adjust poll_interval based on TARGET_FPS poll_interval = 1000 / TARGET_FPS # Adjust poll_interval based on TARGET_FPS
logging.info(f"Poll interval: {poll_interval}ms") logging.info(f"Poll interval: {poll_interval}ms")
max_streams = config.get("max_streams", 5) max_streams = config.get("max_streams", 5)
max_retries = config.get("max_retries", 3)
# Configure logging # Configure logging
logging.basicConfig( logging.basicConfig(
@ -44,89 +48,147 @@ async def detect(websocket: WebSocket):
logging.info("WebSocket connection accepted") logging.info("WebSocket connection accepted")
# Move streams inside the detect function streams = {}
streams = {} # Modify to store {camera_id: {'cap': VideoCapture, 'buffer': Queue, 'thread': Thread}}
def frame_reader(camera_id, cap, buffer): def frame_reader(camera_id, cap, buffer, stop_event):
while True: import time
ret, frame = cap.read() retries = 0
if not ret: while not stop_event.is_set():
logging.warning(f"Failed to read frame from camera: {camera_id}") 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:
logging.error(f"Max retries reached for camera: {camera_id}")
break
# Re-open the VideoCapture
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 # Reset on success
if not buffer.empty():
try:
buffer.get_nowait() # Discard the old frame
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:
logging.error(f"Max retries reached after OpenCV error for camera: {camera_id}")
break
# Re-open the VideoCapture
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 break
if not buffer.empty():
try:
buffer.get_nowait() # Discard the old frame
except queue.Empty:
pass
buffer.put(frame)
async def process_streams(): async def process_streams():
logging.info("Started processing streams") logging.info("Started processing streams")
while True: try:
start_time = time.time() while True:
# Round-robin processing start_time = time.time()
for camera_id, stream in list(streams.items()): # Round-robin processing
buffer = stream['buffer'] for camera_id, stream in list(streams.items()):
if not buffer.empty(): buffer = stream['buffer']
frame = buffer.get() if not buffer.empty():
results = model(frame, stream=False) frame = buffer.get()
boxes = [] results = model(frame, stream=False)
for r in results: boxes = []
for box in r.boxes: for r in results:
boxes.append({ for box in r.boxes:
"class": class_names[int(box.cls[0])], # Use class name boxes.append({
"confidence": float(box.conf[0]), "class": class_names[int(box.cls[0])],
}) "confidence": float(box.conf[0]),
# Broadcast to all subscribers of this URL })
detection_data = { # Broadcast to all subscribers of this URL
"type": "imageDetection", detection_data = {
"cameraIdentifier": camera_id, "type": "imageDetection",
"timestamp": time.time(), "cameraIdentifier": camera_id,
"data": { "timestamp": time.time(),
"detections": boxes "data": {
"detections": boxes
}
} }
} logging.debug(f"Sending detection data for camera {camera_id}: {detection_data}")
logging.debug(f"Sending detection data for camera {camera_id}: {detection_data}") await websocket.send_json(detection_data)
await websocket.send_json(detection_data) elapsed_time = (time.time() - start_time) * 1000 # in ms
elapsed_time = (time.time() - start_time) * 1000 # in ms sleep_time = max(poll_interval - elapsed_time, 0)
sleep_time = max(poll_interval - elapsed_time, 0) logging.debug(f"Elapsed time: {elapsed_time}ms, sleeping for: {sleep_time}ms")
logging.debug(f"Elapsed time: {elapsed_time}ms, sleeping for: {sleep_time}ms") await asyncio.sleep(sleep_time / 1000.0)
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}")
await websocket.accept() await websocket.accept()
task = asyncio.create_task(process_streams()) task = asyncio.create_task(process_streams())
try: try:
while True: while True:
msg = await websocket.receive_text() try:
logging.debug(f"Received message: {msg}") msg = await websocket.receive_text()
data = json.loads(msg) logging.debug(f"Received message: {msg}")
camera_id = data.get("cameraIdentifier") data = json.loads(msg)
rtsp_url = data.get("rtspUrl") camera_id = data.get("cameraIdentifier")
rtsp_url = data.get("rtspUrl")
if camera_id and rtsp_url: if camera_id and rtsp_url:
if camera_id not in streams and len(streams) < max_streams: if camera_id not in streams and len(streams) < max_streams:
cap = cv2.VideoCapture(rtsp_url) cap = cv2.VideoCapture(rtsp_url)
buffer = queue.Queue(maxsize=1) if not cap.isOpened():
thread = threading.Thread(target=frame_reader, args=(camera_id, cap, buffer)) logging.error(f"Failed to open RTSP stream for camera {camera_id}")
thread.daemon = True continue
thread.start() buffer = queue.Queue(maxsize=1)
streams[camera_id] = {'cap': cap, 'buffer': buffer, 'thread': thread} stop_event = threading.Event()
logging.info(f"Subscribed to camera {camera_id} with URL {rtsp_url}") thread = threading.Thread(target=frame_reader, args=(camera_id, cap, buffer, stop_event))
elif camera_id and camera_id in streams: thread.daemon = True
stream = streams.pop(camera_id) thread.start()
stream['cap'].release() streams[camera_id] = {
logging.info(f"Unsubscribed from camera {camera_id}") 'cap': cap,
elif data.get("command") == "stop": 'buffer': buffer,
logging.info("Received stop command") 'thread': thread,
'rtsp_url': rtsp_url,
'stop_event': stop_event
}
logging.info(f"Subscribed to camera {camera_id} with URL {rtsp_url}")
elif camera_id and camera_id in streams:
stream = streams.pop(camera_id)
stream['cap'].release()
logging.info(f"Unsubscribed from camera {camera_id}")
elif data.get("command") == "stop":
logging.info("Received stop command")
break
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 break
except Exception as e: except Exception as e:
logging.error(f"Error in WebSocket connection: {e}") logging.error(f"Unexpected error in WebSocket connection: {e}")
finally: finally:
task.cancel() task.cancel()
await task
for camera_id, stream in streams.items(): for camera_id, stream in streams.items():
stream['stop_event'].set()
stream['thread'].join()
stream['cap'].release() stream['cap'].release()
logging.info(f"Released camera {camera_id}") stream['buffer'].queue.clear()
logging.info(f"Released camera {camera_id} and cleaned up resources")
streams.clear() streams.clear()
logging.info("WebSocket connection closed") logging.info("WebSocket connection closed")
await websocket.close()

4
config.json
View file

@ -1,5 +1,7 @@
{ {
"poll_interval_ms": 100, "poll_interval_ms": 100,
"max_streams": 5, "max_streams": 5,
"target_fps": 2 "target_fps": 2,
"reconnect_interval_sec": 5,
"max_retries": 3
} }