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."""
import io
import logging
from pathlib import Path
from PIL import Image, ImageDraw
import simplehound.core as hound
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.core import split_entity_id
import homeassistant.helpers.config_validation as cv
from homeassistant.util.pil import draw_box
_LOGGER = logging.getLogger(__name__)
@ -22,6 +26,7 @@ EVENT_PERSON_DETECTED = "sighthound.person_detected"
ATTR_BOUNDING_BOX = "bounding_box"
ATTR_PEOPLE = "people"
CONF_ACCOUNT_TYPE = "account_type"
CONF_SAVE_FILE_FOLDER = "save_file_folder"
DEV = "dev"
PROD = "prod"
@ -29,6 +34,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
{
vol.Required(CONF_API_KEY): cv.string,
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)
return
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]:
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)
add_entities(entities)
@ -57,7 +67,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
class SighthoundEntity(ImageProcessingEntity):
"""Create a sighthound entity."""
def __init__(self, api, camera_entity, name):
def __init__(self, api, camera_entity, name, save_file_folder):
"""Init."""
self._api = api
self._camera = camera_entity
@ -69,6 +79,7 @@ class SighthoundEntity(ImageProcessingEntity):
self._state = None
self._image_width = None
self._image_height = None
self._save_file_folder = save_file_folder
def process_image(self, image):
"""Process an image."""
@ -81,6 +92,8 @@ class SighthoundEntity(ImageProcessingEntity):
self._image_height = metadata["image_height"]
for person in people:
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):
"""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
def camera_entity(self):
"""Return camera entity id from process pictures."""

View File

@ -1,4 +1,6 @@
"""Tests for the Sighthound integration."""
from copy import deepcopy
import os
from unittest.mock import patch
import pytest
@ -10,6 +12,8 @@ from homeassistant.const import ATTR_ENTITY_ID, CONF_API_KEY
from homeassistant.core import callback
from homeassistant.setup import async_setup_component
TEST_DIR = os.path.dirname(__file__)
VALID_CONFIG = {
ip.DOMAIN: {
"platform": "sighthound",
@ -91,3 +95,23 @@ async def test_process_image(hass, mock_image, mock_detections):
state = hass.states.get(VALID_ENTITY_ID)
assert state.state == "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