Modular nvmsgbroker pipeline

docs/nvmsgbroker_integration.md

@@ -26,16 +26,16 @@ It processes **RTSP video streams**, runs **classification inference**, and stor
 
 ## 🚀 How to Run
 
+
 ```bash
 # 1. Start database services (MongoDB + MQTT broker)
-docker compose -f env/mongodb/docker-compose.yml up -d
+./database/scripts/up.sh
 
 # 2. Clear MongoDB (for fresh start)
 docker exec -it agstream_mongo mongosh agstream --eval "db.predictions.deleteMany({})"
 
 # 3. Start Consumer (builds image automatically on first run)
 ./database/start_consumer.sh
-
 # 4. Verify Consumer is running
 docker logs mqtt_consumer -f
 

env/mongodb/Dockerfile.consumer

@@ -1,5 +1,5 @@
 FROM python:3.9-slim
 WORKDIR /app
 RUN pip install paho-mqtt pymongo
-COPY consumer.py .
+COPY database/consumer.py .
 CMD ["python", "-u", "consumer.py"]

src/nvmsgbroker_pipeline.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import ctypes
+from typing import List, Dict, Any, Tuple
 import gi
 gi.require_version('Gst', '1.0')
 from gi.repository import Gst, GLib
@@ -13,7 +14,7 @@
 MSGCONV_CONFIG = "/workspace/src/configs/nvmsgbroker_msgconv_config.txt"
 
 # Load class labels
-def load_class_labels():
+def load_class_labels() -> List[str]:
     try:
         with open("/workspace/src/configs/labels.txt", "r") as f:
             return [line.strip() for line in f.readlines()]
@@ -22,6 +23,61 @@ def load_class_labels():
 
 CLASS_LABELS = load_class_labels()
 
+def apply_softmax_normalization(probs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Apply softmax normalization and return top 5 indices."""
+    exp_probs = np.exp(probs - np.max(probs))
+    normalized_probs = exp_probs / np.sum(exp_probs)
+    top_indices = np.argsort(normalized_probs)[-5:][::-1]
+    return normalized_probs, top_indices
+
+def process_tensor_operations(frame_meta) -> List[Dict[str, Any]]:
+    """Extract and process tensor data from frame metadata."""
+    classification_results = []
+    l_user = frame_meta.frame_user_meta_list
+    while l_user is not None:
+        user_meta = pyds.NvDsUserMeta.cast(l_user.data)
+        if user_meta.base_meta.meta_type == pyds.NvDsMetaType.NVDSINFER_TENSOR_OUTPUT_META:
+            tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data)
+            layer = pyds.get_nvds_LayerInfo(tensor_meta, 0)
+            ptr = ctypes.cast(pyds.get_ptr(layer), ctypes.POINTER(ctypes.c_float))
+            probs = np.ctypeslib.as_array(ptr, shape=(83,))
+
+            normalized_probs, top_indices = apply_softmax_normalization(probs)
+
+            for idx in top_indices:
+                if normalized_probs[idx] > 0.001:
+                    class_name = CLASS_LABELS[idx] if idx < len(CLASS_LABELS) else f"unknown_{idx}"
+                    classification_results.append({
+                        "class_id": int(idx),
+                        "confidence": float(normalized_probs[idx]),
+                        "class_name": class_name
+                    })
+        l_user = l_user.next
+    return classification_results
+
+def create_and_add_message_metadata(batch_meta, frame_meta, frame_number: int, classification_results: List[Dict[str, Any]]) -> None:
+    """Create nvmsgbroker metadata and add it to frame metadata."""
+    user_event_meta = pyds.nvds_acquire_user_meta_from_pool(batch_meta)
+    if user_event_meta:
+        msg_meta = pyds.alloc_nvds_event_msg_meta(user_event_meta)
+        msg_meta.frameId = frame_number
+        msg_meta.sensorId = 0
+        msg_meta.placeId = 0
+        msg_meta.moduleId = 0
+        msg_meta.sensorStr = "sensor-0"
+        msg_meta.ts = pyds.alloc_buffer(32)
+        pyds.generate_ts_rfc3339(msg_meta.ts, 32)
+
+        if classification_results:
+            best_result = classification_results[0]
+            msg_meta.objectId = str(best_result["class_id"])
+            msg_meta.confidence = best_result["confidence"]
+            print(f"📡 Frame {frame_number}: {best_result['class_name']} -> nvmsgbroker")
+
+        user_event_meta.user_meta_data = msg_meta
+        user_event_meta.base_meta.meta_type = pyds.NvDsMetaType.NVDS_EVENT_MSG_META
+        pyds.nvds_add_user_meta_to_frame(frame_meta, user_event_meta)
+
 def buffer_probe(pad, info, u_data):
     gst_buffer = info.get_buffer()
     batch_meta = pyds.gst_buffer_get_nvds_batch_meta(hash(gst_buffer))
@@ -31,52 +87,10 @@ def buffer_probe(pad, info, u_data):
         frame_meta = pyds.NvDsFrameMeta.cast(l_frame.data)
         frame_number = frame_meta.frame_num
 
-        # Get classification results
-        classification_results = []
-        l_user = frame_meta.frame_user_meta_list
-        while l_user is not None:
-            user_meta = pyds.NvDsUserMeta.cast(l_user.data)
-            if user_meta.base_meta.meta_type == pyds.NvDsMetaType.NVDSINFER_TENSOR_OUTPUT_META:
-                tensor_meta = pyds.NvDsInferTensorMeta.cast(user_meta.user_meta_data)
-                layer = pyds.get_nvds_LayerInfo(tensor_meta, 0)
-                ptr = ctypes.cast(pyds.get_ptr(layer), ctypes.POINTER(ctypes.c_float))
-                probs = np.ctypeslib.as_array(ptr, shape=(83,))
-                exp_probs = np.exp(probs - np.max(probs))
-                normalized_probs = exp_probs / np.sum(exp_probs)
-                top_indices = np.argsort(normalized_probs)[-5:][::-1]
-
-                for idx in top_indices:
-                    if normalized_probs[idx] > 0.001:
-                        class_name = CLASS_LABELS[idx] if idx < len(CLASS_LABELS) else f"unknown_{idx}"
-                        classification_results.append({
-                            "class_id": int(idx),
-                            "confidence": float(normalized_probs[idx]),
-                            "class_name": class_name
-                        })
-            l_user = l_user.next
+        classification_results = process_tensor_operations(frame_meta)
 
-        # Create nvmsgbroker metadata every 30 frames
         if (frame_number % 30) == 0:
-            user_event_meta = pyds.nvds_acquire_user_meta_from_pool(batch_meta)
-            if user_event_meta:
-                msg_meta = pyds.alloc_nvds_event_msg_meta(user_event_meta)
-                msg_meta.frameId = frame_number
-                msg_meta.sensorId = 0
-                msg_meta.placeId = 0
-                msg_meta.moduleId = 0
-                msg_meta.sensorStr = "sensor-0"
-                msg_meta.ts = pyds.alloc_buffer(32)
-                pyds.generate_ts_rfc3339(msg_meta.ts, 32)
-
-                if classification_results:
-                    best_result = classification_results[0]
-                    msg_meta.objectId = str(best_result["class_id"])
-                    msg_meta.confidence = best_result["confidence"]
-                    print(f"📡 Frame {frame_number}: {best_result['class_name']} -> nvmsgbroker")
-
-                user_event_meta.user_meta_data = msg_meta
-                user_event_meta.base_meta.meta_type = pyds.NvDsMetaType.NVDS_EVENT_MSG_META
-                pyds.nvds_add_user_meta_to_frame(frame_meta, user_event_meta)
+            create_and_add_message_metadata(batch_meta, frame_meta, frame_number, classification_results)
 
         l_frame = l_frame.next
     return Gst.PadProbeReturn.OK