Merge pull request #32 from Sammy1Am/mqtt-support

Improved MQTT Support

README.md

@@ -90,7 +90,20 @@ Every request to DOODS involves the Detect Request JSON object that looks like t
     // only the first region (including the global detection) to match an object will be used.
     {"id": "someregion", "top": 0.1, "left": 0.1, "bottom": 0.9, "right": 0.9, "detect": {"*":50}, "covers": false}
     ...
-  ]
+  ],
+
+  // NOTE: Below fields are only available in requests configured as part of the MQTT configuration
+
+  // If separate_detections is true each detected object will be published separately into 
+  // a sub-topic based on its type (e.g doods/detect/requestid/regionid/person).  When False, the default,
+  // the whole DetectResponse object will be published to the request topic (e.g. doods/detect/requestid).
+  "separate_detections" : false,
+  // If crop is true and separate_detections is true requested images will be cropped to 
+  // the decection box.  Has no effect if separate_detections is false.
+  "crop": false,
+  // If binary_images is true requested images will be pubished as binary data 
+  // to a separate topic (e.g. doods/image/requestid) instead of base64 encoded into the response.
+  "binary_images" : false,
 }  
 ```
 

config.py

@@ -1,6 +1,6 @@
 from typing import List, Optional
 from pydantic import BaseSettings, Extra
-from odrpc import DetectRequest
+from odrpc import MqttDetectRequest
 
 class DoodsDetectorConfig(BaseSettings):
     name: str
@@ -69,7 +69,7 @@ class Config:
 
 class MqttConfig(BaseSettings):
     broker: MqttBrokerConfig
-    requests: List[DetectRequest]
+    requests: List[MqttDetectRequest]
     metrics: Optional[bool] = True
     class Config:
         env_prefix = 'mqtt_'

config.yaml

@@ -42,6 +42,9 @@ mqtt:
     - id: firstrequest
       detector_name: default
       preprocess: []
+      separate_detections: false
+      crop: false
+      binary_images: false
       detect:
         "*": 50
       regions:

mqtt.py

@@ -4,7 +4,9 @@
 import logging
 import threading
 import time
+import cv2
 import paho.mqtt.client as mqtt
+import numpy as np
 from streamer import Streamer
 from fastapi import FastAPI
 from prometheus_fastapi_instrumentator import Instrumentator
@@ -18,19 +20,67 @@ def __init__(self, config, doods, metrics_server_config=None):
         # Borrow the uvicorn logger because it's pretty.
         self.logger = logging.getLogger("doods.mqtt")
 
-    def stream(self, detect_request: str = '{}'):
+    def stream(self, mqtt_detect_request: str = '{}'):
         streamer = None
         try:
             # Run the stream detector and return the results.
-            streamer = Streamer(self.doods).start_stream(detect_request)
+            streamer = Streamer(self.doods).start_stream(mqtt_detect_request)
             for detect_response in streamer:
-                # If we requested an image, base64 encode it back to the user
-                if detect_request.image:
-                    detect_response.image = base64.b64encode(detect_response.image).decode('utf-8')
-                for detection in detect_response.detections:
+                # If separate_detections, iterate over each detection and process it separately
+                if mqtt_detect_request.separate_detections:
+
+                    #If we're going to be cropping, do this processing only once (rather than for each detection)
+                    if mqtt_detect_request.image and mqtt_detect_request.crop:
+                        detect_image_bytes = np.frombuffer(detect_response.image, dtype=np.uint8)
+                        detect_image = cv2.imdecode(detect_image_bytes, cv2.IMREAD_COLOR)
+                        di_height, di_width = detect_image.shape[:2]
+
+
+                    for detection in detect_response.detections:
+                        # If an image was requested
+                        if mqtt_detect_request.image:
+                            # Crop image to detection box if requested
+                            if mqtt_detect_request.crop:
+                                cropped_image = detect_image[
+                                int(detection.top*di_height):int(detection.bottom*di_height), 
+                                int(detection.left*di_width):int(detection.right*di_width)]
+                                mqtt_image = cv2.imencode(mqtt_detect_request.image, cropped_image)[1].tostring()
+                            else:
+                                mqtt_image = detect_response.image
+
+
+                            # For binary images, publish the image to its own topic
+                            if mqtt_detect_request.binary_images:
+                                self.mqtt_client.publish(
+                                    f"doods/image/{mqtt_detect_request.id}{'' if detection.region_id is None else '/'+detection.region_id}/{detection.label or 'object'}", 
+                                    payload=mqtt_image, qos=0, retain=False)
+                            # Otherwise add base64-encoded image to the detection
+                            else:
+                                detection.image = base64.b64encode(mqtt_image).decode('utf-8')
+
+                        self.mqtt_client.publish(
+                            f"doods/detect/{mqtt_detect_request.id}{'' if detection.region_id is None else '/'+detection.region_id}/{detection.label or 'object'}", 
+                            payload=json.dumps(detection.asdict(include_none=False)), qos=0, retain=False)
+
+                # Otherwise, publish the collected detections together
+                else:
+                    # If an image was requested
+                    if mqtt_detect_request.image:
+                        # If binary_images, move the image from the response and publish it to a separate topic
+                        if mqtt_detect_request.binary_images:
+                            mqtt_image = detect_response.image
+                            detect_response.image = None
+                            self.mqtt_client.publish(
+                                f"doods/image/{mqtt_detect_request.id}", 
+                                payload=detect_response.image, qos=0, retain=False)
+                        # Otherwise, inlcude the base64-encoded image in the response
+                        else:
+                            detect_response.image = base64.b64encode(detect_response.image).decode('utf-8')
+
                     self.mqtt_client.publish(
-                        f"doods/detect/{detect_request.id}{'' if detection.region_id is None else '/'+detection.region_id}", 
-                        payload=json.dumps(detection.asdict(include_none=False)), qos=0, retain=False)
+                            f"doods/detect/{mqtt_detect_request.id}", 
+                            payload=json.dumps(detect_response.asdict(include_none=False)), qos=0, retain=False)
+
 
         except Exception as e:
             self.logger.info(e)

odrpc.py

@@ -26,6 +26,12 @@ class DetectRequest:
     detect: Dict[str, float] = field(default_factory=dict)
     regions: List[DetectRegion] = field(default_factory=list)    
 
+@dataclass
+class MqttDetectRequest(DetectRequest):
+    separate_detections: Optional[bool] = False
+    crop: Optional[bool] = False
+    binary_images: Optional[bool] = False
+
 @dataclass
 class Detector:
     name: str
@@ -55,6 +61,7 @@ class Detection:
     right: float = 0.0
     label: str = ""
     confidence: float = 0.0
+    image: Optional[str] = None
 
     def asdict(self, include_none=True):
         ret = asdict(self)