Skip to content
This repository was archived by the owner on Mar 19, 2023. It is now read-only.

Add image scale option #201

Merged
merged 3 commits into from
Jan 22, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ Configuration variables:
- **roi_x_max**: (optional, default 1), range 0-1, must be more than roi_x_min
- **roi_y_min**: (optional, default 0), range 0-1, must be less than roi_y_max
- **roi_y_max**: (optional, default 1), range 0-1, must be more than roi_y_min
- **scale**: (optional, default 1.0), range 0.1-1.0, scaling factor for the images that are pulled from the camera stream. This results in smaller images files, especially when using hires-cameras.
- **source**: Must be a camera.
- **targets**: The list of target object names and/or `object_type`, default `person`. Optionally a `confidence` can be set for this target, if not the default confidence is used. Note the minimum possible confidence is 10%.

Expand Down
46 changes: 37 additions & 9 deletions custom_components/deepstack_object/image_processing.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
CONF_ROI_X_MIN = "roi_x_min"
CONF_ROI_Y_MAX = "roi_y_max"
CONF_ROI_X_MAX = "roi_x_max"
CONF_SCALE = "scale"
CONF_CUSTOM_MODEL = "custom_model"

DATETIME_FORMAT = "%Y-%m-%d_%H-%M-%S"
Expand All @@ -86,6 +87,7 @@
DEFAULT_ROI_Y_MAX = 1.0
DEFAULT_ROI_X_MIN = 0.0
DEFAULT_ROI_X_MAX = 1.0
DEAULT_SCALE = 1.0
DEFAULT_ROI = (
DEFAULT_ROI_Y_MIN,
DEFAULT_ROI_X_MIN,
Expand Down Expand Up @@ -127,6 +129,9 @@
vol.Optional(CONF_ROI_X_MIN, default=DEFAULT_ROI_X_MIN): cv.small_float,
vol.Optional(CONF_ROI_Y_MAX, default=DEFAULT_ROI_Y_MAX): cv.small_float,
vol.Optional(CONF_ROI_X_MAX, default=DEFAULT_ROI_X_MAX): cv.small_float,
vol.Optional(CONF_SCALE, default=DEAULT_SCALE): vol.All(
vol.Coerce(float, vol.Range(min=0.1, max=1))
),
vol.Optional(CONF_SAVE_FILE_FOLDER): cv.isdir,
vol.Optional(CONF_SAVE_TIMESTAMPTED_FILE, default=False): cv.boolean,
vol.Optional(CONF_SHOW_BOXES, default=True): cv.boolean,
Expand Down Expand Up @@ -223,6 +228,7 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
roi_x_min=config[CONF_ROI_X_MIN],
roi_y_max=config[CONF_ROI_Y_MAX],
roi_x_max=config[CONF_ROI_X_MAX],
scale=config[CONF_SCALE],
show_boxes=config[CONF_SHOW_BOXES],
save_file_folder=save_file_folder,
save_timestamped_file=config.get(CONF_SAVE_TIMESTAMPTED_FILE),
Expand All @@ -234,7 +240,7 @@ def setup_platform(hass, config, add_devices, discovery_info=None):


class ObjectClassifyEntity(ImageProcessingEntity):
"""Perform a face classification."""
"""Perform a object classification."""

def __init__(
self,
Expand All @@ -249,6 +255,7 @@ def __init__(
roi_x_min,
roi_y_max,
roi_x_max,
scale,
show_boxes,
save_file_folder,
save_timestamped_file,
Expand Down Expand Up @@ -291,19 +298,34 @@ def __init__(
"y_max": roi_y_max,
"x_max": roi_x_max,
}

self._scale = scale
self._show_boxes = show_boxes
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
self._image = None

def process_image(self, image):
"""Process an image."""
self._image_width, self._image_height = Image.open(
io.BytesIO(bytearray(image))
).size
self._image = Image.open(io.BytesIO(bytearray(image)))
self._image_width, self._image_height = self._image.size

# resize image if different then default
if self._scale != DEAULT_SCALE:
newsize = (self._image_width * self._scale, self._image_width * self._scale)
self._image.thumbnail(newsize, Image.ANTIALIAS)
self._image_width, self._image_height = self._image.size
with io.BytesIO() as output:
self._image.save(output, format="JPEG")
image = output.getvalue()
_LOGGER.debug(
(
f"Image scaled with : {self._scale} W={self._image_width} H={self._image_height}"
)
)

self._state = None
self._objects = [] # The parsed raw data
self._targets_found = []
Expand Down Expand Up @@ -344,7 +366,8 @@ def process_image(self, image):

if self._save_file_folder and self._state > 0:
saved_image_path = self.save_image(
image, self._targets_found, self._save_file_folder,
self._targets_found,
self._save_file_folder,
)

# Fire events
Expand Down Expand Up @@ -396,13 +419,13 @@ def device_state_attributes(self) -> Dict:
attr[CONF_SAVE_TIMESTAMPTED_FILE] = self._save_timestamped_file
return attr

def save_image(self, image, targets, directory) -> str:
def save_image(self, targets, directory) -> str:
"""Draws the actual bounding box of the detected objects.

Returns: saved_image_path, which is the path to the saved timestamped file if configured, else the default saved image.
"""
try:
img = Image.open(io.BytesIO(bytearray(image))).convert("RGB")
img = self._image.convert("RGB")
except UnidentifiedImageError:
_LOGGER.warning("Deepstack unable to process image, bad data")
return
Expand All @@ -411,7 +434,12 @@ def save_image(self, image, targets, directory) -> str:
roi_tuple = tuple(self._roi_dict.values())
if roi_tuple != DEFAULT_ROI and self._show_boxes:
draw_box(
draw, roi_tuple, img.width, img.height, text="ROI", color=GREEN,
draw,
roi_tuple,
img.width,
img.height,
text="ROI",
color=GREEN,
)

for obj in targets:
Expand Down