diff --git a/docs/nvmsgbroker_integration.md b/docs/nvmsgbroker_integration.md index c07b4435..851efc74 100644 --- a/docs/nvmsgbroker_integration.md +++ b/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 diff --git a/env/mongodb/Dockerfile.consumer b/env/mongodb/Dockerfile.consumer index 8d154931..f0567803 100644 --- a/env/mongodb/Dockerfile.consumer +++ b/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"] diff --git a/src/nvmsgbroker_pipeline.py b/src/nvmsgbroker_pipeline.py index f2263de5..0f593ac8 100644 --- a/src/nvmsgbroker_pipeline.py +++ b/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