catch key error when saving image (#18365)

This commit is contained in:
Jason Hunter 2018-11-11 11:44:41 -05:00 committed by Paulus Schoutsen
parent f7c99ada9d
commit 466d3a5ef8

View File

@ -246,7 +246,8 @@ class TensorFlowImageProcessor(ImageProcessingEntity):
for category, values in matches.items():
# Draw custom category regions/areas
if self._category_areas[category] != [0, 0, 1, 1]:
if (category in self._category_areas
and self._category_areas[category] != [0, 0, 1, 1]):
label = "{} Detection Area".format(category.capitalize())
draw_box(draw, self._category_areas[category], img_width,
img_height, label, (0, 255, 0))