feat: refactoring code, adding contaxtmanager, create pipeline

api/dependencies.py

@@ -15,4 +15,4 @@ def get_safety_detection_model(request: HTTPConnection):
 
 
 def get_redis(request: HTTPConnection):
-    return request.app.state.redis
+    return request.app.state.redis

api/routers/camera_stream.py

@@ -1,3 +1,5 @@
+from backend.utils.profiling import profile_step
+from backend.services.pipeline import ProcessingPipeline
 from domain.detection_box_center import calculate_detection_box_center
 from api.dependencies import get_safety_detection_model
 from api.dependencies import get_detection_model, get_depth_model
@@ -52,8 +54,9 @@ async def websocket_detect(
     logger.info(f"Client ID >>{camera_id}<< Connected...")
 
     step_counter = itertools.count()
+    pipeline = ProcessingPipeline(detector, depth_model, safety_detector, redis)
+
 
-    loop = asyncio.get_running_loop()
     # Queue removing old images in case they were being stacked
     frame_queue: asyncio.Queue = asyncio.Queue(maxsize=1)
 
@@ -78,106 +81,21 @@ async def websocket_detect(
         try:
             logger.info(f"Camera {camera_id} start sending frames...")
 
-            def decode_frame(fb):
-                return cv.imdecode(np.frombuffer(fb, np.uint8), cv.IMREAD_COLOR)
-
             # Keep receiving messages in a loop until disconnection.
             while True:
                 frame_bytes = await frame_queue.get()
 
-                # Profiling
-                t0 = time.time()
-                image_array = await loop.run_in_executor(
-                    None, decode_frame, frame_bytes
-                )
-                decode_duration_seconds.labels(camera_id).observe(
-                    round(time.time() - t0, 3)
-                )
-                mlflow.log_metric(
-                    "frame_processing_time",
-                    round(time.time() - t0, 3),
-                    next(step_counter),
-                )
-
-                # Apply detection models
-                t0 = time.time()
-                detection_task = loop.run_in_executor(
-                    None, detector.detect, image_array
-                )
-                safety_task = loop.run_in_executor(
-                    None, safety_detector.detect, image_array
-                )
-                detections, safety_detection = await asyncio.gather(
-                    detection_task, safety_task
-                )
-                detection_duration_seconds.labels(camera_id).observe(
-                    round(time.time() - t0, 3)
-                )
-                mlflow.log_metric(
-                    "detection_duration_seconds",
-                    round(time.time() - t0, 3),
-                    next(step_counter),
-                )
-
-                # Profiling
-                frame_processing_duration_seconds.labels(camera_id).observe(
-                    round(time.time() - t0, 3)
-                )
-                logger.debug("Frame processed", camera_id=camera_id)
-                mlflow.log_metric(
-                    "frame_processing duration time",
-                    round(time.time() - t0, 3),
-                    next(step_counter),
-                )
-
-                boxes_center, boxes_center_ratio = calculate_detection_box_center(detections.detections, image_array.shape[1])
-
-                t0 = time.time()
-                depth_points = (
-                    await loop.run_in_executor(
-                        None, depth_model.calculate_depth, image_array, boxes_center
-                    )
-                    if boxes_center
-                    else []
-                )
-                depth_duration_seconds.labels(camera_id).observe(
-                    round(time.time() - t0, 3)
-                )
-                mlflow.log_metric(
-                    "depth_duration_seconds",
-                    round(time.time() - t0, 3),
-                    next(step_counter),
-                )
-
-                detection_metadata = [
-                    DetectionMetadata(depth=depth, xRatio=xRatio)
-                    for depth, xRatio in zip(depth_points, boxes_center_ratio)
-                ]
-                metadata = CameraMetadata(
-                    camera_id=camera_id,
-                    is_danger=True if safety_detection else False,
-                    detection_metadata=detection_metadata,
-                )
-
-                await redis.publish("dashboard_stream", metadata.model_dump_json())
-                # Even if the camera was disconnected, redis is still going to show its data, which is not accurate.
-                # Instead, we set expiry date for the camera data.
-                await redis.setex(
-                    f"camera:{camera_id}:latest",  # And this is the key, or tag
-                    10,  # in seconds
-                    metadata.model_dump_json(),
-                )
+
+                results = await pipeline.run(camera_id, frame_bytes, next(step_counter))
 
                 # Note that JSONResponse doesn't work here, as it is for HTTP
-                await websocket.send_json({"status": 200, "camera_id": camera_id})
+                await websocket.send_json(results)
 
         except Exception as e:
             logger.error(f"Processing Error: {e}", camera_id=camera_id)
             raise
 
-    with mlflow.start_run(
-        run_name=f"camera_{camera_id}", nested=True, parent_run_id=state.mlflow_run_id
-    ):
+    with mlflow.start_run(run_name=f"camera_{camera_id}", nested=True, parent_run_id=state.mlflow_run_id):
         log_config()
 
         try:
@@ -197,4 +115,4 @@ async def websocket_detect(
             await redis.srem(
                 "cameras:active", camera_id
             )  # Remove the camera from redis connected cameras
-            active_cameras.dec()
+            active_cameras.dec()

services/pipeline.py

@@ -1,3 +1,15 @@
+from backend.api.routers.metrics import depth_duration_seconds
+from backend.api.routers.metrics import detection_duration_seconds
+from backend.api.routers.metrics import decode_duration_seconds
+from backend.utils.profiling import profile_step
+from backend.domain.detection_box_center import calculate_detection_box_center
+import asyncio
+from backend.contracts.camera_metadata import DetectionMetadata
+from backend.contracts.camera_metadata import CameraMetadata
+import cv2 as cv
+import numpy as np
+
+
 class ProcessingPipeline:
     def __init__(self, detector, depth_model, safety_detector, redis_client):
         self.detector = detector
@@ -5,9 +17,48 @@ class ProcessingPipeline:
         self.safety_detector = safety_detector
         self.redis_client = redis_client
 
-    async def run(self, camera_id:str, image_array):
+    def _decode_frame(fb):
+        return cv.imdecode(np.frombuffer(fb, np.uint8), cv.IMREAD_COLOR)
+
+    def _camera_metadata(self, camera_id, safety_detection, depth_points, boxes_center_ratio) -> CameraMetadata:
+        detection_metadata = [
+            DetectionMetadata(depth=depth, xRatio=xRatio) for depth, xRatio in zip(depth_points, boxes_center_ratio)
+        ]
+        metadata = CameraMetadata(
+            camera_id=camera_id,
+            is_danger=True if safety_detection else False,
+            detection_metadata=detection_metadata,
+        )
+        return metadata
+
+    async def run(self, camera_id:str, image_array, frame_count):
+        loop = asyncio.get_running_loop()
+
+        with profile_step("frame_processing_time", decode_duration_seconds, camera_id, frame_count):
+            image_array = await loop.run_in_executor(None, self._decode_frame, image_array)
+
+        with profile_step("detection_duration_seconds", detection_duration_seconds, camera_id, frame_count):
+            detection_task = loop.run_in_executor(None, self.detector.detect, image_array)
+            safety_task = loop.run_in_executor(None, self.safety_detector.detect, image_array)
+            detections, safety_detection = await asyncio.gather(detection_task, safety_task)
+
+        boxes_center, boxes_center_ratio = calculate_detection_box_center(detections.detections, image_array.shape[1])
+
+        depth_points = []
+        if boxes_center: 
+            with profile_step("depth_duration_seconds", depth_duration_seconds, camera_id, frame_count):
+                depth_points = await loop.run_in_executor(None, self.depth_model.calculate_depth, image_array, boxes_center)
+
+        metadata = self._camera_metadata(camera_id, safety_detection, depth_points, boxes_center_ratio)
+
+        await self.redis.publish("dashboard_stream", metadata.model_dump_json())
+        # Even if the camera was disconnected, redis is still going to show its data, which is not accurate.
+        # Instead, we set expiry date for the camera data.
+        await self.redis.setex(
+            f"camera:{camera_id}:latest",  # And this is the key, or tag
+            10,  # in seconds
+            metadata.model_dump_json(),
+        )
 
-        # Run ai models
-        # Domain logic
-        # save to infra
-        pass
+        # Note that JSONResponse doesn't work here, as it is for HTTP
+        return {"status": 200, "camera_id": camera_id}

utils/profiling.py

@@ -0,0 +1,20 @@
+from contextlib import contextmanager
+import time
+import mlflow
+
+@contextmanager
+def profile_step(expr_name: str, prometheus_logger, camera_id, frame_count=None):
+    """With statement utility to time block of code"""
+    start_time = time.time()
+
+    try:
+        # Code inside with statement
+        yield
+    finally:
+        duration = round(time.time() - start_time, 4)
+        prometheus_logger.labels(camera_id).observe(duration)
+        mlflow.log_metric(
+                    expr_name,
+                    duration,
+                    frame_count,
+                )