add pipeline overlay on/off

Signed-off-by: Max Xiang <[email protected]>

Author: maxxgx

application/backend/src/alembic/versions/7a213a27d666_initial_schema.py

@@ -136,6 +136,7 @@ def upgrade() -> None:
         sa.Column("sink_id", sa.Text(), nullable=True),
         sa.Column("model_id", sa.Text(), nullable=True),
         sa.Column("inference_device", sa.String(length=64), nullable=True),
+        sa.Column("overlay", sa.Boolean(), nullable=True),
         sa.Column("is_running", sa.Boolean(), nullable=False),
         sa.Column("is_active", sa.Boolean(), nullable=False),
         sa.Column("data_collection_policies", sa.JSON(), nullable=False),

application/backend/src/db/schema.py

@@ -90,6 +90,7 @@ class PipelineDB(Base):
     sink_id: Mapped[str | None] = mapped_column(Text, ForeignKey("sinks.id", ondelete="RESTRICT"))
     model_id: Mapped[str | None] = mapped_column(Text, ForeignKey("models.id", ondelete="RESTRICT"))
     inference_device: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    overlay: Mapped[bool | None] = mapped_column(Boolean, default=True)
     is_running: Mapped[bool] = mapped_column(Boolean, default=False)
     is_active: Mapped[bool] = mapped_column(Boolean, default=False)
     data_collection_policies: Mapped[list] = mapped_column(JSON, nullable=False, default=list)

application/backend/src/pydantic_models/pipeline.py

@@ -48,6 +48,7 @@ class Pipeline(BaseModel):
     )  # ID of the model, used for DB mapping, not exposed in API
     status: PipelineStatus = PipelineStatus.IDLE  # Current status of the pipeline
     inference_device: str | None = Field(default=None)
+    overlay: bool | None = Field(default=None)
 
     # TODO: can be confused with status.is_running / is_active, consider refactoring
     is_running: bool | None = Field(default=None, exclude=True)  # If set will overwrite status

application/backend/src/repositories/mappers/pipeline_mapper.py

@@ -25,6 +25,7 @@ def from_schema(pipeline_db: PipelineDB) -> Pipeline:
             source_id=UUID(pipeline_db.source_id) if pipeline_db.source_id else None,
             status=PipelineStatus.from_bool(pipeline_db.is_running, pipeline_db.is_active),
             inference_device=pipeline_db.inference_device.upper() if pipeline_db.inference_device else None,
+            overlay=pipeline_db.overlay,
         )
 
     @staticmethod
@@ -38,4 +39,5 @@ def to_schema(pipeline: Pipeline) -> PipelineDB:
             is_running=pipeline.status.is_running,
             is_active=pipeline.status.is_active,
             inference_device=pipeline.inference_device.upper() if pipeline.inference_device else None,
+            overlay=pipeline.overlay,
         )

application/backend/src/services/pipeline_service.py

@@ -74,7 +74,7 @@ async def update_pipeline(self, project_id: UUID, partial_config: dict) -> Pipel
                 await self._notify_pipeline_changed()
                 # Intentionally call activate_model on status change regardless of whether a model exists.
                 self._model_service.activate_model()
-            if updated.inference_device != pipeline.inference_device:
+            if updated.inference_device != pipeline.inference_device or updated.overlay != pipeline.overlay:
                 # reload model on device change
                 self._model_service.activate_model()
             return updated

application/backend/src/workers/inference.py

@@ -12,6 +12,7 @@
 
 if TYPE_CHECKING:
     import multiprocessing as mp
+    from collections.abc import Callable
     from multiprocessing.synchronize import Event as EventClass
     from multiprocessing.synchronize import Lock
 
@@ -57,6 +58,7 @@ def __init__(
         self._cached_models: dict[Any, object] = {}
         self._model_check_interval: float = 5.0  # seconds between model refresh checks
         self._is_passthrough_mode: bool = False
+        self._overlay = True
 
     def setup(self) -> None:
         super().setup()
@@ -71,6 +73,7 @@ async def _get_active_model(self) -> LoadedModel | None:
                 self._is_passthrough_mode = pipeline is None or (
                     pipeline.status.is_active and not pipeline.status.is_running
                 )
+                self._overlay = pipeline.overlay if pipeline and pipeline.overlay is not None else self._overlay
                 logger.info(f"Passthrough mode {'activated' if self._is_passthrough_mode else 'disabled'}.")
                 if pipeline is None or pipeline.model is None:
                     return None
@@ -224,7 +227,11 @@ async def _process_frame_with_model(self, stream_data: StreamData) -> None:
             raise RuntimeError("Inference failed or model became unavailable")
 
         # Build visualization
-        vis_frame: np.ndarray = Visualizer.overlay_predictions(frame, prediction_response)
+        overlays: list[Callable] = []
+        if self._overlay:
+            overlays.append(Visualizer.overlay_anomaly_heatmap)
+        overlays.append(Visualizer.draw_prediction_label)
+        vis_frame: np.ndarray = Visualizer.overlay_predictions(frame, prediction_response, *overlays)
 
         # Package inference data
         stream_data.inference_data = InferenceData(