Merge pull request #191 from robmarkcole/add-object-type

Add object_type to events

README.md

@@ -7,9 +7,9 @@ docker run -e VISION-DETECTION=True -e API-KEY="mysecretkey" -v localstorage:/da
 ```
 
 ## Usage of this component
-The `deepstack_object` component adds an `image_processing` entity where the state of the entity is the total count of target objects that are above a `confidence` threshold which has a default value of 80%. You can have a single target object class, or multiple. The time of the last detection of any target object is in the `last target detection` attribute. The type and number of objects (of any confidence) is listed in the `summary` attributes. Optionally a region of interest (ROI) can be configured, and only objects with their center (represented by a `x`) within the ROI will be included in the state count. The ROI will be displayed as a green box, and objects with their center in the ROI have a red box, whilst objects with their center outside the ROI have a yellow box.
+The `deepstack_object` component adds an `image_processing` entity where the state of the entity is the total count of target objects that are above a `confidence` threshold which has a default value of 80%. You can have a single target object class, or multiple. The time of the last detection of any target object is in the `last target detection` attribute. The type and number of objects (of any confidence) is listed in the `summary` attributes. Optionally a region of interest (ROI) can be configured, and only objects with their center (represented by a `x`) within the ROI will be included in the state count. The ROI will be displayed as a green box, and objects with their center in the ROI have a red box.
 
-Also optionally the processed image can be saved to disk, with bounding boxes showing the location of detected objects. If `save_file_folder` is configured, an image with filename of format `deepstack_object_{source name}_latest.jpg` is over-written on each new detection of a target. Optionally this image can also be saved with a timestamp in the filename, if `save_timestamped_file` is configured as `True`. An event `deepstack.object_detected` is fired for each object detected. If you are a power user with advanced needs such as zoning detections or you want to track multiple object types, you will need to use the `deepstack.object_detected` events.
+Also optionally the processed image can be saved to disk, with bounding boxes showing the location of detected objects. If `save_file_folder` is configured, an image with filename of format `deepstack_object_{source name}_latest.jpg` is over-written on each new detection of a target. Optionally this image can also be saved with a timestamp in the filename, if `save_timestamped_file` is configured as `True`. An event `deepstack.object_detected` is fired for each object detected that is in the targets list, and meets the confidence and ROI criteria. If you are a power user with advanced needs such as zoning detections or you want to track multiple object types, you will need to use the `deepstack.object_detected` events.
 
 **Note** that by default the component will **not** automatically scan images, but requires you to call the `image_processing.scan` service e.g. using an automation triggered by motion.
 
@@ -77,6 +77,7 @@ An example use case for event is to get an alert when some rarely appearing obje
 
 - `entity_id` : the entity id responsible for the event
 - `name` : the name of the type of object detected
+- `object_type` : the type of the object, from `person`, `vehicle`, `animal` or `other`
 - `confidence` : the confidence in detection in the range 0 - 100%
 - `box` : the bounding box of the object
 - `centroid` : the centre point of the object
@@ -170,8 +171,8 @@ A6: This can happen when you are running in Docker/Hassio, and indicates that on
 
 ------
 
-## Objects list
-The following list is [from the deepstack docs](https://python.deepstack.cc/object-detection):
+## Objects
+The following lists all target objects:
 ```
 person,   bicycle,   car,   motorcycle,   airplane,
 bus,   train,   truck,   boat,   traffic light,   fire hydrant,   stop_sign,
@@ -185,7 +186,11 @@ toilet,   tv,   laptop,   mouse,   remote,   keyboard,   cell phone,   microwave
 oven,   toaster,   sink,   refrigerator,   book,   clock,   vase,   scissors,   teddy bear,
 hair dryer, toothbrush.
 ```
-
+Objects are grouped by the following `object_type`:
+- **person**: person
+- **animal**: bird, cat, dog, horse, sheep, cow, elephant, bear, zebra, giraffe
+- **vehicle**: bicycle, car, motorcycle, airplane, bus, train, truck
+- **other**: any object that is not in `person`, `animal` or `vehicle`
 
 ## Development
 Currently only the helper functions are tested, using pytest.

custom_components/deepstack_object/image_processing.py

@@ -43,6 +43,24 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+ANIMAL = "animal"
+ANIMALS = [
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+]
+OTHER = "other"
+PERSON = "person"
+VEHICLE = "vehicle"
+VEHICLES = ["bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck"]
+
 CONF_API_KEY = "api_key"
 CONF_TARGETS = "targets"
 CONF_TIMEOUT = "timeout"
@@ -57,7 +75,7 @@
 
 DATETIME_FORMAT = "%Y-%m-%d_%H-%M-%S"
 DEFAULT_API_KEY = ""
-DEFAULT_TARGETS = ["person"]
+DEFAULT_TARGETS = [PERSON]
 DEFAULT_TIMEOUT = 10
 DEFAULT_ROI_Y_MIN = 0.0
 DEFAULT_ROI_Y_MAX = 1.0
@@ -124,6 +142,17 @@ def get_valid_filename(name: str) -> str:
     return re.sub(r"(?u)[^-\w.]", "", str(name).strip().replace(" ", "_"))
 
 
+def get_object_type(object_name: str) -> str:
+    if object_name == PERSON:
+        return PERSON
+    elif object_name in ANIMALS:
+        return ANIMAL
+    elif object_name in VEHICLES:
+        return VEHICLE
+    else:
+        return OTHER
+
+
 def get_objects(predictions: list, img_width: int, img_height: int) -> List[Dict]:
     """Return objects with formatting and extra info."""
     objects = []
@@ -145,6 +174,7 @@ def get_objects(predictions: list, img_width: int, img_height: int) -> List[Dict
             "y": round(box["y_min"] + (box["height"] / 2), decimal_places),
         }
         name = pred["label"]
+        object_type = get_object_type(name)
         confidence = round(pred["confidence"] * 100, decimal_places)
 
         objects.append(
@@ -153,6 +183,7 @@ def get_objects(predictions: list, img_width: int, img_height: int) -> List[Dict
                 "box_area": box_area,
                 "centroid": centroid,
                 "name": name,
+                "object_type": object_type,
                 "confidence": confidence,
             }
         )
@@ -365,25 +396,20 @@ def save_image(self, image, targets, directory) -> str:
             centroid = obj["centroid"]
             box_label = f"{name}: {confidence:.1f}%"
 
-            if object_in_roi(self._roi_dict, centroid):
-                box_colour = RED
-            else:
-                box_colour = YELLOW
-
             draw_box(
                 draw,
                 (box["y_min"], box["x_min"], box["y_max"], box["x_max"]),
                 img.width,
                 img.height,
                 text=box_label,
-                color=box_colour,
+                color=RED,
             )
 
             # draw bullseye
             draw.text(
                 (centroid["x"] * img.width, centroid["y"] * img.height),
                 text="X",
-                fill=box_colour,
+                fill=RED,
             )
 
         # Save images, returning the path of saved image as str