Add sighthound save image (#32103)

* Adds save_image

* Update test_image_processing.py

* Update test

* Tidy test

* update image_processing with reviewer comments

* Update test_image_processing.py

* Ammend tests

not passing

* Patch convert

* remove join

* Use valid_config_save_file
This commit is contained in:
Robin 2020-02-23 18:11:05 +00:00 committed by GitHub
parent f6fbecf963
commit 4cc4f070f5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 52 additions and 2 deletions

View File

@ -1,6 +1,9 @@
"""Person detection using Sighthound cloud service.""" """Person detection using Sighthound cloud service."""
import io
import logging import logging
from pathlib import Path
from PIL import Image, ImageDraw
import simplehound.core as hound import simplehound.core as hound
import voluptuous as vol import voluptuous as vol
@ -14,6 +17,7 @@ from homeassistant.components.image_processing import (
from homeassistant.const import ATTR_ENTITY_ID, CONF_API_KEY from homeassistant.const import ATTR_ENTITY_ID, CONF_API_KEY
from homeassistant.core import split_entity_id from homeassistant.core import split_entity_id
import homeassistant.helpers.config_validation as cv import homeassistant.helpers.config_validation as cv
from homeassistant.util.pil import draw_box
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
@ -22,6 +26,7 @@ EVENT_PERSON_DETECTED = "sighthound.person_detected"
ATTR_BOUNDING_BOX = "bounding_box" ATTR_BOUNDING_BOX = "bounding_box"
ATTR_PEOPLE = "people" ATTR_PEOPLE = "people"
CONF_ACCOUNT_TYPE = "account_type" CONF_ACCOUNT_TYPE = "account_type"
CONF_SAVE_FILE_FOLDER = "save_file_folder"
DEV = "dev" DEV = "dev"
PROD = "prod" PROD = "prod"
@ -29,6 +34,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
{ {
vol.Required(CONF_API_KEY): cv.string, vol.Required(CONF_API_KEY): cv.string,
vol.Optional(CONF_ACCOUNT_TYPE, default=DEV): vol.In([DEV, PROD]), vol.Optional(CONF_ACCOUNT_TYPE, default=DEV): vol.In([DEV, PROD]),
vol.Optional(CONF_SAVE_FILE_FOLDER): cv.isdir,
} }
) )
@ -45,10 +51,14 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
_LOGGER.error("Sighthound error %s setup aborted", exc) _LOGGER.error("Sighthound error %s setup aborted", exc)
return return
save_file_folder = config.get(CONF_SAVE_FILE_FOLDER)
if save_file_folder:
save_file_folder = Path(save_file_folder)
entities = [] entities = []
for camera in config[CONF_SOURCE]: for camera in config[CONF_SOURCE]:
sighthound = SighthoundEntity( sighthound = SighthoundEntity(
api, camera[CONF_ENTITY_ID], camera.get(CONF_NAME) api, camera[CONF_ENTITY_ID], camera.get(CONF_NAME), save_file_folder
) )
entities.append(sighthound) entities.append(sighthound)
add_entities(entities) add_entities(entities)
@ -57,7 +67,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
class SighthoundEntity(ImageProcessingEntity): class SighthoundEntity(ImageProcessingEntity):
"""Create a sighthound entity.""" """Create a sighthound entity."""
def __init__(self, api, camera_entity, name): def __init__(self, api, camera_entity, name, save_file_folder):
"""Init.""" """Init."""
self._api = api self._api = api
self._camera = camera_entity self._camera = camera_entity
@ -69,6 +79,7 @@ class SighthoundEntity(ImageProcessingEntity):
self._state = None self._state = None
self._image_width = None self._image_width = None
self._image_height = None self._image_height = None
self._save_file_folder = save_file_folder
def process_image(self, image): def process_image(self, image):
"""Process an image.""" """Process an image."""
@ -81,6 +92,8 @@ class SighthoundEntity(ImageProcessingEntity):
self._image_height = metadata["image_height"] self._image_height = metadata["image_height"]
for person in people: for person in people:
self.fire_person_detected_event(person) self.fire_person_detected_event(person)
if self._save_file_folder and self._state > 0:
self.save_image(image, people, self._save_file_folder)
def fire_person_detected_event(self, person): def fire_person_detected_event(self, person):
"""Send event with detected total_persons.""" """Send event with detected total_persons."""
@ -94,6 +107,19 @@ class SighthoundEntity(ImageProcessingEntity):
}, },
) )
def save_image(self, image, people, directory):
"""Save a timestamped image with bounding boxes around targets."""
img = Image.open(io.BytesIO(bytearray(image))).convert("RGB")
draw = ImageDraw.Draw(img)
for person in people:
box = hound.bbox_to_tf_style(
person["boundingBox"], self._image_width, self._image_height
)
draw_box(draw, box, self._image_width, self._image_height)
latest_save_path = directory / f"{self._name}_latest.jpg"
img.save(latest_save_path)
@property @property
def camera_entity(self): def camera_entity(self):
"""Return camera entity id from process pictures.""" """Return camera entity id from process pictures."""

View File

@ -1,4 +1,6 @@
"""Tests for the Sighthound integration.""" """Tests for the Sighthound integration."""
from copy import deepcopy
import os
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
@ -10,6 +12,8 @@ from homeassistant.const import ATTR_ENTITY_ID, CONF_API_KEY
from homeassistant.core import callback from homeassistant.core import callback
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
TEST_DIR = os.path.dirname(__file__)
VALID_CONFIG = { VALID_CONFIG = {
ip.DOMAIN: { ip.DOMAIN: {
"platform": "sighthound", "platform": "sighthound",
@ -91,3 +95,23 @@ async def test_process_image(hass, mock_image, mock_detections):
state = hass.states.get(VALID_ENTITY_ID) state = hass.states.get(VALID_ENTITY_ID)
assert state.state == "2" assert state.state == "2"
assert len(person_events) == 2 assert len(person_events) == 2
async def test_save_image(hass, mock_image, mock_detections):
"""Save a processed image."""
valid_config_save_file = deepcopy(VALID_CONFIG)
valid_config_save_file[ip.DOMAIN].update({sh.CONF_SAVE_FILE_FOLDER: TEST_DIR})
await async_setup_component(hass, ip.DOMAIN, valid_config_save_file)
assert hass.states.get(VALID_ENTITY_ID)
with patch(
"homeassistant.components.sighthound.image_processing.Image.open"
) as pil_img_open:
pil_img = pil_img_open.return_value
pil_img = pil_img.convert.return_value
data = {ATTR_ENTITY_ID: VALID_ENTITY_ID}
await hass.services.async_call(ip.DOMAIN, ip.SERVICE_SCAN, service_data=data)
await hass.async_block_till_done()
state = hass.states.get(VALID_ENTITY_ID)
assert state.state == "2"
assert pil_img.save.call_count == 1