diff --git a/azure-pipelines-wheels.yml b/azure-pipelines-wheels.yml
index c89435954294ec4bd36c815ce421102207becc88..ebe704f12e2c58b069b89344116fbb532bd86ad2 100644
--- a/azure-pipelines-wheels.yml
+++ b/azure-pipelines-wheels.yml
@@ -89,5 +89,6 @@ jobs:
sed -i "s|# py_noaa|py_noaa|g" ${requirement_file}
sed -i "s|# bme680|bme680|g" ${requirement_file}
sed -i "s|# python-gammu|python-gammu|g" ${requirement_file}
+ sed -i "s|# tf-models-official|tf-models-official|g" ${requirement_file}
done
displayName: 'Prepare requirements files for Home Assistant wheels'
diff --git a/homeassistant/components/tensorflow/image_processing.py b/homeassistant/components/tensorflow/image_processing.py
index f4eb5342c4653e660e922273423064c8614807d8..d6d20c63f56de3bbebf4c8de54db3d08fc999824 100644
--- a/homeassistant/components/tensorflow/image_processing.py
+++ b/homeassistant/components/tensorflow/image_processing.py
@@ -3,9 +3,11 @@ import io
import logging
import os
import sys
+import time
from PIL import Image, ImageDraw, UnidentifiedImageError
import numpy as np
+import tensorflow as tf
import voluptuous as vol
from homeassistant.components.image_processing import (
@@ -16,16 +18,21 @@ from homeassistant.components.image_processing import (
PLATFORM_SCHEMA,
ImageProcessingEntity,
)
+from homeassistant.const import EVENT_HOMEASSISTANT_START
from homeassistant.core import split_entity_id
from homeassistant.helpers import template
import homeassistant.helpers.config_validation as cv
from homeassistant.util.pil import draw_box
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+
+DOMAIN = "tensorflow"
_LOGGER = logging.getLogger(__name__)
ATTR_MATCHES = "matches"
ATTR_SUMMARY = "summary"
ATTR_TOTAL_MATCHES = "total_matches"
+ATTR_PROCESS_TIME = "process_time"
CONF_AREA = "area"
CONF_BOTTOM = "bottom"
@@ -34,6 +41,7 @@ CONF_CATEGORY = "category"
CONF_FILE_OUT = "file_out"
CONF_GRAPH = "graph"
CONF_LABELS = "labels"
+CONF_LABEL_OFFSET = "label_offset"
CONF_LEFT = "left"
CONF_MODEL = "model"
CONF_MODEL_DIR = "model_dir"
@@ -58,12 +66,13 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
vol.Optional(CONF_FILE_OUT, default=[]): vol.All(cv.ensure_list, [cv.template]),
vol.Required(CONF_MODEL): vol.Schema(
{
- vol.Required(CONF_GRAPH): cv.isfile,
+ vol.Required(CONF_GRAPH): cv.isdir,
vol.Optional(CONF_AREA): AREA_SCHEMA,
vol.Optional(CONF_CATEGORIES, default=[]): vol.All(
cv.ensure_list, [vol.Any(cv.string, CATEGORY_SCHEMA)]
),
vol.Optional(CONF_LABELS): cv.isfile,
+ vol.Optional(CONF_LABEL_OFFSET, default=1): int,
vol.Optional(CONF_MODEL_DIR): cv.isdir,
}
),
@@ -71,17 +80,40 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
)
+def get_model_detection_function(model):
+ """Get a tf.function for detection."""
+
+ @tf.function
+ def detect_fn(image):
+ """Detect objects in image."""
+
+ image, shapes = model.preprocess(image)
+ prediction_dict = model.predict(image, shapes)
+ detections = model.postprocess(prediction_dict, shapes)
+
+ return detections
+
+ return detect_fn
+
+
def setup_platform(hass, config, add_entities, discovery_info=None):
"""Set up the TensorFlow image processing platform."""
- model_config = config.get(CONF_MODEL)
+ model_config = config[CONF_MODEL]
model_dir = model_config.get(CONF_MODEL_DIR) or hass.config.path("tensorflow")
labels = model_config.get(CONF_LABELS) or hass.config.path(
"tensorflow", "object_detection", "data", "mscoco_label_map.pbtxt"
)
+ checkpoint = os.path.join(model_config[CONF_GRAPH], "checkpoint")
+ pipeline_config = os.path.join(model_config[CONF_GRAPH], "pipeline.config")
# Make sure locations exist
- if not os.path.isdir(model_dir) or not os.path.exists(labels):
- _LOGGER.error("Unable to locate tensorflow models or label map")
+ if (
+ not os.path.isdir(model_dir)
+ or not os.path.isdir(checkpoint)
+ or not os.path.exists(pipeline_config)
+ or not os.path.exists(labels)
+ ):
+ _LOGGER.error("Unable to locate tensorflow model or label map")
return
# append custom model path to sys.path
@@ -89,18 +121,17 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
try:
# Verify that the TensorFlow Object Detection API is pre-installed
- os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
# These imports shouldn't be moved to the top, because they depend on code from the model_dir.
# (The model_dir is created during the manual setup process. See integration docs.)
- import tensorflow as tf # pylint: disable=import-outside-toplevel
# pylint: disable=import-outside-toplevel
- from object_detection.utils import label_map_util
+ from object_detection.utils import config_util, label_map_util
+ from object_detection.builders import model_builder
except ImportError:
_LOGGER.error(
"No TensorFlow Object Detection library found! Install or compile "
"for your system following instructions here: "
- "https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/installation.md"
+ "https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2.md#installation"
)
return
@@ -113,22 +144,45 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
"PIL at reduced resolution"
)
- # Set up Tensorflow graph, session, and label map to pass to processor
- # pylint: disable=no-member
- detection_graph = tf.Graph()
- with detection_graph.as_default():
- od_graph_def = tf.GraphDef()
- with tf.gfile.GFile(model_config.get(CONF_GRAPH), "rb") as fid:
- serialized_graph = fid.read()
- od_graph_def.ParseFromString(serialized_graph)
- tf.import_graph_def(od_graph_def, name="")
-
- session = tf.Session(graph=detection_graph)
- label_map = label_map_util.load_labelmap(labels)
- categories = label_map_util.convert_label_map_to_categories(
- label_map, max_num_classes=90, use_display_name=True
+ hass.data[DOMAIN] = {CONF_MODEL: None}
+
+ def tensorflow_hass_start(_event):
+ """Set up TensorFlow model on hass start."""
+ start = time.perf_counter()
+
+ # Load pipeline config and build a detection model
+ pipeline_configs = config_util.get_configs_from_pipeline_file(pipeline_config)
+ detection_model = model_builder.build(
+ model_config=pipeline_configs["model"], is_training=False
+ )
+
+ # Restore checkpoint
+ ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)
+ ckpt.restore(os.path.join(checkpoint, "ckpt-0")).expect_partial()
+
+ _LOGGER.debug(
+ "Model checkpoint restore took %d seconds", time.perf_counter() - start
+ )
+
+ model = get_model_detection_function(detection_model)
+
+ # Preload model cache with empty image tensor
+ inp = np.zeros([2160, 3840, 3], dtype=np.uint8)
+ # The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
+ input_tensor = tf.convert_to_tensor(inp, dtype=tf.float32)
+ # The model expects a batch of images, so add an axis with `tf.newaxis`.
+ input_tensor = input_tensor[tf.newaxis, ...]
+ # Run inference
+ model(input_tensor)
+
+ _LOGGER.debug("Model load took %d seconds", time.perf_counter() - start)
+ hass.data[DOMAIN][CONF_MODEL] = model
+
+ hass.bus.listen_once(EVENT_HOMEASSISTANT_START, tensorflow_hass_start)
+
+ category_index = label_map_util.create_category_index_from_labelmap(
+ labels, use_display_name=True
)
- category_index = label_map_util.create_category_index(categories)
entities = []
@@ -138,8 +192,6 @@ def setup_platform(hass, config, add_entities, discovery_info=None):
hass,
camera[CONF_ENTITY_ID],
camera.get(CONF_NAME),
- session,
- detection_graph,
category_index,
config,
)
@@ -152,14 +204,7 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
"""Representation of an TensorFlow image processor."""
def __init__(
- self,
- hass,
- camera_entity,
- name,
- session,
- detection_graph,
- category_index,
- config,
+ self, hass, camera_entity, name, category_index, config,
):
"""Initialize the TensorFlow entity."""
model_config = config.get(CONF_MODEL)
@@ -169,13 +214,12 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
self._name = name
else:
self._name = "TensorFlow {}".format(split_entity_id(camera_entity)[1])
- self._session = session
- self._graph = detection_graph
self._category_index = category_index
self._min_confidence = config.get(CONF_CONFIDENCE)
self._file_out = config.get(CONF_FILE_OUT)
# handle categories and specific detection areas
+ self._label_id_offset = model_config.get(CONF_LABEL_OFFSET)
categories = model_config.get(CONF_CATEGORIES)
self._include_categories = []
self._category_areas = {}
@@ -212,6 +256,7 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
self._matches = {}
self._total_matches = 0
self._last_image = None
+ self._process_time = 0
@property
def camera_entity(self):
@@ -237,6 +282,7 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
category: len(values) for category, values in self._matches.items()
},
ATTR_TOTAL_MATCHES: self._total_matches,
+ ATTR_PROCESS_TIME: self._process_time,
}
def _save_image(self, image, matches, paths):
@@ -281,10 +327,16 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
def process_image(self, image):
"""Process the image."""
+ model = self.hass.data[DOMAIN][CONF_MODEL]
+ if not model:
+ _LOGGER.debug("Model not yet ready.")
+ return
+ start = time.perf_counter()
try:
import cv2 # pylint: disable=import-error, import-outside-toplevel
+ # pylint: disable=no-member
img = cv2.imdecode(np.asarray(bytearray(image)), cv2.IMREAD_UNCHANGED)
inp = img[:, :, [2, 1, 0]] # BGR->RGB
inp_expanded = inp.reshape(1, inp.shape[0], inp.shape[1], 3)
@@ -303,15 +355,15 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
)
inp_expanded = np.expand_dims(inp, axis=0)
- image_tensor = self._graph.get_tensor_by_name("image_tensor:0")
- boxes = self._graph.get_tensor_by_name("detection_boxes:0")
- scores = self._graph.get_tensor_by_name("detection_scores:0")
- classes = self._graph.get_tensor_by_name("detection_classes:0")
- boxes, scores, classes = self._session.run(
- [boxes, scores, classes], feed_dict={image_tensor: inp_expanded}
- )
- boxes, scores, classes = map(np.squeeze, [boxes, scores, classes])
- classes = classes.astype(int)
+ # The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
+ input_tensor = tf.convert_to_tensor(inp_expanded, dtype=tf.float32)
+
+ detections = model(input_tensor)
+ boxes = detections["detection_boxes"][0].numpy()
+ scores = detections["detection_scores"][0].numpy()
+ classes = (
+ detections["detection_classes"][0].numpy() + self._label_id_offset
+ ).astype(int)
matches = {}
total_matches = 0
@@ -367,3 +419,4 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
self._matches = matches
self._total_matches = total_matches
+ self._process_time = time.perf_counter() - start
diff --git a/homeassistant/components/tensorflow/manifest.json b/homeassistant/components/tensorflow/manifest.json
index b7d0361d0875b40c69f6dd33bea4d2b255f3c792..fc87b5cdbff89b86971579f2d367db44d97cf966 100644
--- a/homeassistant/components/tensorflow/manifest.json
+++ b/homeassistant/components/tensorflow/manifest.json
@@ -3,9 +3,12 @@
"name": "TensorFlow",
"documentation": "https://www.home-assistant.io/integrations/tensorflow",
"requirements": [
- "tensorflow==1.13.2",
+ "tensorflow==2.2.0",
+ "tf-slim==1.1.0",
+ "tf-models-official==2.2.1",
+ "pycocotools==2.0.1",
"numpy==1.19.1",
- "protobuf==3.6.1",
+ "protobuf==3.12.2",
"pillow==7.1.2"
],
"codeowners": []
diff --git a/pylintrc b/pylintrc
index df53c2f67a23b8c23113b2b1b00ea9a5cc7473c7..f2860026cd80793b94af4b29a770c1c38cd49ec6 100644
--- a/pylintrc
+++ b/pylintrc
@@ -5,7 +5,7 @@ ignore=tests
jobs=2
load-plugins=pylint_strict_informational
persistent=no
-extension-pkg-whitelist=ciso8601
+extension-pkg-whitelist=ciso8601,cv2
[BASIC]
good-names=id,i,j,k,ex,Run,_,fp,T,ev
diff --git a/requirements_all.txt b/requirements_all.txt
index fea3c46d4a48078022f8bf90957dfacdb9bfcec4..5a9a0ea64c9d55a4e878311608c3dd194f4142c7 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -1120,7 +1120,7 @@ proliphix==0.4.1
prometheus_client==0.7.1
# homeassistant.components.tensorflow
-protobuf==3.6.1
+protobuf==3.12.2
# homeassistant.components.proxmoxve
proxmoxer==1.1.1
@@ -1261,6 +1261,9 @@ pychromecast==7.2.0
# homeassistant.components.cmus
pycmus==0.1.1
+# homeassistant.components.tensorflow
+pycocotools==2.0.1
+
# homeassistant.components.comfoconnect
pycomfoconnect==0.3
@@ -2098,7 +2101,7 @@ temescal==0.1
temperusb==1.5.3
# homeassistant.components.tensorflow
-# tensorflow==1.13.2
+# tensorflow==2.2.0
# homeassistant.components.powerwall
tesla-powerwall==0.2.12
@@ -2106,6 +2109,12 @@ tesla-powerwall==0.2.12
# homeassistant.components.tesla
teslajsonpy==0.10.1
+# homeassistant.components.tensorflow
+# tf-models-official==2.2.1
+
+# homeassistant.components.tensorflow
+tf-slim==1.1.0
+
# homeassistant.components.thermoworks_smoke
thermoworks_smoke==0.1.8
diff --git a/script/gen_requirements_all.py b/script/gen_requirements_all.py
index 4625924da29fc29f1e4b1f69649188eab1085d57..772b9af503440aab0857705bb7c2323cfa2e5a17 100755
--- a/script/gen_requirements_all.py
+++ b/script/gen_requirements_all.py
@@ -41,6 +41,7 @@ COMMENT_REQUIREMENTS = (
"RPi.GPIO",
"smbus-cffi",
"tensorflow",
+ "tf-models-official",
"VL53L1X2",
)