Merge pull request #7 from robmarkcole/dev

Add saving images

README.md

@@ -12,10 +12,19 @@ Place the `custom_components` folder in your configuration directory (or add its
 image_processing:
   - platform: platerecognizer
     api_token: your_token
+    save_file_folder: /config/images/platerecognizer/
+    save_timestamped_file: True
+    always_save_latest_jpg: True
     source:
       - entity_id: camera.yours
 ```
 
+Configuration variables:
+- **api_key**: Your api key.
+- **save_file_folder**: (Optional) The folder to save processed images to. Note that folder path should be added to [whitelist_external_dirs](https://www.home-assistant.io/docs/configuration/basic/)
+- **save_timestamped_file**: (Optional, default `False`, requires `save_file_folder` to be configured) Save the processed image with the time of detection in the filename.
+- **always_save_latest_jpg**: (Optional, default `False`, requires `save_file_folder` to be configured) Always save the last processed image, no matter there were detections or not.
+- **source**: Must be a camera.
 
 ## Development
 Currently only the helper functions are tested, using pytest.

custom_components/platerecognizer/image_processing.py

@@ -2,6 +2,11 @@
 import logging
 import requests
 import voluptuous as vol
+import re
+import io
+
+from PIL import Image, ImageDraw, UnidentifiedImageError
+from pathlib import Path
 
 from homeassistant.components.image_processing import (
     CONF_ENTITY_ID,
@@ -14,11 +19,12 @@
 from homeassistant.core import split_entity_id
 import homeassistant.helpers.config_validation as cv
 import homeassistant.util.dt as dt_util
+from homeassistant.util.pil import draw_box
 
 _LOGGER = logging.getLogger(__name__)
 
 PLATE_READER_URL = "https://api.platerecognizer.com/v1/plate-reader/"
-STATS_URL = 'https://api.platerecognizer.com/v1/statistics/'
+STATS_URL = "https://api.platerecognizer.com/v1/statistics/"
 
 EVENT_VEHICLE_DETECTED = "platerecognizer.vehicle_detected"
 
@@ -28,25 +34,42 @@
 ATTR_VEHICLE_TYPE = "vehicle_type"
 
 CONF_API_TOKEN = "api_token"
+CONF_SAVE_FILE_FOLDER = "save_file_folder"
+CONF_SAVE_TIMESTAMPTED_FILE = "save_timestamped_file"
+CONF_ALWAYS_SAVE_LATEST_JPG = "always_save_latest_jpg"
 
 DATETIME_FORMAT = "%Y-%m-%d_%H-%M-%S"
+RED = (255, 0, 0)  # For objects within the ROI
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Required(CONF_API_TOKEN): cv.string,
+        vol.Optional(CONF_SAVE_FILE_FOLDER): cv.isdir,
+        vol.Optional(CONF_SAVE_TIMESTAMPTED_FILE, default=False): cv.boolean,
+        vol.Optional(CONF_ALWAYS_SAVE_LATEST_JPG, default=False): cv.boolean,
     }
 )
 
+# def get_valid_filename(name: str) -> str:
+#     return re.sub(r"(?u)[^-\w.]", "", str(name).strip().replace(" ", "_"))
+
 
 def setup_platform(hass, config, add_entities, discovery_info=None):
     """Set up the platform."""
     # Validate credentials by processing image.
+    save_file_folder = config.get(CONF_SAVE_FILE_FOLDER)
+    if save_file_folder:
+        save_file_folder = Path(save_file_folder)
+
     entities = []
     for camera in config[CONF_SOURCE]:
         platerecognizer = PlateRecognizerEntity(
-            config.get(CONF_API_TOKEN),
-            camera[CONF_ENTITY_ID],
-            camera.get(CONF_NAME),
+            api_token=config.get(CONF_API_TOKEN),
+            save_file_folder=save_file_folder,
+            save_timestamped_file=config.get(CONF_SAVE_TIMESTAMPTED_FILE),
+            always_save_latest_jpg=config.get(CONF_ALWAYS_SAVE_LATEST_JPG),
+            camera_entity=camera[CONF_ENTITY_ID],
+            name=camera.get(CONF_NAME),
         )
         entities.append(platerecognizer)
     add_entities(entities)
@@ -55,28 +78,42 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
 class PlateRecognizerEntity(ImageProcessingEntity):
     """Create entity."""
 
-    def __init__(self, api_token, camera_entity, name):
+    def __init__(
+        self,
+        api_token,
+        save_file_folder,
+        save_timestamped_file,
+        always_save_latest_jpg,
+        camera_entity,
+        name,
+    ):
         """Init."""
-        self._headers = {
-            "Authorization": f"Token {api_token}",
-        }
+        self._headers = {"Authorization": f"Token {api_token}"}
         self._camera = camera_entity
         if name:
             self._name = name
         else:
             camera_name = split_entity_id(camera_entity)[1]
             self._name = f"platerecognizer_{camera_name}"
+        self._save_file_folder = save_file_folder
+        self._save_timestamped_file = save_timestamped_file
+        self._always_save_latest_jpg = always_save_latest_jpg
         self._state = None
         self._results = {}
         self._vehicles = [{}]
         self._statistics = {}
         self._last_detection = None
+        self._image_width = None
+        self._image_height = None
+        self._image = None
         self.get_statistics()
 
     def process_image(self, image):
         """Process an image."""
         self._results = {}
         self._vehicles = [{}]
+        self._image = Image.open(io.BytesIO(bytearray(image)))
+        self._image_width, self._image_height = self._image.size
         try:
             self._results = requests.post(
                 PLATE_READER_URL, files={"upload": image}, headers=self._headers
@@ -98,13 +135,16 @@ def process_image(self, image):
             self._last_detection = dt_util.now().strftime(DATETIME_FORMAT)
             for vehicle in self._vehicles:
                 self.fire_vehicle_detected_event(vehicle)
+        if self._save_file_folder:
+            if self._state > 0 or self._always_save_latest_jpg:
+                self.save_image()
         self.get_statistics()
 
     def get_statistics(self):
         try:
             response = requests.get(STATS_URL, headers=self._headers).json()
-            calls_remaining = response['total_calls'] - response['usage']['calls']
-            response.update({'calls_remaining': calls_remaining})
+            calls_remaining = response["total_calls"] - response["usage"]["calls"]
+            response.update({"calls_remaining": calls_remaining})
             self._statistics = response.copy()
         except Exception as exc:
             _LOGGER.error("platerecognizer error getting statistics: %s", exc)
@@ -115,6 +155,36 @@ def fire_vehicle_detected_event(self, vehicle):
         vehicle_copy.update({ATTR_ENTITY_ID: self.entity_id})
         self.hass.bus.fire(EVENT_VEHICLE_DETECTED, vehicle_copy)
 
+    def save_image(self):
+        """Save a timestamped image with bounding boxes around plates."""
+        draw = ImageDraw.Draw(self._image)
+
+        decimal_places = 3
+        for vehicle in self._results:
+            box = (
+                    round(vehicle['box']["ymin"] / self._image_height, decimal_places),
+                    round(vehicle['box']["xmin"] / self._image_width, decimal_places),
+                    round(vehicle['box']["ymax"] / self._image_height, decimal_places),
+                    round(vehicle['box']["xmax"] / self._image_width, decimal_places),
+            )
+            text = vehicle['plate']
+            draw_box(
+                draw,
+                box,
+                self._image_width,
+                self._image_height,
+                text=text,
+                color=RED,
+                )
+
+        latest_save_path = self._save_file_folder / f"{self._name}_latest.png"
+        self._image.save(latest_save_path)
+
+        if self._save_timestamped_file:
+            timestamp_save_path = self._save_file_folder / f"{self._name}_{self._last_detection}.png"
+            self._image.save(timestamp_save_path)
+            _LOGGER.info("platerecognizer saved file %s", timestamp_save_path)
+
     @property
     def camera_entity(self):
         """Return camera entity id from process pictures."""
@@ -147,4 +217,8 @@ def device_state_attributes(self):
         attr.update({"last_detection": self._last_detection})
         attr.update({"vehicles": self._vehicles})
         attr.update({"statistics": self._statistics})
+        if self._save_file_folder:
+            attr[CONF_SAVE_FILE_FOLDER] = str(self._save_file_folder)
+            attr[CONF_SAVE_TIMESTAMPTED_FILE] = self._save_timestamped_file
+            attr[CONF_ALWAYS_SAVE_LATEST_JPG] = self._always_save_latest_jpg
         return attr

custom_components/platerecognizer/manifest.json

@@ -2,7 +2,7 @@
         "domain": "platerecognizer",
         "name": "platerecognizer custom integration",
         "documentation": "https://github.com/robmarkcole/HASS-plate-recognizer",
-        "version": "0.3.0",
+        "version": "0.4.0",
         "requirements": ["pillow", "requests"],
         "dependencies": [],
         "codeowners": [