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

thread safety

Browse source
This commit is contained in:
Siwat Sirichai 2025-01-14 23:54:07 +07:00
parent e52efabbb7
commit ffe9c90747
1 changed files with 199 additions and 240 deletions
Download patch file Download diff file Expand all files Collapse all files

439
app.py
View file

@ -34,7 +34,7 @@ max_retries = config.get("max_retries", 3)
# Configure logging # Configure logging
logging.basicConfig( logging.basicConfig(
level=logging.INFO, level=logging.DEBUG,
format="%(asctime)s [%(levelname)s] %(message)s", format="%(asctime)s [%(levelname)s] %(message)s",
handlers=[ handlers=[
logging.FileHandler("app.log"), logging.FileHandler("app.log"),
@ -49,6 +49,10 @@ os.makedirs("models", exist_ok=True)
HEARTBEAT_INTERVAL = 2 # seconds HEARTBEAT_INTERVAL = 2 # seconds
WORKER_TIMEOUT_MS = 10000 WORKER_TIMEOUT_MS = 10000
# Add a lock for thread-safe operations on shared resources
streams_lock = threading.Lock()
models_lock = threading.Lock()
@app.websocket("/") @app.websocket("/")
async def detect(websocket: WebSocket): async def detect(websocket: WebSocket):
import asyncio import asyncio
@ -61,77 +65,84 @@ async def detect(websocket: WebSocket):
# This function is user-modifiable # This function is user-modifiable
# Save data you want to persist across frames in the persistent_data dictionary # Save data you want to persist across frames in the persistent_data dictionary
async def handle_detection(camera_id, stream, frame, websocket, model: YOLO, persistent_data): async def handle_detection(camera_id, stream, frame, websocket, model: YOLO, persistent_data):
boxes = [] try:
for r in model.track(frame, stream=False, persist=True): boxes = []
for box in r.boxes: for r in model.track(frame, stream=False, persist=True):
track_id = None for box in r.boxes:
if hasattr(box, "id") and box.id is not None: track_id = None
track_id = box.id.item() if hasattr(box, "id") and box.id is not None:
box_cpu = box.cpu() track_id = box.id.item()
boxes.append({ box_cpu = box.cpu()
"class": model.names[int(box_cpu.cls[0])], boxes.append({
"confidence": float(box_cpu.conf[0]), "class": model.names[int(box_cpu.cls[0])],
"id": track_id, "confidence": float(box_cpu.conf[0]),
}) "id": track_id,
# 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": {
"modelId": stream['modelId'], "detections": boxes,
"modelName": stream['modelName'] "modelId": stream['modelId'],
"modelName": stream['modelName']
}
} }
} 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) return persistent_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): def frame_reader(camera_id, cap, buffer, stop_event):
import time import time
retries = 0 retries = 0
while not stop_event.is_set(): try:
try: while not stop_event.is_set():
ret, frame = cap.read() try:
if not ret: ret, frame = cap.read()
logging.warning(f"Connection lost for camera: {camera_id}, retry {retries+1}/{max_retries}") 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() cap.release()
time.sleep(reconnect_interval) time.sleep(reconnect_interval)
retries += 1 retries += 1
if retries > max_retries: if retries > max_retries and max_retries != -1:
logging.error(f"Max retries reached for camera: {camera_id}") logging.error(f"Max retries reached after OpenCV error for camera: {camera_id}")
break break
# Re-open the VideoCapture # Re-open the VideoCapture
cap = cv2.VideoCapture(streams[camera_id]['rtsp_url']) cap = cv2.VideoCapture(streams[camera_id]['rtsp_url'])
if not cap.isOpened(): if not cap.isOpened():
logging.error(f"Failed to reopen RTSP stream for camera: {camera_id}") logging.error(f"Failed to reopen RTSP stream for camera {camera_id} after OpenCV error")
continue continue
continue except Exception as e:
retries = 0 # Reset on success logging.error(f"Unexpected error for camera {camera_id}: {e}")
if not buffer.empty(): cap.release()
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 and max_retries != -1:
logging.error(f"Max retries reached after OpenCV error for camera: {camera_id}")
break break
# Re-open the VideoCapture except Exception as e:
cap = cv2.VideoCapture(streams[camera_id]['rtsp_url']) logging.error(f"Error in frame_reader thread for camera {camera_id}: {e}")
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
async def process_streams(): async def process_streams():
global models global models
@ -141,11 +152,14 @@ async def detect(websocket: WebSocket):
while True: while True:
start_time = time.time() start_time = time.time()
# Round-robin processing # Round-robin processing
for camera_id, stream in list(streams.items()): with streams_lock:
current_streams = list(streams.items())
for camera_id, stream in current_streams:
buffer = stream['buffer'] buffer = stream['buffer']
if not buffer.empty(): if not buffer.empty():
frame = buffer.get() frame = buffer.get()
model = models.get(camera_id, {}).get(stream['modelId']) with models_lock:
model = models.get(camera_id, {}).get(stream['modelId'])
key = (camera_id, stream['modelId']) key = (camera_id, stream['modelId'])
persistent_data = persistent_data_dict.get(key, {}) persistent_data = persistent_data_dict.get(key, {})
updated_persistent_data = await handle_detection(camera_id, stream, frame, websocket, model, persistent_data) updated_persistent_data = await handle_detection(camera_id, stream, frame, websocket, model, persistent_data)
@ -198,185 +212,120 @@ async def detect(websocket: WebSocket):
async def on_message(): async def on_message():
global models global models
while True:
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")
modelId = payload.get("modelId")
modelName = payload.get("modelName")
if model_url:
if camera_id not in models:
models[camera_id] = {}
if modelId not in models[camera_id]:
print(f"Downloading model from {model_url}")
parsed_url = urlparse(model_url)
filename = os.path.basename(parsed_url.path)
model_filename = os.path.join("models", filename)
# Download the model
response = requests.get(model_url, stream=True)
if response.status_code == 200:
with open(model_filename, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
logging.info(f"Downloaded model from {model_url} to {model_filename}")
model = YOLO(model_filename)
if torch.cuda.is_available():
model.to('cuda')
models[camera_id][modelId] = model
logging.info(f"Loaded model {modelId} for camera {camera_id}")
else:
logging.error(f"Failed to download model from {model_url}")
continue
if camera_id and rtsp_url:
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} and 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}")
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")
if camera_id and camera_id in streams:
stream = streams.pop(camera_id)
stream['cap'].release()
logging.info(f"Unsubscribed from camera {camera_id}")
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 == "requestState":
# Handle state request
cpu_usage = psutil.cpu_percent()
memory_usage = psutil.virtual_memory().percent
if torch.cuda.is_available():
gpu_usage = torch.cuda.memory_allocated() / (1024 ** 2) # Convert to MB
gpu_memory_usage = torch.cuda.memory_reserved() / (1024 ** 2) # Convert to 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))
else:
logging.error(f"Unknown message type: {msg_type}")
await websocket.accept()
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)
model = None
model_path = None
try:
while True: while True:
try: try:
msg = await websocket.receive_text() msg = await websocket.receive_text()
logging.debug(f"Received message: {msg}") logging.debug(f"Received message: {msg}")
data = json.loads(msg) data = json.loads(msg)
camera_id = data.get("cameraIdentifier") msg_type = data.get("type")
rtsp_url = data.get("rtspUrl")
model_url = data.get("modelUrl")
modelId = data.get("modelId")
modelName = data.get("modelName")
if model_url: if msg_type == "subscribe":
print(f"Downloading model from {model_url}") payload = data.get("payload", {})
parsed_url = urlparse(model_url) camera_id = payload.get("cameraIdentifier")
filename = os.path.basename(parsed_url.path) rtsp_url = payload.get("rtspUrl")
model_filename = os.path.join("models", filename) model_url = payload.get("modelUrl")
# Download the model modelId = payload.get("modelId")
response = requests.get(model_url, stream=True) modelName = payload.get("modelName")
if response.status_code == 200:
with open(model_filename, 'wb') as f: if model_url:
for chunk in response.iter_content(chunk_size=8192): with models_lock:
f.write(chunk) if camera_id not in models:
logging.info(f"Downloaded model from {model_url} to {model_filename}") models[camera_id] = {}
model = YOLO(model_filename) if modelId not in models[camera_id]:
if torch.cuda.is_available(): print(f"Downloading model from {model_url}")
model.to('cuda') parsed_url = urlparse(model_url)
class_names = model.names filename = os.path.basename(parsed_url.path)
model_filename = os.path.join("models", filename)
# Download the model
response = requests.get(model_url, stream=True)
if response.status_code == 200:
with open(model_filename, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
logging.info(f"Downloaded model from {model_url} to {model_filename}")
model = YOLO(model_filename)
if torch.cuda.is_available():
model.to('cuda')
models[camera_id][modelId] = model
logging.info(f"Loaded model {modelId} for camera {camera_id}")
else:
logging.error(f"Failed to download model from {model_url}")
continue
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} and 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}")
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['cap'].release()
logging.info(f"Unsubscribed from camera {camera_id}")
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 == "requestState":
# Handle state request
cpu_usage = psutil.cpu_percent()
memory_usage = psutil.virtual_memory().percent
if torch.cuda.is_available():
gpu_usage = torch.cuda.memory_allocated() / (1024 ** 2) # Convert to MB
gpu_memory_usage = torch.cuda.memory_reserved() / (1024 ** 2) # Convert to MB
else: else:
logging.error(f"Failed to download model from {model_url}") gpu_usage = None
continue gpu_memory_usage = None
if camera_id and rtsp_url:
if camera_id not in streams and len(streams) < max_streams: camera_connections = [
cap = cv2.VideoCapture(rtsp_url) {
if not cap.isOpened(): "cameraIdentifier": camera_id,
logging.error(f"Failed to open RTSP stream for camera {camera_id}") "modelId": stream['modelId'],
continue "modelName": stream['modelName'],
buffer = queue.Queue(maxsize=1) "online": True
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} and URL {rtsp_url}") for camera_id, stream in streams.items()
elif camera_id and camera_id in streams: ]
stream = streams.pop(camera_id)
stream['cap'].release() state_report = {
logging.info(f"Unsubscribed from camera {camera_id}") "type": "stateReport",
elif data.get("command") == "stop": "cpuUsage": cpu_usage,
logging.info("Received stop command") "memoryUsage": memory_usage,
break "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: except json.JSONDecodeError:
logging.error("Received invalid JSON message") logging.error("Received invalid JSON message")
except (WebSocketDisconnect, ConnectionClosedError) as e: except (WebSocketDisconnect, ConnectionClosedError) as e:
@ -385,17 +334,27 @@ async def detect(websocket: WebSocket):
except Exception as e: except Exception as e:
logging.error(f"Error handling message: {e}") logging.error(f"Error handling message: {e}")
break break
try:
await websocket.accept()
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: except Exception as e:
logging.error(f"Unexpected error in WebSocket connection: {e}") logging.error(f"Error in detect websocket: {e}")
finally: finally:
task.cancel() task.cancel()
await task await task
for camera_id, stream in streams.items(): with streams_lock:
stream['stop_event'].set() for camera_id, stream in streams.items():
stream['thread'].join() stream['stop_event'].set()
stream['cap'].release() stream['thread'].join()
stream['buffer'].queue.clear() stream['cap'].release()
logging.info(f"Released camera {camera_id} and cleaned up resources") stream['buffer'].queue.clear()
streams.clear() logging.info(f"Released camera {camera_id} and cleaned up resources")
models.clear() streams.clear()
with models_lock:
models.clear()
logging.info("WebSocket connection closed") logging.info("WebSocket connection closed")