Add sighthound timestamped file (#32202)
* Update image_processing.py Adds save timestamp file and adds last_detection attribute * Update test_image_processing.py Adds test * Adds assert pil_img.save.call_args * Test timestamp filename * Add test bad data * Update test_image_processing.py * Fix bad image data test * Update homeassistant/components/sighthound/image_processing.py Co-Authored-By: Martin Hjelmare <marhje52@gmail.com> Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
parent
e416f17e4d
commit
3ca97a0517
2 changed files with 111 additions and 9 deletions
|
@ -3,7 +3,7 @@ import io
|
|||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
from PIL import Image, ImageDraw, UnidentifiedImageError
|
||||
import simplehound.core as hound
|
||||
import voluptuous as vol
|
||||
|
||||
|
@ -17,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
|
||||
import homeassistant.util.dt as dt_util
|
||||
from homeassistant.util.pil import draw_box
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
|
@ -27,6 +28,8 @@ ATTR_BOUNDING_BOX = "bounding_box"
|
|||
ATTR_PEOPLE = "people"
|
||||
CONF_ACCOUNT_TYPE = "account_type"
|
||||
CONF_SAVE_FILE_FOLDER = "save_file_folder"
|
||||
CONF_SAVE_TIMESTAMPTED_FILE = "save_timestamped_file"
|
||||
DATETIME_FORMAT = "%Y-%m-%d_%H:%M:%S"
|
||||
DEV = "dev"
|
||||
PROD = "prod"
|
||||
|
||||
|
@ -35,6 +38,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,
|
||||
vol.Optional(CONF_SAVE_TIMESTAMPTED_FILE, default=False): cv.boolean,
|
||||
}
|
||||
)
|
||||
|
||||
|
@ -58,7 +62,11 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
|
|||
entities = []
|
||||
for camera in config[CONF_SOURCE]:
|
||||
sighthound = SighthoundEntity(
|
||||
api, camera[CONF_ENTITY_ID], camera.get(CONF_NAME), save_file_folder
|
||||
api,
|
||||
camera[CONF_ENTITY_ID],
|
||||
camera.get(CONF_NAME),
|
||||
save_file_folder,
|
||||
config[CONF_SAVE_TIMESTAMPTED_FILE],
|
||||
)
|
||||
entities.append(sighthound)
|
||||
add_entities(entities)
|
||||
|
@ -67,7 +75,9 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
|
|||
class SighthoundEntity(ImageProcessingEntity):
|
||||
"""Create a sighthound entity."""
|
||||
|
||||
def __init__(self, api, camera_entity, name, save_file_folder):
|
||||
def __init__(
|
||||
self, api, camera_entity, name, save_file_folder, save_timestamped_file
|
||||
):
|
||||
"""Init."""
|
||||
self._api = api
|
||||
self._camera = camera_entity
|
||||
|
@ -77,15 +87,19 @@ class SighthoundEntity(ImageProcessingEntity):
|
|||
camera_name = split_entity_id(camera_entity)[1]
|
||||
self._name = f"sighthound_{camera_name}"
|
||||
self._state = None
|
||||
self._last_detection = None
|
||||
self._image_width = None
|
||||
self._image_height = None
|
||||
self._save_file_folder = save_file_folder
|
||||
self._save_timestamped_file = save_timestamped_file
|
||||
|
||||
def process_image(self, image):
|
||||
"""Process an image."""
|
||||
detections = self._api.detect(image)
|
||||
people = hound.get_people(detections)
|
||||
self._state = len(people)
|
||||
if self._state > 0:
|
||||
self._last_detection = dt_util.now().strftime(DATETIME_FORMAT)
|
||||
|
||||
metadata = hound.get_metadata(detections)
|
||||
self._image_width = metadata["image_width"]
|
||||
|
@ -109,7 +123,11 @@ 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")
|
||||
try:
|
||||
img = Image.open(io.BytesIO(bytearray(image))).convert("RGB")
|
||||
except UnidentifiedImageError:
|
||||
_LOGGER.warning("Sighthound unable to process image, bad data")
|
||||
return
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
for person in people:
|
||||
|
@ -117,9 +135,15 @@ class SighthoundEntity(ImageProcessingEntity):
|
|||
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)
|
||||
|
||||
if self._save_timestamped_file:
|
||||
timestamp_save_path = directory / f"{self._name}_{self._last_detection}.jpg"
|
||||
img.save(timestamp_save_path)
|
||||
_LOGGER.info("Sighthound saved file %s", timestamp_save_path)
|
||||
|
||||
@property
|
||||
def camera_entity(self):
|
||||
"""Return camera entity id from process pictures."""
|
||||
|
@ -144,3 +168,11 @@ class SighthoundEntity(ImageProcessingEntity):
|
|||
def unit_of_measurement(self):
|
||||
"""Return the unit of measurement."""
|
||||
return ATTR_PEOPLE
|
||||
|
||||
@property
|
||||
def device_state_attributes(self):
|
||||
"""Return the attributes."""
|
||||
attr = {}
|
||||
if self._last_detection:
|
||||
attr["last_person"] = self._last_detection
|
||||
return attr
|
||||
|
|
|
@ -1,8 +1,11 @@
|
|||
"""Tests for the Sighthound integration."""
|
||||
from copy import deepcopy
|
||||
import datetime
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from PIL import UnidentifiedImageError
|
||||
import pytest
|
||||
import simplehound.core as hound
|
||||
|
||||
|
@ -40,11 +43,13 @@ MOCK_DETECTIONS = {
|
|||
"requestId": "545cec700eac4d389743e2266264e84b",
|
||||
}
|
||||
|
||||
MOCK_NOW = datetime.datetime(2020, 2, 20, 10, 5, 3)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_detections():
|
||||
"""Return a mock detection."""
|
||||
with patch(
|
||||
with mock.patch(
|
||||
"simplehound.core.cloud.detect", return_value=MOCK_DETECTIONS
|
||||
) as detection:
|
||||
yield detection
|
||||
|
@ -53,16 +58,35 @@ def mock_detections():
|
|||
@pytest.fixture
|
||||
def mock_image():
|
||||
"""Return a mock camera image."""
|
||||
with patch(
|
||||
with mock.patch(
|
||||
"homeassistant.components.demo.camera.DemoCamera.camera_image",
|
||||
return_value=b"Test",
|
||||
) as image:
|
||||
yield image
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_bad_image_data():
|
||||
"""Mock bad image data."""
|
||||
with mock.patch(
|
||||
"homeassistant.components.sighthound.image_processing.Image.open",
|
||||
side_effect=UnidentifiedImageError,
|
||||
) as bad_data:
|
||||
yield bad_data
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_now():
|
||||
"""Return a mock now datetime."""
|
||||
with mock.patch("homeassistant.util.dt.now", return_value=MOCK_NOW) as now_dt:
|
||||
yield now_dt
|
||||
|
||||
|
||||
async def test_bad_api_key(hass, caplog):
|
||||
"""Catch bad api key."""
|
||||
with patch("simplehound.core.cloud.detect", side_effect=hound.SimplehoundException):
|
||||
with mock.patch(
|
||||
"simplehound.core.cloud.detect", side_effect=hound.SimplehoundException
|
||||
):
|
||||
await async_setup_component(hass, ip.DOMAIN, VALID_CONFIG)
|
||||
assert "Sighthound error" in caplog.text
|
||||
assert not hass.states.get(VALID_ENTITY_ID)
|
||||
|
@ -97,6 +121,21 @@ async def test_process_image(hass, mock_image, mock_detections):
|
|||
assert len(person_events) == 2
|
||||
|
||||
|
||||
async def test_catch_bad_image(
|
||||
hass, caplog, mock_image, mock_detections, mock_bad_image_data
|
||||
):
|
||||
"""Process an 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)
|
||||
|
||||
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()
|
||||
assert "Sighthound unable to process image" in caplog.text
|
||||
|
||||
|
||||
async def test_save_image(hass, mock_image, mock_detections):
|
||||
"""Save a processed image."""
|
||||
valid_config_save_file = deepcopy(VALID_CONFIG)
|
||||
|
@ -104,7 +143,7 @@ async def test_save_image(hass, mock_image, mock_detections):
|
|||
await async_setup_component(hass, ip.DOMAIN, valid_config_save_file)
|
||||
assert hass.states.get(VALID_ENTITY_ID)
|
||||
|
||||
with patch(
|
||||
with mock.patch(
|
||||
"homeassistant.components.sighthound.image_processing.Image.open"
|
||||
) as pil_img_open:
|
||||
pil_img = pil_img_open.return_value
|
||||
|
@ -115,3 +154,34 @@ async def test_save_image(hass, mock_image, mock_detections):
|
|||
state = hass.states.get(VALID_ENTITY_ID)
|
||||
assert state.state == "2"
|
||||
assert pil_img.save.call_count == 1
|
||||
|
||||
directory = Path(TEST_DIR)
|
||||
latest_save_path = directory / "sighthound_demo_camera_latest.jpg"
|
||||
assert pil_img.save.call_args_list[0] == mock.call(latest_save_path)
|
||||
|
||||
|
||||
async def test_save_timestamped_image(hass, mock_image, mock_detections, mock_now):
|
||||
"""Save a processed image."""
|
||||
valid_config_save_ts_file = deepcopy(VALID_CONFIG)
|
||||
valid_config_save_ts_file[ip.DOMAIN].update({sh.CONF_SAVE_FILE_FOLDER: TEST_DIR})
|
||||
valid_config_save_ts_file[ip.DOMAIN].update({sh.CONF_SAVE_TIMESTAMPTED_FILE: True})
|
||||
await async_setup_component(hass, ip.DOMAIN, valid_config_save_ts_file)
|
||||
assert hass.states.get(VALID_ENTITY_ID)
|
||||
|
||||
with mock.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 == 2
|
||||
|
||||
directory = Path(TEST_DIR)
|
||||
timestamp_save_path = (
|
||||
directory / "sighthound_demo_camera_2020-02-20_10:05:03.jpg"
|
||||
)
|
||||
assert pil_img.save.call_args_list[1] == mock.call(timestamp_save_path)
|
||||
|
|
Loading…
Add table
Reference in a new issue