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

Commit

Permalink
Merge pull request #201 from jodur/master
Browse files Browse the repository at this point in the history
Add image scale option
  • Loading branch information
robmarkcole authored Jan 22, 2021
2 parents 823f3be + 6245cd1 commit df9765f
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 9 deletions.
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

0 comments on commit df9765f

Please sign in to comment.