thread safety
This commit is contained in:
parent
e52efabbb7
commit
ffe9c90747
1 changed files with 199 additions and 240 deletions
105
app.py
105
app.py
|
|
@ -34,7 +34,7 @@ max_retries = config.get("max_retries", 3)
|
|||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
level=logging.DEBUG,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||
handlers=[
|
||||
logging.FileHandler("app.log"),
|
||||
|
|
@ -49,6 +49,10 @@ os.makedirs("models", exist_ok=True)
|
|||
HEARTBEAT_INTERVAL = 2 # seconds
|
||||
WORKER_TIMEOUT_MS = 10000
|
||||
|
||||
# Add a lock for thread-safe operations on shared resources
|
||||
streams_lock = threading.Lock()
|
||||
models_lock = threading.Lock()
|
||||
|
||||
@app.websocket("/")
|
||||
async def detect(websocket: WebSocket):
|
||||
import asyncio
|
||||
|
|
@ -61,6 +65,7 @@ async def detect(websocket: WebSocket):
|
|||
# This function is user-modifiable
|
||||
# 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):
|
||||
try:
|
||||
boxes = []
|
||||
for r in model.track(frame, stream=False, persist=True):
|
||||
for box in r.boxes:
|
||||
|
|
@ -87,10 +92,14 @@ async def detect(websocket: WebSocket):
|
|||
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):
|
||||
import time
|
||||
retries = 0
|
||||
try:
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
ret, frame = cap.read()
|
||||
|
|
@ -132,6 +141,8 @@ async def detect(websocket: WebSocket):
|
|||
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():
|
||||
global models
|
||||
|
|
@ -141,10 +152,13 @@ async def detect(websocket: WebSocket):
|
|||
while True:
|
||||
start_time = time.time()
|
||||
# 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']
|
||||
if not buffer.empty():
|
||||
frame = buffer.get()
|
||||
with models_lock:
|
||||
model = models.get(camera_id, {}).get(stream['modelId'])
|
||||
key = (camera_id, stream['modelId'])
|
||||
persistent_data = persistent_data_dict.get(key, {})
|
||||
|
|
@ -199,6 +213,7 @@ async def detect(websocket: WebSocket):
|
|||
async def on_message():
|
||||
global models
|
||||
while True:
|
||||
try:
|
||||
msg = await websocket.receive_text()
|
||||
logging.debug(f"Received message: {msg}")
|
||||
data = json.loads(msg)
|
||||
|
|
@ -213,6 +228,7 @@ async def detect(websocket: WebSocket):
|
|||
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]:
|
||||
|
|
@ -236,6 +252,7 @@ async def detect(websocket: WebSocket):
|
|||
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():
|
||||
|
|
@ -267,6 +284,8 @@ async def detect(websocket: WebSocket):
|
|||
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()
|
||||
|
|
@ -307,76 +326,6 @@ async def detect(websocket: WebSocket):
|
|||
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:
|
||||
try:
|
||||
msg = await websocket.receive_text()
|
||||
logging.debug(f"Received message: {msg}")
|
||||
data = json.loads(msg)
|
||||
camera_id = data.get("cameraIdentifier")
|
||||
rtsp_url = data.get("rtspUrl")
|
||||
model_url = data.get("modelUrl")
|
||||
modelId = data.get("modelId")
|
||||
modelName = data.get("modelName")
|
||||
|
||||
if model_url:
|
||||
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')
|
||||
class_names = model.names
|
||||
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}")
|
||||
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:
|
||||
|
|
@ -385,11 +334,20 @@ async def detect(websocket: WebSocket):
|
|||
except Exception as e:
|
||||
logging.error(f"Error handling message: {e}")
|
||||
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:
|
||||
logging.error(f"Unexpected error in WebSocket connection: {e}")
|
||||
logging.error(f"Error in detect websocket: {e}")
|
||||
finally:
|
||||
task.cancel()
|
||||
await task
|
||||
with streams_lock:
|
||||
for camera_id, stream in streams.items():
|
||||
stream['stop_event'].set()
|
||||
stream['thread'].join()
|
||||
|
|
@ -397,5 +355,6 @@ async def detect(websocket: WebSocket):
|
|||
stream['buffer'].queue.clear()
|
||||
logging.info(f"Released camera {camera_id} and cleaned up resources")
|
||||
streams.clear()
|
||||
with models_lock:
|
||||
models.clear()
|
||||
logging.info("WebSocket connection closed")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue