Merge pull request #224 from robmarkcole/ISSUE221

Issue221

Author: robmarkcole

README.md

@@ -29,8 +29,9 @@ image_processing:
     # custom_model: mask
     # confidence: 80
     save_file_folder: /config/snapshots/
+    save_file_format: png
     save_timestamped_file: True
-    always_save_latest_jpg: True
+    always_save_latest_file: True
     scale: 0.75
     # roi_x_min: 0.35
     roi_x_max: 0.8
@@ -54,8 +55,9 @@ Configuration variables:
 - **custom_model**: (Optional) The name of a custom model if you are using one. Don't forget to add the targets from the custom model below
 - **confidence**: (Optional) The confidence (in %) above which detected targets are counted in the sensor state. Default value: 80
 - **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_file_format**: (Optional, default `jpg`, alternatively `png`) The file format to save images as. `png` generally results in easier to read annotations.
 - **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, even if there were no detections.
+- **always_save_latest_file**: (Optional, default `False`, requires `save_file_folder` to be configured) Always save the last processed image, even if there were no detections.
 - **scale**: (optional, default 1.0), range 0.1-1.0, applies a scaling factor to the images that are saved. This reduces the disk space used by saved images, and is especially beneficial when using high resolution cameras.
 - **show_boxes**: (optional, default `True`), if `False` bounding boxes are not shown on saved images
 - **roi_x_min**: (optional, default 0), range 0-1, must be less than roi_x_max

custom_components/deepstack_object/image_processing.py

@@ -69,9 +69,10 @@
 CONF_TARGET = "target"
 CONF_TARGETS = "targets"
 CONF_TIMEOUT = "timeout"
+CONF_SAVE_FILE_FORMAT = "save_file_format"
 CONF_SAVE_FILE_FOLDER = "save_file_folder"
 CONF_SAVE_TIMESTAMPTED_FILE = "save_timestamped_file"
-CONF_ALWAYS_SAVE_LATEST_JPG = "always_save_latest_jpg"
+CONF_ALWAYS_SAVE_LATEST_FILE = "always_save_latest_file"
 CONF_SHOW_BOXES = "show_boxes"
 CONF_ROI_Y_MIN = "roi_y_min"
 CONF_ROI_X_MIN = "roi_x_min"
@@ -102,6 +103,8 @@
 OBJECT = "object"
 SAVED_FILE = "saved_file"
 MIN_CONFIDENCE = 0.1
+JPG = "jpg"
+PNG = "png"
 
 # rgb(red, green, blue)
 RED = (255, 0, 0)  # For objects within the ROI
@@ -134,8 +137,9 @@
             vol.Coerce(float, vol.Range(min=0.1, max=1))
         ),
         vol.Optional(CONF_SAVE_FILE_FOLDER): cv.isdir,
+        vol.Optional(CONF_SAVE_FILE_FORMAT, default=JPG): vol.In([JPG, PNG]),
         vol.Optional(CONF_SAVE_TIMESTAMPTED_FILE, default=False): cv.boolean,
-        vol.Optional(CONF_ALWAYS_SAVE_LATEST_JPG, default=False): cv.boolean,
+        vol.Optional(CONF_ALWAYS_SAVE_LATEST_FILE, default=False): cv.boolean,
         vol.Optional(CONF_SHOW_BOXES, default=True): cv.boolean,
     }
 )
@@ -233,8 +237,9 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
             scale=config[CONF_SCALE],
             show_boxes=config[CONF_SHOW_BOXES],
             save_file_folder=save_file_folder,
+            save_file_format=config[CONF_SAVE_FILE_FORMAT],
             save_timestamped_file=config.get(CONF_SAVE_TIMESTAMPTED_FILE),
-            always_save_latest_jpg=config.get(CONF_ALWAYS_SAVE_LATEST_JPG),
+            always_save_latest_file=config.get(CONF_ALWAYS_SAVE_LATEST_FILE),
             camera_entity=camera.get(CONF_ENTITY_ID),
             name=camera.get(CONF_NAME),
         )
@@ -261,8 +266,9 @@ def __init__(
         scale,
         show_boxes,
         save_file_folder,
+        save_file_format,
         save_timestamped_file,
-        always_save_latest_jpg,
+        always_save_latest_file,
         camera_entity,
         name=None,
     ):
@@ -296,6 +302,7 @@ def __init__(
         self._state = None
         self._objects = []  # The parsed raw data
         self._targets_found = []
+        self._last_detection = None
 
         self._roi_dict = {
             "y_min": roi_y_min,
@@ -305,12 +312,13 @@ def __init__(
         }
         self._scale = scale
         self._show_boxes = show_boxes
-        self._last_detection = None
         self._image_width = None
         self._image_height = None
         self._save_file_folder = save_file_folder
+        self._save_file_format = save_file_format
+        self._always_save_latest_file = always_save_latest_file
         self._save_timestamped_file = save_timestamped_file
-        self._always_save_latest_jpg = always_save_latest_jpg
+        self._always_save_latest_file = always_save_latest_file
         self._image = None
 
     def process_image(self, image):
@@ -377,7 +385,7 @@ def process_image(self, image):
         self._summary = dict(Counter(targets_found))  # e.g. {'car':2, 'person':1}
 
         if self._save_file_folder:
-            if self._state > 0 or self._always_save_latest_jpg:
+            if self._state > 0 or self._always_save_latest_file:
                 saved_image_path = self.save_image(
                     self._targets_found, self._save_file_folder,
                 )
@@ -405,6 +413,11 @@ def name(self):
         """Return the name of the sensor."""
         return self._name
 
+    @property
+    def unit_of_measurement(self):
+        """Return the unit of measurement."""
+        return "targets"
+
     @property
     def should_poll(self):
         """Return the polling state."""
@@ -428,8 +441,9 @@ def device_state_attributes(self) -> Dict:
         ]
         if self._save_file_folder:
             attr[CONF_SAVE_FILE_FOLDER] = str(self._save_file_folder)
+            attr[CONF_SAVE_FILE_FORMAT] = self._save_file_format
             attr[CONF_SAVE_TIMESTAMPTED_FILE] = self._save_timestamped_file
-            attr[CONF_ALWAYS_SAVE_LATEST_JPG] = self._always_save_latest_jpg
+            attr[CONF_ALWAYS_SAVE_LATEST_FILE] = self._always_save_latest_file
         return attr
 
     def save_image(self, targets, directory) -> str:
@@ -477,14 +491,14 @@ def save_image(self, targets, directory) -> str:
 
         # Save images, returning the path of saved image as str
         latest_save_path = (
-            directory / f"{get_valid_filename(self._name).lower()}_latest.jpg"
+            directory / f"{get_valid_filename(self._name).lower()}_latest.{self._save_file_format}"
         )
-        _LOGGER.info("Deepstack saved file %s", latest_save_path)
         img.save(latest_save_path)
+        _LOGGER.info("Deepstack saved file %s", latest_save_path)
         saved_image_path = latest_save_path
 
         if self._save_timestamped_file:
-            timestamp_save_path = directory / f"{self._name}_{self._last_detection}.jpg"
+            timestamp_save_path = directory / f"{self._name}_{self._last_detection}.{self._save_file_format}"
             img.save(timestamp_save_path)
             _LOGGER.info("Deepstack saved file %s", timestamp_save_path)
             saved_image_path = timestamp_save_path