Add doods contains flags on areas to allow specifying overlap (#27035)

* Add support for the contains flags on areas to allow specifying overlap vs contains

* Remove draw_box

* Add timeout option

* Fix import for CONF_TIMEOUT

* Change contains to covers
This commit is contained in:
Zach 2019-10-04 20:48:45 -04:00 committed by Paulus Schoutsen
parent 6ae908b883
commit bbd2078986

View file

@ -7,6 +7,7 @@ import voluptuous as vol
from PIL import Image, ImageDraw
from pydoods import PyDOODS
from homeassistant.const import CONF_TIMEOUT
from homeassistant.components.image_processing import (
CONF_CONFIDENCE,
CONF_ENTITY_ID,
@ -31,6 +32,7 @@ CONF_AUTH_KEY = "auth_key"
CONF_DETECTOR = "detector"
CONF_LABELS = "labels"
CONF_AREA = "area"
CONF_COVERS = "covers"
CONF_TOP = "top"
CONF_BOTTOM = "bottom"
CONF_RIGHT = "right"
@ -43,6 +45,7 @@ AREA_SCHEMA = vol.Schema(
vol.Optional(CONF_LEFT, default=0): cv.small_float,
vol.Optional(CONF_RIGHT, default=1): cv.small_float,
vol.Optional(CONF_TOP, default=0): cv.small_float,
vol.Optional(CONF_COVERS, default=True): cv.boolean,
}
)
@ -50,7 +53,7 @@ LABEL_SCHEMA = vol.Schema(
{
vol.Required(CONF_NAME): cv.string,
vol.Optional(CONF_AREA): AREA_SCHEMA,
vol.Optional(CONF_CONFIDENCE, default=0.0): vol.Range(min=0, max=100),
vol.Optional(CONF_CONFIDENCE): vol.Range(min=0, max=100),
}
)
@ -58,6 +61,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
{
vol.Required(CONF_URL): cv.string,
vol.Required(CONF_DETECTOR): cv.string,
vol.Required(CONF_TIMEOUT, default=90): cv.positive_int,
vol.Optional(CONF_AUTH_KEY, default=""): cv.string,
vol.Optional(CONF_FILE_OUT, default=[]): vol.All(cv.ensure_list, [cv.template]),
vol.Optional(CONF_CONFIDENCE, default=0.0): vol.Range(min=0, max=100),
@ -74,8 +78,9 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
url = config[CONF_URL]
auth_key = config[CONF_AUTH_KEY]
detector_name = config[CONF_DETECTOR]
timeout = config[CONF_TIMEOUT]
doods = PyDOODS(url, auth_key)
doods = PyDOODS(url, auth_key, timeout)
response = doods.get_detectors()
if not isinstance(response, dict):
_LOGGER.warning("Could not connect to doods server: %s", url)
@ -140,6 +145,7 @@ class Doods(ImageProcessingEntity):
# handle labels and specific detection areas
labels = config[CONF_LABELS]
self._label_areas = {}
self._label_covers = {}
for label in labels:
if isinstance(label, dict):
label_name = label[CONF_NAME]
@ -147,14 +153,17 @@ class Doods(ImageProcessingEntity):
_LOGGER.warning("Detector does not support label %s", label_name)
continue
# Label Confidence
label_confidence = label[CONF_CONFIDENCE]
# If label confidence is not specified, use global confidence
label_confidence = label.get(CONF_CONFIDENCE)
if not label_confidence:
label_confidence = confidence
if label_name not in dconfig or dconfig[label_name] > label_confidence:
dconfig[label_name] = label_confidence
# Label area
label_area = label.get(CONF_AREA)
self._label_areas[label_name] = [0, 0, 1, 1]
self._label_covers[label_name] = True
if label_area:
self._label_areas[label_name] = [
label_area[CONF_TOP],
@ -162,6 +171,7 @@ class Doods(ImageProcessingEntity):
label_area[CONF_BOTTOM],
label_area[CONF_RIGHT],
]
self._label_covers[label_name] = label_area[CONF_COVERS]
else:
if label not in detector["labels"] and label != "*":
_LOGGER.warning("Detector does not support label %s", label)
@ -175,6 +185,7 @@ class Doods(ImageProcessingEntity):
# Handle global detection area
self._area = [0, 0, 1, 1]
self._covers = True
area_config = config.get(CONF_AREA)
if area_config:
self._area = [
@ -183,6 +194,7 @@ class Doods(ImageProcessingEntity):
area_config[CONF_BOTTOM],
area_config[CONF_RIGHT],
]
self._covers = area_config[CONF_COVERS]
template.attach(hass, self._file_out)
@ -308,6 +320,7 @@ class Doods(ImageProcessingEntity):
continue
# Exclude matches outside global area definition
if self._covers:
if (
boxes[0] < self._area[0]
or boxes[1] < self._area[1]
@ -315,15 +328,33 @@ class Doods(ImageProcessingEntity):
or boxes[3] > self._area[3]
):
continue
else:
if (
boxes[0] > self._area[2]
or boxes[1] > self._area[3]
or boxes[2] < self._area[0]
or boxes[3] < self._area[1]
):
continue
# Exclude matches outside label specific area definition
if self._label_areas and (
if self._label_areas.get(label):
if self._label_covers[label]:
if (
boxes[0] < self._label_areas[label][0]
or boxes[1] < self._label_areas[label][1]
or boxes[2] > self._label_areas[label][2]
or boxes[3] > self._label_areas[label][3]
):
continue
else:
if (
boxes[0] > self._label_areas[label][2]
or boxes[1] > self._label_areas[label][3]
or boxes[2] < self._label_areas[label][0]
or boxes[3] < self._label_areas[label][1]
):
continue
if label not in matches:
matches[label] = []