Support Tensorflow plugin as motion sensor

elad-bar · May 14, 2021 · 4a906d5 · 4a906d5
1 parent 9fde8a1
commit 4a906d5
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## v1.0.2
+
+- Support Tensorflow plugin as motion sensor
+
 ## v1.0.1
 
 - Fix hassfest validation error - Added iot_class=local_polling to manifest

diff --git a/custom_components/shinobi/helpers/const.py b/custom_components/shinobi/helpers/const.py
@@ -195,13 +195,15 @@
     ATTR_CAMERA_DETAILS_FPS: ATTR_FPS
 }
 
+TRIGGER_PLUG = "plug"
 TRIGGER_NAME = "name"
 TRIGGER_DETAILS = "details"
 TRIGGER_DETAILS_PLUG = "plug"
 TRIGGER_DETAILS_REASON = "reason"
 TRIGGER_DETAILS_MATRICES = "matrices"
 TRIGGER_DETAILS_MATRICES_TAG = "tag"
 
+TRIGGER_PLUG_TENSORFLOW = "Tensorflow"
 TRIGGER_PLUG_YOLO = "Yolo"
 TRIGGER_PLUG_DB = "audio"
 
@@ -213,11 +215,13 @@
 SOUND_DETECTION = "Sound Detection"
 
 SENSOR_DEVICE_CLASS = {
+    TRIGGER_PLUG_TENSORFLOW: "motion",
     TRIGGER_PLUG_YOLO: "motion",
     TRIGGER_PLUG_DB: "sound"
 }
 
 TRIGGER_DURATION = {
+    TRIGGER_PLUG_TENSORFLOW: 20,
     TRIGGER_PLUG_YOLO: 20,
     TRIGGER_PLUG_DB: 10
 }

diff --git a/custom_components/shinobi/managers/entity_manager.py b/custom_components/shinobi/managers/entity_manager.py
@@ -243,9 +243,7 @@ def get_camera_entity(self, camera: CameraData, sensor_type) -> EntityData:
         try:
             device_name = self.device_manager.get_camera_device_name(camera)
 
-            device_class = SENSOR_DEVICE_CLASS.get(sensor_type, sensor_type)
-
-            entity_name = f"{self.integration_title} {camera.name} {device_class.capitalize()}"
+            entity_name = f"{self.integration_title} {camera.name} {sensor_type.capitalize()}"
             unique_id = f"{DOMAIN}-{DOMAIN_BINARY_SENSOR}-{entity_name}"
 
             state_topic = f"{MQTT_ALL_TOPIC}/{self.api.group_id}/{camera.monitorId}/trigger"
@@ -274,7 +272,7 @@ def get_camera_entity(self, camera: CameraData, sensor_type) -> EntityData:
             entity.device_name = device_name
             entity.topic = state_topic
             entity.event = sensor_type
-            entity.device_class = device_class.lower()
+            entity.device_class = sensor_type
             entity.type = sensor_type
 
             self.mqtt_manager.set_state(state_topic, sensor_type, mqtt_state)
@@ -292,11 +290,14 @@ def generate_camera_binary_sensors(self, camera: CameraData):
 
         try:
             supported_sensors = []
+            audio_event = SENSOR_DEVICE_CLASS[TRIGGER_PLUG_DB]
+            motion_event = SENSOR_DEVICE_CLASS[TRIGGER_PLUG_YOLO]
+
             if camera.has_audio and camera.has_audio_detector:
-                supported_sensors.append(TRIGGER_PLUG_DB)
+                supported_sensors.append(audio_event)
 
             if camera.has_motion_detector:
-                supported_sensors.append(TRIGGER_PLUG_YOLO)
+                supported_sensors.append(motion_event)
 
             for sensor_type_name in supported_sensors:
                 entity = self.get_camera_entity(camera, sensor_type_name)

diff --git a/custom_components/shinobi/managers/mqtt_manager.py b/custom_components/shinobi/managers/mqtt_manager.py
@@ -83,15 +83,18 @@ def _state_message_received(self, message: Message):
 
         value = {
             TRIGGER_NAME: trigger_name,
+            TRIGGER_PLUG: trigger_plug,
             TRIGGER_DETAILS_REASON: trigger_reason,
             TRIGGER_TAGS: trigger_tags,
             TRIGGER_STATE: STATE_ON,
             TRIGGER_TIMESTAMP: datetime.now().timestamp()
         }
 
-        _LOGGER.debug(f"Topic: {topic} for {trigger_plug}: {value}")
+        trigger_event_name = SENSOR_DEVICE_CLASS[trigger_plug]
 
-        self.set_state(topic, trigger_plug, value)
+        _LOGGER.debug(f"Topic: {topic} for {trigger_plug} ({trigger_event_name}): {value}")
+
+        self.set_state(topic, trigger_event_name, value)
         self.event_handler()
 
     def get_state(self, topic, event_type):