From 87626e406c7cec5a377115bc03f19ab1d3b27452 Mon Sep 17 00:00:00 2001 From: Luke Date: Sat, 25 Mar 2023 15:30:08 -0400 Subject: [PATCH] feat: Moved map parser to package for future --- pyproject.toml | 1 + roborock/map_parser/__init__.py | 5 + roborock/map_parser/image_handler.py | 638 ++++++++++++++++++++++++ roborock/map_parser/map_data.py | 377 ++++++++++++++ roborock/map_parser/map_data_parser.py | 539 ++++++++++++++++++++ roborock/map_parser/map_parser_const.py | 276 ++++++++++ roborock/map_parser/types.py | 9 + 7 files changed, 1845 insertions(+) create mode 100644 roborock/map_parser/__init__.py create mode 100644 roborock/map_parser/image_handler.py create mode 100644 roborock/map_parser/map_data.py create mode 100644 roborock/map_parser/map_data_parser.py create mode 100644 roborock/map_parser/map_parser_const.py create mode 100644 roborock/map_parser/types.py diff --git a/pyproject.toml b/pyproject.toml index d65581c5..9ac7e8d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ async-timeout = "*" pycryptodome = "~3.17.0" pycryptodomex = {version = "~3.17.0", markers = "sys_platform == 'darwin'"} paho-mqtt = "~1.6.1" +Pillow = "*" [build-system] diff --git a/roborock/map_parser/__init__.py b/roborock/map_parser/__init__.py new file mode 100644 index 00000000..1d292c4f --- /dev/null +++ b/roborock/map_parser/__init__.py @@ -0,0 +1,5 @@ +from .image_handler import * +from .map_data import * +from .map_data_parser import * +from .map_parser_const import * +from .types import * diff --git a/roborock/map_parser/image_handler.py b/roborock/map_parser/image_handler.py new file mode 100644 index 00000000..c871a330 --- /dev/null +++ b/roborock/map_parser/image_handler.py @@ -0,0 +1,638 @@ +import logging +import math +from typing import Tuple, List, Dict, Callable, Set + +from PIL import Image, ImageDraw, ImageFont +from PIL.Image import Image as ImageType + +from .map_data import ( + ImageData, + Path, + Area, + Wall, + Zone, + Point, + Obstacle, + Room, +) +from .map_parser_const import * +from .types import ( + Colors, + ImageConfig, + Sizes, + Color, + Texts, +) + +_LOGGER = logging.getLogger(__name__) + + +class ImageHandlerRoborock: + MAP_OUTSIDE = 0x00 + MAP_WALL = 0x01 + MAP_INSIDE = 0xFF + MAP_SCAN = 0x07 + + COLORS = { + COLOR_MAP_INSIDE: (32, 115, 185), + COLOR_MAP_OUTSIDE: (0, 0, 0, 0), + COLOR_MAP_WALL: (109, 110, 112), + COLOR_MAP_WALL_V2: (109, 110, 112), + COLOR_GREY_WALL: (0, 0, 0, 0), + COLOR_CLEANED_AREA: (127, 127, 127, 127), + COLOR_PATH: (255, 255, 255), + COLOR_MOP_PATH: (255, 255, 255, 0x5F), + COLOR_GOTO_PATH: (0, 255, 0), + COLOR_PREDICTED_PATH: (255, 255, 0), + COLOR_ZONES: (0xAD, 0xD8, 0xFF, 0x8F), + COLOR_ZONES_OUTLINE: (0xAD, 0xD8, 0xFF), + COLOR_VIRTUAL_WALLS: (255, 0, 0), + COLOR_NEW_DISCOVERED_AREA: (64, 64, 64), + COLOR_CARPETS: (0, 0, 0, 51), + COLOR_NO_CARPET_ZONES: (255, 33, 55, 127), + COLOR_NO_CARPET_ZONES_OUTLINE: (255, 0, 0), + COLOR_NO_GO_ZONES: (255, 94, 73, 102), + COLOR_NO_GO_ZONES_OUTLINE: (255, 94, 73), + COLOR_NO_MOPPING_ZONES: (163, 130, 211, 127), + COLOR_NO_MOPPING_ZONES_OUTLINE: (163, 130, 211), + COLOR_CHARGER: (86, 85, 210), + COLOR_CHARGER_OUTLINE: (255, 255, 255), + COLOR_ROBO: (0xFF, 0xFF, 0xFF), + COLOR_ROBO_OUTLINE: (0, 0, 0), + COLOR_ROOM_NAMES: (0, 0, 0), + COLOR_OBSTACLE: (63, 159, 254), + COLOR_IGNORED_OBSTACLE: (63, 159, 254), + COLOR_OBSTACLE_WITH_PHOTO: (63, 159, 254), + COLOR_IGNORED_OBSTACLE_WITH_PHOTO: (63, 159, 254), + COLOR_OBSTACLE_OUTLINE: (255, 255, 255), + COLOR_UNKNOWN: (0, 0, 0), + COLOR_SCAN: (0xDF, 0xDF, 0xDF), + COLOR_ROOM_1: (240, 178, 122), + COLOR_ROOM_2: (133, 193, 233), + COLOR_ROOM_3: (217, 136, 128), + COLOR_ROOM_4: (52, 152, 219), + COLOR_ROOM_5: (205, 97, 85), + COLOR_ROOM_6: (243, 156, 18), + COLOR_ROOM_7: (88, 214, 141), + COLOR_ROOM_8: (245, 176, 65), + COLOR_ROOM_9: (252, 212, 81), + COLOR_ROOM_10: (72, 201, 176), + COLOR_ROOM_11: (84, 153, 199), + COLOR_ROOM_12: (133, 193, 233), + COLOR_ROOM_13: (245, 176, 65), + COLOR_ROOM_14: (82, 190, 128), + COLOR_ROOM_15: (72, 201, 176), + COLOR_ROOM_16: (165, 105, 189), + } + ROOM_COLORS = [ + COLOR_ROOM_1, + COLOR_ROOM_2, + COLOR_ROOM_3, + COLOR_ROOM_4, + COLOR_ROOM_5, + COLOR_ROOM_6, + COLOR_ROOM_7, + COLOR_ROOM_8, + COLOR_ROOM_9, + COLOR_ROOM_10, + COLOR_ROOM_11, + COLOR_ROOM_12, + COLOR_ROOM_13, + COLOR_ROOM_14, + COLOR_ROOM_15, + COLOR_ROOM_16, + ] + + @staticmethod + def create_empty_map_image(colors: Colors, text: str = "NO MAP") -> ImageType: + color = ImageHandlerRoborock.__get_color__(COLOR_MAP_OUTSIDE, colors) + image = Image.new("RGBA", (300, 200), color=color) + if sum(color[0:3]) > 382: + text_color = (0, 0, 0) + else: + text_color = (255, 255, 255) + draw = ImageDraw.Draw(image, "RGBA") + w, h = draw.textsize(text) + draw.text( + ((image.size[0] - w) / 2, (image.size[1] - h) / 2), text, fill=text_color + ) + return image + + @staticmethod + def draw_path( + image: ImageData, path: Path, sizes: Sizes, colors: Colors, scale: float + ): + ImageHandlerRoborock.__draw_path__( + image, + path, + sizes[CONF_SIZE_PATH_WIDTH], + ImageHandlerRoborock.__get_color__(COLOR_PATH, colors), + scale, + ) + + @staticmethod + def draw_goto_path( + image: ImageData, path: Path, sizes: Sizes, colors: Colors, scale: float + ): + ImageHandlerRoborock.__draw_path__( + image, + path, + sizes[CONF_SIZE_PATH_WIDTH], + ImageHandlerRoborock.__get_color__(COLOR_GOTO_PATH, colors), + scale, + ) + + @staticmethod + def draw_predicted_path( + image: ImageData, path: Path, sizes: Sizes, colors: Colors, scale: float + ): + ImageHandlerRoborock.__draw_path__( + image, + path, + sizes[CONF_SIZE_PATH_WIDTH], + ImageHandlerRoborock.__get_color__(COLOR_PREDICTED_PATH, colors), + scale, + ) + + @staticmethod + def draw_mop_path( + image: ImageData, path: Path, sizes: Sizes, colors: Colors, scale: float + ): + ImageHandlerRoborock.__draw_path__( + image, + path, + sizes[CONF_SIZE_MOP_PATH_WIDTH], + ImageHandlerRoborock.__get_color__(COLOR_MOP_PATH, colors), + scale, + ) + + @staticmethod + def draw_no_carpet_areas(image: ImageData, areas: List[Area], colors: Colors): + ImageHandlerRoborock.__draw_areas__( + image, + areas, + ImageHandlerRoborock.__get_color__(COLOR_NO_CARPET_ZONES, colors), + ImageHandlerRoborock.__get_color__(COLOR_NO_CARPET_ZONES_OUTLINE, colors), + ) + + @staticmethod + def draw_no_go_areas(image: ImageData, areas: List[Area], colors: Colors): + ImageHandlerRoborock.__draw_areas__( + image, + areas, + ImageHandlerRoborock.__get_color__(COLOR_NO_GO_ZONES, colors), + ImageHandlerRoborock.__get_color__(COLOR_NO_GO_ZONES_OUTLINE, colors), + ) + + @staticmethod + def draw_no_mopping_areas(image: ImageData, areas: List[Area], colors: Colors): + ImageHandlerRoborock.__draw_areas__( + image, + areas, + ImageHandlerRoborock.__get_color__(COLOR_NO_MOPPING_ZONES, colors), + ImageHandlerRoborock.__get_color__(COLOR_NO_MOPPING_ZONES_OUTLINE, colors), + ) + + @staticmethod + def draw_walls(image: ImageData, walls: List[Wall], colors: Colors): + draw = ImageDraw.Draw(image.data, "RGBA") + for wall in walls: + draw.line( + wall.to_img(image.dimensions).as_list(), + ImageHandlerRoborock.__get_color__(COLOR_VIRTUAL_WALLS, colors), + width=2, + ) + + @staticmethod + def draw_zones(image: ImageData, zones: List[Zone], colors: Colors): + areas = [z.as_area() for z in zones] + ImageHandlerRoborock.__draw_areas__( + image, + areas, + ImageHandlerRoborock.__get_color__(COLOR_ZONES, colors), + ImageHandlerRoborock.__get_color__(COLOR_ZONES_OUTLINE, colors), + ) + + @staticmethod + def draw_charger(image: ImageData, charger: Point, sizes: Sizes, colors: Colors): + color = ImageHandlerRoborock.__get_color__(COLOR_CHARGER, colors) + outline = ImageHandlerRoborock.__get_color__(COLOR_CHARGER_OUTLINE, colors) + radius = sizes[CONF_SIZE_CHARGER_RADIUS] + ImageHandlerRoborock.__draw_pieslice__(image, charger, radius, outline, color) + + @staticmethod + def draw_obstacles(image: ImageData, obstacles, sizes: Sizes, colors: Colors): + color = ImageHandlerRoborock.__get_color__(COLOR_OBSTACLE, colors) + radius = sizes[CONF_SIZE_OBSTACLE_RADIUS] + ImageHandlerRoborock.draw_all_obstacles(image, obstacles, radius, color, colors) + + @staticmethod + def draw_ignored_obstacles( + image: ImageData, obstacles: List[Obstacle], sizes: Sizes, colors: Colors + ): + color = ImageHandlerRoborock.__get_color__(COLOR_IGNORED_OBSTACLE, colors) + radius = sizes[CONF_SIZE_IGNORED_OBSTACLE_RADIUS] + ImageHandlerRoborock.draw_all_obstacles(image, obstacles, radius, color, colors) + + @staticmethod + def draw_obstacles_with_photo( + image: ImageData, obstacles: List[Obstacle], sizes: Sizes, colors: Colors + ): + color = ImageHandlerRoborock.__get_color__(COLOR_OBSTACLE_WITH_PHOTO, colors) + radius = sizes[CONF_SIZE_OBSTACLE_WITH_PHOTO_RADIUS] + ImageHandlerRoborock.draw_all_obstacles(image, obstacles, radius, color, colors) + + @staticmethod + def draw_ignored_obstacles_with_photo( + image: ImageData, obstacles: List[Obstacle], sizes: Sizes, colors: Colors + ): + color = ImageHandlerRoborock.__get_color__( + COLOR_IGNORED_OBSTACLE_WITH_PHOTO, colors + ) + radius = sizes[CONF_SIZE_IGNORED_OBSTACLE_WITH_PHOTO_RADIUS] + ImageHandlerRoborock.draw_all_obstacles(image, obstacles, radius, color, colors) + + @staticmethod + def draw_all_obstacles( + image: ImageData, + obstacles: List[Obstacle], + radius: float, + color: Color, + colors: Colors, + ): + outline_color = ImageHandlerRoborock.__get_color__( + COLOR_OBSTACLE_OUTLINE, colors + ) + for obstacle in obstacles: + ImageHandlerRoborock.__draw_circle__( + image, obstacle, radius, outline_color, color + ) + + @staticmethod + def draw_vacuum_position( + image: ImageData, vacuum_position: Point, sizes: Sizes, colors: Colors + ): + color = ImageHandlerRoborock.__get_color__(COLOR_ROBO, colors) + outline = ImageHandlerRoborock.__get_color__(COLOR_ROBO_OUTLINE, colors) + radius = sizes[CONF_SIZE_VACUUM_RADIUS] + ImageHandlerRoborock.__draw_vacuum__( + image, vacuum_position, radius, outline, color + ) + + @staticmethod + def draw_room_names(image: ImageData, rooms: Dict[int, Room], colors: Colors): + color = ImageHandlerRoborock.__get_color__(COLOR_ROOM_NAMES, colors) + for room in rooms.values(): + p = room.point() + if p is not None: + point = p.to_img(image.dimensions) + ImageHandlerRoborock.__draw_text__( + image, room.name, point.x, point.y, color + ) + + @staticmethod + def rotate(image: ImageData): + if image.dimensions.rotation == 90: + image.data = image.data.transpose(Image.ROTATE_90) + if image.dimensions.rotation == 180: + image.data = image.data.transpose(Image.ROTATE_180) + if image.dimensions.rotation == 270: + image.data = image.data.transpose(Image.ROTATE_270) + + @staticmethod + def draw_texts(image: ImageData, texts: Texts): + for text_config in texts: + x = text_config[CONF_X] * image.data.size[0] / 100 + y = text_config[CONF_Y] * image.data.size[1] / 100 + ImageHandlerRoborock.__draw_text__( + image, + text_config[CONF_TEXT], + x, + y, + text_config[CONF_COLOR], + text_config[CONF_FONT], + text_config[CONF_FONT_SIZE], + ) + + @staticmethod + def draw_layer(image: ImageData, layer_name: str): + ImageHandlerRoborock.__draw_layer__(image, image.additional_layers[layer_name]) + + @staticmethod + def __use_transparency__(*colors): + return any(len(color) > 3 for color in colors) + + @staticmethod + def __draw_vacuum__(image: ImageData, vacuum_pos, r, outline, fill): + def draw_func(draw: ImageDraw): + if vacuum_pos.a is None: + vacuum_pos.a = 0 + point = vacuum_pos.to_img(image.dimensions) + r_scaled = r / 16 + # main outline + coords = [point.x - r, point.y - r, point.x + r, point.y + r] + draw.ellipse(coords, outline=outline, fill=fill) + if r >= 8: + # secondary outline + r2 = r_scaled * 14 + x = point.x + y = point.y + coords = [x - r2, y - r2, x + r2, y + r2] + draw.ellipse(coords, outline=outline, fill=None) + # bin cover + a1 = (vacuum_pos.a + 104) / 180 * math.pi + a2 = (vacuum_pos.a - 104) / 180 * math.pi + r2 = r_scaled * 13 + x1 = point.x - r2 * math.cos(a1) + y1 = point.y + r2 * math.sin(a1) + x2 = point.x - r2 * math.cos(a2) + y2 = point.y + r2 * math.sin(a2) + draw.line([x1, y1, x2, y2], width=1, fill=outline) + # lidar + angle = vacuum_pos.a / 180 * math.pi + r2 = r_scaled * 3 + x = point.x + r2 * math.cos(angle) + y = point.y - r2 * math.sin(angle) + r2 = r_scaled * 4 + coords = [x - r2, y - r2, x + r2, y + r2] + draw.ellipse(coords, outline=outline, fill=fill) + # button + half_color = ( + (outline[0] + fill[0]) // 2, + (outline[1] + fill[1]) // 2, + (outline[2] + fill[2]) // 2, + ) + r2 = r_scaled * 10 + x = point.x + r2 * math.cos(angle) + y = point.y - r2 * math.sin(angle) + r2 = r_scaled * 2 + coords = [x - r2, y - r2, x + r2, y + r2] + draw.ellipse(coords, outline=half_color, fill=half_color) + + ImageHandlerRoborock.__draw_on_new_layer__( + image, + draw_func, + 1, + ImageHandlerRoborock.__use_transparency__(outline, fill), + ) + + @staticmethod + def __draw_circle__( + image: ImageData, center: Point, r: float, outline: Color, fill: Color + ): + def draw_func(draw: ImageDraw): + point = center.to_img(image.dimensions) + coords = [point.x - r, point.y - r, point.x + r, point.y + r] + draw.ellipse(coords, outline=outline, fill=fill) + + ImageHandlerRoborock.__draw_on_new_layer__( + image, + draw_func, + 1, + ImageHandlerRoborock.__use_transparency__(outline, fill), + ) + + @staticmethod + def __draw_pieslice__(image: ImageData, position, r, outline, fill): + def draw_func(draw: ImageDraw): + point = position.to_img(image.dimensions) + angle = -position.a if position.a is not None else 0 + coords = [point.x - r, point.y - r, point.x + r, point.y + r] + draw.pieslice(coords, angle + 90, angle - 90, outline="black", fill=fill) + + ImageHandlerRoborock.__draw_on_new_layer__( + image, + draw_func, + 1, + ImageHandlerRoborock.__use_transparency__(outline, fill), + ) + + @staticmethod + def __draw_areas__( + image: ImageData, areas: List[Area], fill: Color, outline: Color + ): + if len(areas) == 0: + return + + use_transparency = ImageHandlerRoborock.__use_transparency__(outline, fill) + for area in areas: + current_area = area + + def draw_func(draw: ImageDraw): + draw.polygon( + current_area.to_img(image.dimensions).as_list(), fill, outline + ) + + ImageHandlerRoborock.__draw_on_new_layer__( + image, draw_func, 1, use_transparency + ) + + @staticmethod + def __draw_path__( + image: ImageData, path: Path, path_width: int, color: Color, scale: float + ): + if len(path.path) < 1: + return + + def draw_func(draw: ImageDraw): + for current_path in path.path: + if len(current_path) > 1: + s = current_path[0].to_img(image.dimensions) + coords = None + for point in current_path[1:]: + e = point.to_img(image.dimensions) + sx = s.x * scale + sy = s.y * scale + ex = e.x * scale + ey = e.y * scale + draw.line( + [sx, sy, ex, ey], width=int(scale * path_width), fill=color + ) + if path_width > 4: + r = scale * path_width / 2 + if not coords: + coords = [sx - r, sy - r, sx + r, sy + r] + draw.pieslice(coords, 0, 360, outline=color, fill=color) + coords = [ex - r, ey - r, ex + r, ey + r] + draw.pieslice(coords, 0, 360, outline=color, fill=color) + s = e + + ImageHandlerRoborock.__draw_on_new_layer__( + image, draw_func, scale, ImageHandlerRoborock.__use_transparency__(color) + ) + + @staticmethod + def __draw_text__( + image: ImageData, + text: str, + x: float, + y: float, + color: Color, + font_file=None, + font_size=None, + ): + def draw_func(draw: ImageDraw): + font = ImageFont.load_default() + try: + if font_file is not None and font_size > 0: + font = ImageFont.truetype(font_file, font_size) + except OSError: + _LOGGER.warning("Unable to find font file: %s", font_file) + except ImportError: + _LOGGER.warning("Unable to open font: %s", font_file) + finally: + w, h = draw.textsize(text, font) + draw.text((x - w / 2, y - h / 2), text, font=font, fill=color) + + ImageHandlerRoborock.__draw_on_new_layer__( + image, draw_func, 1, ImageHandlerRoborock.__use_transparency__(color) + ) + + @staticmethod + def __get_color__(name, colors: Colors, default_name: str = None) -> Color: + if name in colors: + return colors[name] + if default_name is None: + return ImageHandlerRoborock.COLORS[name] + return ImageHandlerRoborock.COLORS[default_name] + + @staticmethod + def __draw_on_new_layer__( + image: ImageData, + draw_function: Callable, + scale: float = 1, + use_transparency=False, + ): + if scale == 1 and not use_transparency: + draw = ImageDraw.Draw(image.data, "RGBA") + draw_function(draw) + else: + size = (int(image.data.size[0] * scale), int(image.data.size[1] * scale)) + layer = Image.new("RGBA", size, (255, 255, 255, 0)) + draw = ImageDraw.Draw(layer, "RGBA") + draw_function(draw) + if scale != 1: + layer = layer.resize(image.data.size, resample=Image.BOX) + ImageHandlerRoborock.__draw_layer__(image, layer) + + @staticmethod + def __draw_layer__(image: ImageData, layer: ImageType): + image.data = Image.alpha_composite(image.data, layer) + + @staticmethod + def parse( + raw_data: bytes, + width: int, + height: int, + carpet_map: Set[int], + colors: Colors, + image_config: ImageConfig, + ) -> Tuple[ImageType, dict]: + rooms = {} + scale = image_config[CONF_SCALE] + trim_left = int(image_config[CONF_TRIM][CONF_LEFT] * width / 100) + trim_right = int(image_config[CONF_TRIM][CONF_RIGHT] * width / 100) + trim_top = int(image_config[CONF_TRIM][CONF_TOP] * height / 100) + trim_bottom = int(image_config[CONF_TRIM][CONF_BOTTOM] * height / 100) + trimmed_height = height - trim_top - trim_bottom + trimmed_width = width - trim_left - trim_right + image = Image.new("RGBA", (trimmed_width, trimmed_height)) + if width == 0 or height == 0: + return ImageHandlerRoborock.create_empty_map_image(colors), {} + pixels = image.load() + for img_y in range(trimmed_height): + for img_x in range(trimmed_width): + idx = img_x + trim_left + width * (img_y + trim_bottom) + pixel_type = raw_data[idx] + x = img_x + y = trimmed_height - img_y - 1 + if pixel_type == ImageHandlerRoborock.MAP_OUTSIDE: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_MAP_OUTSIDE, colors + ) + elif pixel_type == ImageHandlerRoborock.MAP_WALL: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_MAP_WALL, colors + ) + elif pixel_type == ImageHandlerRoborock.MAP_INSIDE: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_MAP_INSIDE, colors + ) + elif pixel_type == ImageHandlerRoborock.MAP_SCAN: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_SCAN, colors + ) + else: + obstacle = pixel_type & 0x07 + if obstacle == 0: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_GREY_WALL, colors + ) + elif obstacle == 1: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_MAP_WALL_V2, colors + ) + elif obstacle == 7: + room_number = (pixel_type & 0xFF) >> 3 + room_x = img_x + trim_left + room_y = img_y + trim_bottom + if room_number not in rooms: + rooms[room_number] = (room_x, room_y, room_x, room_y) + else: + rooms[room_number] = ( + min(rooms[room_number][0], room_x), + min(rooms[room_number][1], room_y), + max(rooms[room_number][2], room_x), + max(rooms[room_number][3], room_y), + ) + default = ImageHandlerRoborock.ROOM_COLORS[room_number >> 1] + pixels[x, y] = ImageHandlerRoborock.__get_color__( + f"{COLOR_ROOM_PREFIX}{room_number}", colors, default + ) + else: + pixels[x, y] = ImageHandlerRoborock.__get_color__( + COLOR_UNKNOWN, colors + ) + + if idx in carpet_map and (x + y) % 2: + + def combine_color_component(base: int, overlay: int, alpha: int): + return int((base * (255 - alpha) + overlay * alpha) / 255) + + carpet_color = ImageHandlerRoborock.__get_color__( + COLOR_CARPETS, colors + ) + + if not pixels[x, y] or len(carpet_color) != 4: + pixels[x, y] = carpet_color + else: + pixels[x, y] = ( + combine_color_component( + pixels[x, y][0], carpet_color[0], carpet_color[3] + ), + combine_color_component( + pixels[x, y][1], carpet_color[1], carpet_color[3] + ), + combine_color_component( + pixels[x, y][2], carpet_color[2], carpet_color[3] + ), + ) + + if image_config["scale"] != 1 and width != 0 and height != 0: + image = image.resize( + (int(trimmed_width * scale), int(trimmed_height * scale)), + resample=Image.NEAREST, + ) + return image, rooms + + @staticmethod + def get_room_at_pixel(raw_data: bytes, width: int, x: int, y: int) -> int: + room_number = None + pixel_type = raw_data[x + width * y] + if pixel_type not in [ + ImageHandlerRoborock.MAP_INSIDE, + ImageHandlerRoborock.MAP_SCAN, + ]: + if pixel_type & 0x07 == 7: + room_number = (pixel_type & 0xFF) >> 3 + return room_number diff --git a/roborock/map_parser/map_data.py b/roborock/map_parser/map_data.py new file mode 100644 index 00000000..2e5f9467 --- /dev/null +++ b/roborock/map_parser/map_data.py @@ -0,0 +1,377 @@ +from __future__ import annotations + +from typing import Any, Callable, Dict, List, Optional, Set + +from PIL.Image import Image as ImageType + +from .map_parser_const import * +from .types import ( + CalibrationPoints, + ImageConfig, +) + + +class Point: + def __init__(self, x: float, y: float, a=None): + self.x = x + self.y = y + self.a = a + + def __str__(self) -> str: + if self.a is None: + return f"({self.x}, {self.y})" + return f"({self.x}, {self.y}, a = {self.a})" + + def __repr__(self) -> str: + return self.__str__() + + def __eq__(self, other: Point) -> bool: + return ( + other is not None + and self.x == other.x + and self.y == other.y + and self.a == other.a + ) + + def as_dict(self) -> Dict[str, Any]: + if self.a is None: + return {ATTR_X: self.x, ATTR_Y: self.y} + return {ATTR_X: self.x, ATTR_Y: self.y, ATTR_A: self.a} + + def to_img(self, image_dimensions) -> Point: + return image_dimensions.to_img(self) + + def rotated(self, image_dimensions) -> Point: + alpha = image_dimensions.rotation + w = int(image_dimensions.width * image_dimensions.scale) + h = int(image_dimensions.height * image_dimensions.scale) + x = self.x + y = self.y + while alpha > 0: + tmp = y + y = w - x + x = tmp + tmp = h + h = w + w = tmp + alpha = alpha - 90 + return Point(x, y) + + def __mul__(self, other) -> Point: + return Point(self.x * other, self.y * other, self.a) + + def __truediv__(self, other) -> Point: + return Point(self.x / other, self.y / other, self.a) + + +class Obstacle(Point): + def __init__(self, x: float, y: float, details: Dict[str, Any]): + super().__init__(x, y) + self.details = details + + def as_dict(self) -> Dict[str, Any]: + return {**super(Obstacle, self).as_dict(), **self.details} + + def __str__(self) -> str: + return f"({self.x}, {self.y}, details = {self.details})" + + +class ImageDimensions: + def __init__( + self, + top: int, + left: int, + height: int, + width: int, + scale: float, + rotation: int, + img_transformation: Callable[[Point], Point], + ): + self.top = top + self.left = left + self.height = height + self.width = width + self.scale = scale + self.rotation = rotation + self.img_transformation = img_transformation + + def to_img(self, point: Point) -> Point: + p = self.img_transformation(point) + return Point( + (p.x - self.left) * self.scale, + (self.height - (p.y - self.top) - 1) * self.scale, + ) + + +class ImageData: + def __init__( + self, + size: int, + top: int, + left: int, + height: int, + width: int, + image_config: ImageConfig, + data: ImageType, + img_transformation: Callable[[Point], Point], + additional_layers: dict = None, + ): + trim_left = int(image_config[CONF_TRIM][CONF_LEFT] * width / 100) + trim_right = int(image_config[CONF_TRIM][CONF_RIGHT] * width / 100) + trim_top = int(image_config[CONF_TRIM][CONF_TOP] * height / 100) + trim_bottom = int(image_config[CONF_TRIM][CONF_BOTTOM] * height / 100) + scale = image_config[CONF_SCALE] + rotation = image_config[CONF_ROTATE] + self.size = size + self.dimensions = ImageDimensions( + top + trim_bottom, + left + trim_left, + height - trim_top - trim_bottom, + width - trim_left - trim_right, + scale, + rotation, + img_transformation, + ) + self.is_empty = height == 0 or width == 0 + self.data = data + if additional_layers is None: + self.additional_layers = {} + else: + self.additional_layers = dict( + filter(lambda l: l[1] is not None, additional_layers.items()) + ) + + def as_dict(self) -> Dict[str, Any]: + return { + ATTR_SIZE: self.size, + ATTR_OFFSET_Y: self.dimensions.top, + ATTR_OFFSET_X: self.dimensions.left, + ATTR_HEIGHT: self.dimensions.height, + ATTR_SCALE: self.dimensions.scale, + ATTR_ROTATION: self.dimensions.rotation, + ATTR_WIDTH: self.dimensions.width, + } + + @staticmethod + def create_empty(data: ImageType) -> ImageData: + image_config = { + CONF_TRIM: {CONF_LEFT: 0, CONF_RIGHT: 0, CONF_TOP: 0, CONF_BOTTOM: 0}, + CONF_SCALE: 2, + CONF_ROTATE: 0, + } + return ImageData(0, 0, 0, 0, 0, image_config, data, lambda p: p) + + +class Path: + def __init__( + self, + point_length: Optional[int], + point_size: Optional[int], + angle: Optional[int], + path: List[List[Point]], + ): + self.point_length = point_length + self.point_size = point_size + self.angle = angle + self.path = path + + def as_dict(self) -> Dict[str, Any]: + return { + ATTR_POINT_LENGTH: self.point_length, + ATTR_POINT_SIZE: self.point_size, + ATTR_ANGLE: self.angle, + ATTR_PATH: self.path, + } + + +class Zone: + def __init__(self, x0: float, y0: float, x1: float, y1: float): + self.x0 = x0 + self.y0 = y0 + self.x1 = x1 + self.y1 = y1 + + def __str__(self) -> str: + return f"[{self.x0}, {self.y0}, {self.x1}, {self.y1}]" + + def __repr__(self) -> str: + return self.__str__() + + def as_dict(self) -> Dict[str, Any]: + return {ATTR_X0: self.x0, ATTR_Y0: self.y0, ATTR_X1: self.x1, ATTR_Y1: self.y1} + + def as_area(self) -> Area: + return Area( + self.x0, self.y0, self.x0, self.y1, self.x1, self.y1, self.x1, self.y0 + ) + + +class Room(Zone): + def __init__( + self, + number: int, + x0: Optional[float], + y0: Optional[float], + x1: Optional[float], + y1: Optional[float], + name: str = None, + pos_x: float = None, + pos_y: float = None, + ): + super().__init__(x0, y0, x1, y1) + self.number = number + self.name = name + self.pos_x = pos_x + self.pos_y = pos_y + + def as_dict(self) -> Dict[str, Any]: + super_dict = {**super(Room, self).as_dict()} + if self.name is not None: + super_dict[ATTR_NAME] = self.name + if self.pos_x is not None: + super_dict[ATTR_X] = self.pos_x + if self.pos_y is not None: + super_dict[ATTR_Y] = self.pos_y + return super_dict + + def __str__(self) -> str: + return f"[number: {self.number}, name: {self.name}, {self.x0}, {self.y0}, {self.x1}, {self.y1}]" + + def __repr__(self) -> str: + return self.__str__() + + def point(self) -> Optional[Point]: + if self.pos_x is not None and self.pos_y is not None and self.name is not None: + return Point(self.pos_x, self.pos_y) + return None + + +class Wall: + def __init__(self, x0: float, y0: float, x1: float, y1: float): + self.x0 = x0 + self.y0 = y0 + self.x1 = x1 + self.y1 = y1 + + def __str__(self) -> str: + return f"[{self.x0}, {self.y0}, {self.x1}, {self.y1}]" + + def __repr__(self) -> str: + return self.__str__() + + def as_dict(self) -> Dict[str, Any]: + return {ATTR_X0: self.x0, ATTR_Y0: self.y0, ATTR_X1: self.x1, ATTR_Y1: self.y1} + + def to_img(self, image_dimensions) -> Wall: + p0 = Point(self.x0, self.y0).to_img(image_dimensions) + p1 = Point(self.x1, self.y1).to_img(image_dimensions) + return Wall(p0.x, p0.y, p1.x, p1.y) + + def as_list(self) -> List[float]: + return [self.x0, self.y0, self.x1, self.y1] + + +class Area: + def __init__( + self, + x0: float, + y0: float, + x1: float, + y1: float, + x2: float, + y2: float, + x3: float, + y3: float, + ): + self.x0 = x0 + self.y0 = y0 + self.x1 = x1 + self.y1 = y1 + self.x2 = x2 + self.y2 = y2 + self.x3 = x3 + self.y3 = y3 + + def __str__(self) -> str: + return f"[{self.x0}, {self.y0}, {self.x1}, {self.y1}, {self.x2}, {self.y2}, {self.x3}, {self.y3}]" + + def __repr__(self) -> str: + return self.__str__() + + def as_dict(self) -> Dict[str, Any]: + return { + ATTR_X0: self.x0, + ATTR_Y0: self.y0, + ATTR_X1: self.x1, + ATTR_Y1: self.y1, + ATTR_X2: self.x2, + ATTR_Y2: self.y2, + ATTR_X3: self.x3, + ATTR_Y3: self.y3, + } + + def as_list(self) -> List[float]: + return [self.x0, self.y0, self.x1, self.y1, self.x2, self.y2, self.x3, self.y3] + + def to_img(self, image_dimensions) -> Area: + p0 = Point(self.x0, self.y0).to_img(image_dimensions) + p1 = Point(self.x1, self.y1).to_img(image_dimensions) + p2 = Point(self.x2, self.y2).to_img(image_dimensions) + p3 = Point(self.x3, self.y3).to_img(image_dimensions) + return Area(p0.x, p0.y, p1.x, p1.y, p2.x, p2.y, p3.x, p3.y) + + +class MapData: + def __init__(self, calibration_center: float = 0, calibration_diff: float = 0): + self._calibration_center = calibration_center + self._calibration_diff = calibration_diff + self.blocks = None + self.charger: Optional[Point] = None + self.goto: Optional[List[Point]] = None + self.goto_path: Optional[Path] = None + self.image: Optional[ImageData] = None + self.no_go_areas: Optional[List[Area]] = None + self.no_mopping_areas: Optional[List[Area]] = None + self.no_carpet_areas: Optional[List[Area]] = None + self.carpet_map: Optional[Set[int]] = [] + self.obstacles: Optional[List[Obstacle]] = None + self.ignored_obstacles: Optional[List[Obstacle]] = None + self.obstacles_with_photo: Optional[List[Obstacle]] = None + self.ignored_obstacles_with_photo: Optional[List[Obstacle]] = None + self.path: Optional[Path] = None + self.predicted_path: Optional[Path] = None + self.mop_path: Optional[Path] = None + self.rooms: Optional[Dict[int, Room]] = None + self.vacuum_position: Optional[Point] = None + self.vacuum_room: Optional[int] = None + self.vacuum_room_name: Optional[str] = None + self.walls: Optional[List[Wall]] = None + self.zones: Optional[List[Zone]] = None + self.cleaned_rooms: Optional[Set[int]] = None + self.map_name: Optional[str] = None + + def calibration(self) -> Optional[CalibrationPoints]: + if self.image.is_empty: + return None + calibration_points = [] + for point in [ + Point(self._calibration_center, self._calibration_center), + Point( + self._calibration_center + self._calibration_diff, + self._calibration_center, + ), + Point( + self._calibration_center, + self._calibration_center + self._calibration_diff, + ), + ]: + img_point = point.to_img(self.image.dimensions).rotated( + self.image.dimensions + ) + calibration_points.append( + { + "vacuum": {"x": point.x, "y": point.y}, + "map": {"x": int(img_point.x), "y": int(img_point.y)}, + } + ) + return calibration_points diff --git a/roborock/map_parser/map_data_parser.py b/roborock/map_parser/map_data_parser.py new file mode 100644 index 00000000..4ac61f99 --- /dev/null +++ b/roborock/map_parser/map_data_parser.py @@ -0,0 +1,539 @@ +import logging +from typing import Tuple + +from .image_handler import ImageHandlerRoborock +from .map_data import * +from .map_parser_const import * +from .types import Colors, Drawables, Sizes, Texts + +_LOGGER = logging.getLogger(__name__) + + +class MapDataParserRoborock: + CHARGER = 1 + IMAGE = 2 + PATH = 3 + GOTO_PATH = 4 + GOTO_PREDICTED_PATH = 5 + CURRENTLY_CLEANED_ZONES = 6 + GOTO_TARGET = 7 + ROBOT_POSITION = 8 + NO_GO_AREAS = 9 + VIRTUAL_WALLS = 10 + BLOCKS = 11 + NO_MOPPING_AREAS = 12 + OBSTACLES = 13 + IGNORED_OBSTACLES = 14 + OBSTACLES_WITH_PHOTO = 15 + IGNORED_OBSTACLES_WITH_PHOTO = 16 + CARPET_MAP = 17 + MOP_PATH = 18 + NO_CARPET_AREAS = 19 + DIGEST = 1024 + SIZE = 1024 + KNOWN_OBSTACLE_TYPES = { + 0: "cable", + 2: "shoes", + 3: "poop", + 5: "extension cord", + 9: "weighting scale", + 10: "clothes", + } + + @staticmethod + def create_empty(colors: Colors, text: str) -> MapData: + map_data = MapData() + empty_map = ImageHandlerRoborock.create_empty_map_image(colors, text) + map_data.image = ImageData.create_empty(empty_map) + return map_data + + @staticmethod + def draw_elements( + colors: Colors, + drawables: Drawables, + sizes: Sizes, + map_data: MapData, + image_config: ImageConfig, + ): + scale = float(image_config[CONF_SCALE]) + for drawable in drawables: + if DRAWABLE_CHARGER == drawable and map_data.charger is not None: + ImageHandlerRoborock.draw_charger( + map_data.image, map_data.charger, sizes, colors + ) + if ( + DRAWABLE_VACUUM_POSITION == drawable + and map_data.vacuum_position is not None + ): + ImageHandlerRoborock.draw_vacuum_position( + map_data.image, map_data.vacuum_position, sizes, colors + ) + if DRAWABLE_OBSTACLES == drawable and map_data.obstacles is not None: + ImageHandlerRoborock.draw_obstacles( + map_data.image, map_data.obstacles, sizes, colors + ) + if ( + DRAWABLE_IGNORED_OBSTACLES == drawable + and map_data.ignored_obstacles is not None + ): + ImageHandlerRoborock.draw_ignored_obstacles( + map_data.image, map_data.ignored_obstacles, sizes, colors + ) + if ( + DRAWABLE_OBSTACLES_WITH_PHOTO == drawable + and map_data.obstacles_with_photo is not None + ): + ImageHandlerRoborock.draw_obstacles_with_photo( + map_data.image, map_data.obstacles_with_photo, sizes, colors + ) + if ( + DRAWABLE_IGNORED_OBSTACLES_WITH_PHOTO == drawable + and map_data.ignored_obstacles_with_photo is not None + ): + ImageHandlerRoborock.draw_ignored_obstacles_with_photo( + map_data.image, map_data.ignored_obstacles_with_photo, sizes, colors + ) + if DRAWABLE_MOP_PATH == drawable and map_data.mop_path is not None: + ImageHandlerRoborock.draw_mop_path( + map_data.image, map_data.mop_path, sizes, colors, scale + ) + if DRAWABLE_PATH == drawable and map_data.path is not None: + ImageHandlerRoborock.draw_path( + map_data.image, map_data.path, sizes, colors, scale + ) + if DRAWABLE_GOTO_PATH == drawable and map_data.goto_path is not None: + ImageHandlerRoborock.draw_goto_path( + map_data.image, map_data.goto_path, sizes, colors, scale + ) + if ( + DRAWABLE_PREDICTED_PATH == drawable + and map_data.predicted_path is not None + ): + ImageHandlerRoborock.draw_predicted_path( + map_data.image, map_data.predicted_path, sizes, colors, scale + ) + if ( + DRAWABLE_NO_CARPET_AREAS == drawable + and map_data.no_carpet_areas is not None + ): + ImageHandlerRoborock.draw_no_carpet_areas( + map_data.image, map_data.no_carpet_areas, colors + ) + if DRAWABLE_NO_GO_AREAS == drawable and map_data.no_go_areas is not None: + ImageHandlerRoborock.draw_no_go_areas( + map_data.image, map_data.no_go_areas, colors + ) + if ( + DRAWABLE_NO_MOPPING_AREAS == drawable + and map_data.no_mopping_areas is not None + ): + ImageHandlerRoborock.draw_no_mopping_areas( + map_data.image, map_data.no_mopping_areas, colors + ) + if DRAWABLE_VIRTUAL_WALLS == drawable and map_data.walls is not None: + ImageHandlerRoborock.draw_walls(map_data.image, map_data.walls, colors) + if DRAWABLE_ZONES == drawable and map_data.zones is not None: + ImageHandlerRoborock.draw_zones(map_data.image, map_data.zones, colors) + if ( + DRAWABLE_CLEANED_AREA == drawable + and DRAWABLE_CLEANED_AREA in map_data.image.additional_layers + ): + ImageHandlerRoborock.draw_layer(map_data.image, drawable) + if DRAWABLE_ROOM_NAMES == drawable and map_data.rooms is not None: + ImageHandlerRoborock.draw_room_names( + map_data.image, map_data.rooms, colors + ) + + @staticmethod + def parse( + raw: bytes, + colors: Colors, + drawables: Drawables, + texts: Texts, + sizes: Sizes, + image_config: ImageConfig, + *args, + **kwargs, + ) -> MapData: + map_data = MapData(25500, 1000) + map_header_length = MapDataParserRoborock.get_int16(raw, 0x02) + map_data.major_version = MapDataParserRoborock.get_int16(raw, 0x08) + map_data.minor_version = MapDataParserRoborock.get_int16(raw, 0x0A) + map_data.map_index = MapDataParserRoborock.get_int32(raw, 0x0C) + map_data.map_sequence = MapDataParserRoborock.get_int32(raw, 0x10) + block_start_position = map_header_length + img_start = None + img_data = None + while block_start_position < len(raw): + block_header_length = MapDataParserRoborock.get_int16( + raw, block_start_position + 0x02 + ) + header = MapDataParserRoborock.get_bytes( + raw, block_start_position, block_header_length + ) + block_type = MapDataParserRoborock.get_int16(header, 0x00) + block_data_length = MapDataParserRoborock.get_int32(header, 0x04) + block_data_start = block_start_position + block_header_length + data = MapDataParserRoborock.get_bytes( + raw, block_data_start, block_data_length + ) + + if block_type == MapDataParserRoborock.CHARGER: + map_data.charger = MapDataParserRoborock.parse_object_position( + block_data_length, data + ) + elif block_type == MapDataParserRoborock.IMAGE: + img_start = block_start_position + img_data_length = block_data_length + img_header_length = block_header_length + img_data = data + img_header = header + elif block_type == MapDataParserRoborock.ROBOT_POSITION: + map_data.vacuum_position = MapDataParserRoborock.parse_object_position( + block_data_length, data + ) + elif block_type == MapDataParserRoborock.PATH: + map_data.path = MapDataParserRoborock.parse_path( + block_start_position, header, raw + ) + elif block_type == MapDataParserRoborock.GOTO_PATH: + map_data.goto_path = MapDataParserRoborock.parse_path( + block_start_position, header, raw + ) + elif block_type == MapDataParserRoborock.GOTO_PREDICTED_PATH: + map_data.predicted_path = MapDataParserRoborock.parse_path( + block_start_position, header, raw + ) + elif block_type == MapDataParserRoborock.CURRENTLY_CLEANED_ZONES: + map_data.zones = MapDataParserRoborock.parse_zones(data, header) + elif block_type == MapDataParserRoborock.GOTO_TARGET: + map_data.goto = MapDataParserRoborock.parse_goto_target(data) + elif block_type == MapDataParserRoborock.DIGEST: + map_data.is_valid = True + elif block_type == MapDataParserRoborock.VIRTUAL_WALLS: + map_data.walls = MapDataParserRoborock.parse_walls(data, header) + elif ( + block_type == MapDataParserRoborock.NO_GO_AREAS + and image_config[CONF_INCLUDE_NOGO] + ): + map_data.no_go_areas = MapDataParserRoborock.parse_area(header, data) + elif block_type == MapDataParserRoborock.NO_MOPPING_AREAS: + map_data.no_mopping_areas = MapDataParserRoborock.parse_area( + header, data + ) + elif block_type == MapDataParserRoborock.OBSTACLES: + map_data.obstacles = MapDataParserRoborock.parse_obstacles(data, header) + elif ( + block_type == MapDataParserRoborock.IGNORED_OBSTACLES + and image_config[CONF_INCLUDE_IGNORED_OBSTACLES] + ): + map_data.ignored_obstacles = MapDataParserRoborock.parse_obstacles( + data, header + ) + elif block_type == MapDataParserRoborock.OBSTACLES_WITH_PHOTO: + map_data.obstacles_with_photo = MapDataParserRoborock.parse_obstacles( + data, header + ) + elif block_type == MapDataParserRoborock.IGNORED_OBSTACLES_WITH_PHOTO: + map_data.ignored_obstacles_with_photo = ( + MapDataParserRoborock.parse_obstacles(data, header) + ) + elif block_type == MapDataParserRoborock.BLOCKS: + block_pairs = MapDataParserRoborock.get_int16(header, 0x08) + map_data.blocks = MapDataParserRoborock.get_bytes(data, 0, block_pairs) + elif block_type == MapDataParserRoborock.MOP_PATH: + points_mask = MapDataParserRoborock.get_bytes( + raw, block_data_start, block_data_length + ) + # only the map_data.path points where points_mask == 1 are in mop_path + map_data.mop_path = MapDataParserRoborock.parse_mop_path( + map_data.path, points_mask + ) + elif block_type == MapDataParserRoborock.CARPET_MAP: + data = MapDataParserRoborock.get_bytes( + raw, block_data_start, block_data_length + ) + # only the indexes where value == 1 are in carpet_map + map_data.carpet_map = MapDataParserRoborock.parse_carpet_map( + data, image_config + ) + elif block_type == MapDataParserRoborock.NO_CARPET_AREAS: + map_data.no_carpet_areas = MapDataParserRoborock.parse_area( + header, data + ) + else: + _LOGGER.debug( + "UNKNOWN BLOCK TYPE: %s, header length %s, data length %s", + block_type, + block_header_length, + block_data_length, + ) + block_start_position = ( + block_start_position + + block_data_length + + MapDataParserRoborock.get_int8(header, 2) + ) + + if img_data: + image, rooms = MapDataParserRoborock.parse_image( + img_data_length, + img_header_length, + img_data, + img_header, + map_data.carpet_map, + colors, + image_config, + ) + map_data.image = image + map_data.rooms = rooms + + if not map_data.image.is_empty: + MapDataParserRoborock.draw_elements( + colors, drawables, sizes, map_data, image_config + ) + if len(map_data.rooms) > 0 and map_data.vacuum_position is not None: + map_data.vacuum_room = MapDataParserRoborock.get_current_vacuum_room( + img_start, raw, map_data.vacuum_position + ) + ImageHandlerRoborock.rotate(map_data.image) + ImageHandlerRoborock.draw_texts(map_data.image, texts) + return map_data + + @staticmethod + def map_to_image(p: Point) -> Point: + return Point(p.x / MM, p.y / MM) + + @staticmethod + def image_to_map(x: float) -> float: + return x * MM + + @staticmethod + def get_current_vacuum_room( + block_start_position: int, raw: bytes, vacuum_position: Point + ) -> int: + block_header_length = MapDataParserRoborock.get_int16( + raw, block_start_position + 0x02 + ) + header = MapDataParserRoborock.get_bytes( + raw, block_start_position, block_header_length + ) + block_data_length = MapDataParserRoborock.get_int32(header, 0x04) + block_data_start = block_start_position + block_header_length + data = MapDataParserRoborock.get_bytes(raw, block_data_start, block_data_length) + image_top = MapDataParserRoborock.get_int32(header, block_header_length - 16) + image_left = MapDataParserRoborock.get_int32(header, block_header_length - 12) + image_width = MapDataParserRoborock.get_int32(header, block_header_length - 4) + p = MapDataParserRoborock.map_to_image(vacuum_position) + room = ImageHandlerRoborock.get_room_at_pixel( + data, image_width, round(p.x - image_left), round(p.y - image_top) + ) + return room + + @staticmethod + def parse_image( + block_data_length: int, + block_header_length: int, + data: bytes, + header: bytes, + carpet_map: Set[int], + colors: Colors, + image_config: ImageConfig, + ) -> Tuple[ImageData, Dict[int, Room]]: + image_size = block_data_length + image_top = MapDataParserRoborock.get_int32(header, block_header_length - 16) + image_left = MapDataParserRoborock.get_int32(header, block_header_length - 12) + image_height = MapDataParserRoborock.get_int32(header, block_header_length - 8) + image_width = MapDataParserRoborock.get_int32(header, block_header_length - 4) + if ( + image_width + - image_width + * (image_config[CONF_TRIM][CONF_LEFT] + image_config[CONF_TRIM][CONF_RIGHT]) + / 100 + < MINIMAL_IMAGE_WIDTH + ): + image_config[CONF_TRIM][CONF_LEFT] = 0 + image_config[CONF_TRIM][CONF_RIGHT] = 0 + if ( + image_height + - image_height + * (image_config[CONF_TRIM][CONF_TOP] + image_config[CONF_TRIM][CONF_BOTTOM]) + / 100 + < MINIMAL_IMAGE_HEIGHT + ): + image_config[CONF_TRIM][CONF_TOP] = 0 + image_config[CONF_TRIM][CONF_BOTTOM] = 0 + image, rooms_raw = ImageHandlerRoborock.parse( + data, image_width, image_height, carpet_map, colors, image_config + ) + rooms = {} + for number, room in rooms_raw.items(): + rooms[number] = Room( + number, + MapDataParserRoborock.image_to_map(room[0] + image_left), + MapDataParserRoborock.image_to_map(room[1] + image_top), + MapDataParserRoborock.image_to_map(room[2] + image_left), + MapDataParserRoborock.image_to_map(room[3] + image_top), + ) + return ( + ImageData( + image_size, + image_top, + image_left, + image_height, + image_width, + image_config, + image, + MapDataParserRoborock.map_to_image, + ), + rooms, + ) + + @staticmethod + def parse_carpet_map(data: bytes, image_config: ImageConfig) -> Set[int]: + carpet_map = set() + + for i, v in enumerate(data): + if v: + carpet_map.add(i) + return carpet_map + + @staticmethod + def parse_goto_target(data: bytes) -> Point: + x = MapDataParserRoborock.get_int16(data, 0x00) + y = MapDataParserRoborock.get_int16(data, 0x02) + return Point(x, y) + + @staticmethod + def parse_object_position(block_data_length: int, data: bytes) -> Point: + x = MapDataParserRoborock.get_int32(data, 0x00) + y = MapDataParserRoborock.get_int32(data, 0x04) + a = None + if block_data_length > 8: + a = MapDataParserRoborock.get_int32(data, 0x08) + if a > 0xFF: + a = (a & 0xFF) - 256 + return Point(x, y, a) + + @staticmethod + def parse_walls(data: bytes, header: bytes) -> List[Wall]: + wall_pairs = MapDataParserRoborock.get_int16(header, 0x08) + walls = [] + for wall_start in range(0, wall_pairs * 8, 8): + x0 = MapDataParserRoborock.get_int16(data, wall_start + 0) + y0 = MapDataParserRoborock.get_int16(data, wall_start + 2) + x1 = MapDataParserRoborock.get_int16(data, wall_start + 4) + y1 = MapDataParserRoborock.get_int16(data, wall_start + 6) + walls.append(Wall(x0, y0, x1, y1)) + return walls + + @staticmethod + def parse_obstacles(data: bytes, header: bytes) -> List[Obstacle]: + obstacle_pairs = MapDataParserRoborock.get_int16(header, 0x08) + obstacles = [] + if obstacle_pairs == 0: + return obstacles + obstacle_size = int(len(data) / obstacle_pairs) + for obstacle_start in range(0, obstacle_pairs * obstacle_size, obstacle_size): + x = MapDataParserRoborock.get_int16(data, obstacle_start + 0) + y = MapDataParserRoborock.get_int16(data, obstacle_start + 2) + details = {} + if obstacle_size >= 6: + details[ATTR_TYPE] = MapDataParserRoborock.get_int16( + data, obstacle_start + 4 + ) + if details[ATTR_TYPE] in MapDataParserRoborock.KNOWN_OBSTACLE_TYPES: + details[ + ATTR_DESCRIPTION + ] = MapDataParserRoborock.KNOWN_OBSTACLE_TYPES[details[ATTR_TYPE]] + if obstacle_size >= 10: + u1 = MapDataParserRoborock.get_int16(data, obstacle_start + 6) + u2 = MapDataParserRoborock.get_int16(data, obstacle_start + 8) + details[ATTR_CONFIDENCE_LEVEL] = 0 if u2 == 0 else u1 * 10.0 / u2 + if obstacle_size == 28 and (data[obstacle_start + 12] & 0xFF) > 0: + txt = MapDataParserRoborock.get_bytes( + data, obstacle_start + 12, 16 + ) + details[ATTR_PHOTO_NAME] = txt.decode("ascii") + obstacles.append(Obstacle(x, y, details)) + return obstacles + + @staticmethod + def parse_zones(data: bytes, header: bytes) -> List[Zone]: + zone_pairs = MapDataParserRoborock.get_int16(header, 0x08) + zones = [] + for zone_start in range(0, zone_pairs * 8, 8): + x0 = MapDataParserRoborock.get_int16(data, zone_start + 0) + y0 = MapDataParserRoborock.get_int16(data, zone_start + 2) + x1 = MapDataParserRoborock.get_int16(data, zone_start + 4) + y1 = MapDataParserRoborock.get_int16(data, zone_start + 6) + zones.append(Zone(x0, y0, x1, y1)) + return zones + + @staticmethod + def parse_path(block_start_position: int, header: bytes, raw: bytes) -> Path: + path_points = [] + end_pos = MapDataParserRoborock.get_int32(header, 0x04) + point_length = MapDataParserRoborock.get_int32(header, 0x08) + point_size = MapDataParserRoborock.get_int32(header, 0x0C) + angle = MapDataParserRoborock.get_int32(header, 0x10) + start_pos = block_start_position + 0x14 + for pos in range(start_pos, start_pos + end_pos, 4): + x = MapDataParserRoborock.get_int16(raw, pos) + y = MapDataParserRoborock.get_int16(raw, pos + 2) + path_points.append(Point(x, y)) + return Path(point_length, point_size, angle, [path_points]) + + @staticmethod + def parse_mop_path(path: Path, mask: bytes) -> Path: + mop_paths = [] + points_num = 0 + for each_path in path.path: + mop_path_points = [] + for i, point in enumerate(each_path): + if mask[i]: + mop_path_points.append(point) + if i + 1 < len(mask) and not mask[i + 1]: + points_num += len(mop_path_points) + mop_paths.append(mop_path_points) + mop_path_points = [] + + points_num += len(mop_path_points) + mop_paths.append(mop_path_points) + return Path(points_num, path.point_size, path.angle, mop_paths) + + @staticmethod + def parse_area(header: bytes, data: bytes) -> List[Area]: + area_pairs = MapDataParserRoborock.get_int16(header, 0x08) + areas = [] + for area_start in range(0, area_pairs * 16, 16): + x0 = MapDataParserRoborock.get_int16(data, area_start + 0) + y0 = MapDataParserRoborock.get_int16(data, area_start + 2) + x1 = MapDataParserRoborock.get_int16(data, area_start + 4) + y1 = MapDataParserRoborock.get_int16(data, area_start + 6) + x2 = MapDataParserRoborock.get_int16(data, area_start + 8) + y2 = MapDataParserRoborock.get_int16(data, area_start + 10) + x3 = MapDataParserRoborock.get_int16(data, area_start + 12) + y3 = MapDataParserRoborock.get_int16(data, area_start + 14) + areas.append(Area(x0, y0, x1, y1, x2, y2, x3, y3)) + return areas + + @staticmethod + def get_bytes(data: bytes, start_index: int, size: int) -> bytes: + return data[start_index: start_index + size] + + @staticmethod + def get_int8(data: bytes, address: int) -> int: + return data[address] & 0xFF + + @staticmethod + def get_int16(data: bytes, address: int) -> int: + return ((data[address + 0] << 0) & 0xFF) | ((data[address + 1] << 8) & 0xFFFF) + + @staticmethod + def get_int32(data: bytes, address: int) -> int: + return ( + ((data[address + 0] << 0) & 0xFF) + | ((data[address + 1] << 8) & 0xFFFF) + | ((data[address + 2] << 16) & 0xFFFFFF) + | ((data[address + 3] << 24) & 0xFFFFFFFF) + ) diff --git a/roborock/map_parser/map_parser_const.py b/roborock/map_parser/map_parser_const.py new file mode 100644 index 00000000..62c20d2e --- /dev/null +++ b/roborock/map_parser/map_parser_const.py @@ -0,0 +1,276 @@ +"""Constants for Roborock Map creation.""" +MINIMAL_IMAGE_WIDTH = 20 +MINIMAL_IMAGE_HEIGHT = 20 + +CONF_INCLUDE_SHARED = "include_shared" +CONF_INCLUDE_NOGO = "include_nogo" +CONF_INCLUDE_IGNORED_OBSTACLES = "include_ignored_obstacles" +CONF_BOTTOM = "bottom" +CONF_COLOR = "color" +CONF_COLORS = "colors" +CONF_COUNTRY = "country" +CONF_DRAW = "draw" +CONF_FORCE_API = "force_api" +CONF_FONT = "font" +CONF_FONT_SIZE = "font_size" +CONF_LEFT = "left" +CONF_MAP_TRANSFORM = "map_transformation" +CONF_RIGHT = "right" +CONF_ROOM_COLORS = "room_colors" +CONF_ROTATE = "rotate" +CONF_SCALE = "scale" +CONF_SIZES = "sizes" +CONF_SIZE_CHARGER_RADIUS = "charger_radius" +CONF_SIZE_IGNORED_OBSTACLE_RADIUS = "ignored_obstacle_radius" +CONF_SIZE_IGNORED_OBSTACLE_WITH_PHOTO_RADIUS = "ignored_obstacle_with_photo_radius" +CONF_SIZE_MOP_PATH_WIDTH = "mop_path_width" +CONF_SIZE_OBSTACLE_RADIUS = "obstacle_radius" +CONF_SIZE_OBSTACLE_WITH_PHOTO_RADIUS = "obstacle_with_photo_radius" +CONF_SIZE_VACUUM_RADIUS = "vacuum_radius" +CONF_SIZE_PATH_WIDTH = "path_width" +CONF_STORE_MAP_RAW = "store_map_raw" +CONF_STORE_MAP_IMAGE = "store_map_image" +CONF_STORE_MAP_PATH = "store_map_path" +CONF_TEXT = "text" +CONF_TEXTS = "texts" +CONF_TOP = "top" +CONF_TRIM = "trim" +CONF_X = "x" +CONF_Y = "y" +CONTENT_TYPE = "image/png" + +DRAWABLE_CHARGER = "charger" +DRAWABLE_CLEANED_AREA = "cleaned_area" +DRAWABLE_GOTO_PATH = "goto_path" +DRAWABLE_IGNORED_OBSTACLES = "ignored_obstacles" +DRAWABLE_IGNORED_OBSTACLES_WITH_PHOTO = "ignored_obstacles_with_photo" +DRAWABLE_MOP_PATH = "mop_path" +DRAWABLE_NO_CARPET_AREAS = "no_carpet_zones" +DRAWABLE_NO_GO_AREAS = "no_go_zones" +DRAWABLE_NO_MOPPING_AREAS = "no_mopping_zones" +DRAWABLE_OBSTACLES = "obstacles" +DRAWABLE_OBSTACLES_WITH_PHOTO = "obstacles_with_photo" +DRAWABLE_PATH = "path" +DRAWABLE_PREDICTED_PATH = "predicted_path" +DRAWABLE_ROOM_NAMES = "room_names" +DRAWABLE_VACUUM_POSITION = "vacuum_position" +DRAWABLE_VIRTUAL_WALLS = "virtual_walls" +DRAWABLE_ZONES = "zones" +CONF_AVAILABLE_DRAWABLES = [ + DRAWABLE_CLEANED_AREA, + DRAWABLE_CHARGER, + DRAWABLE_GOTO_PATH, + DRAWABLE_IGNORED_OBSTACLES, + DRAWABLE_IGNORED_OBSTACLES_WITH_PHOTO, + DRAWABLE_MOP_PATH, + DRAWABLE_NO_CARPET_AREAS, + DRAWABLE_NO_GO_AREAS, + DRAWABLE_NO_MOPPING_AREAS, + DRAWABLE_PATH, + DRAWABLE_OBSTACLES, + DRAWABLE_OBSTACLES_WITH_PHOTO, + DRAWABLE_PREDICTED_PATH, + DRAWABLE_ROOM_NAMES, + DRAWABLE_VACUUM_POSITION, + DRAWABLE_VIRTUAL_WALLS, + DRAWABLE_ZONES, +] + +COLOR_ROOM_PREFIX = "color_room_" + +COLOR_CARPETS = "color_carpets" +COLOR_CHARGER = "color_charger" +COLOR_CHARGER_OUTLINE = "color_charger_outline" +COLOR_CLEANED_AREA = "color_cleaned_area" +COLOR_GOTO_PATH = "color_goto_path" +COLOR_GREY_WALL = "color_grey_wall" +COLOR_IGNORED_OBSTACLE = "color_ignored_obstacle" +COLOR_IGNORED_OBSTACLE_WITH_PHOTO = "color_ignored_obstacle_with_photo" +COLOR_MAP_INSIDE = "color_map_inside" +COLOR_MAP_OUTSIDE = "color_map_outside" +COLOR_MAP_WALL = "color_map_wall" +COLOR_MAP_WALL_V2 = "color_map_wall_v2" +COLOR_MOP_PATH = "color_mop_path" +COLOR_NEW_DISCOVERED_AREA = "color_new_discovered_area" +COLOR_NO_CARPET_ZONES = "color_no_carpet_zones" +COLOR_NO_CARPET_ZONES_OUTLINE = "color_no_carpet_zones_outline" +COLOR_NO_GO_ZONES = "color_no_go_zones" +COLOR_NO_GO_ZONES_OUTLINE = "color_no_go_zones_outline" +COLOR_NO_MOPPING_ZONES = "color_no_mop_zones" +COLOR_NO_MOPPING_ZONES_OUTLINE = "color_no_mop_zones_outline" +COLOR_OBSTACLE = "color_obstacle" +COLOR_OBSTACLE_WITH_PHOTO = "color_obstacle_with_photo" +COLOR_OBSTACLE_OUTLINE = "color_obstacle_outline" +COLOR_PATH = "color_path" +COLOR_PREDICTED_PATH = "color_predicted_path" +COLOR_ROBO = "color_robo" +COLOR_ROBO_OUTLINE = "color_robo_outline" +COLOR_ROOM_NAMES = "color_room_names" +COLOR_SCAN = "color_scan" +COLOR_UNKNOWN = "color_unknown" +COLOR_VIRTUAL_WALLS = "color_virtual_walls" +COLOR_ZONES = "color_zones" +COLOR_ZONES_OUTLINE = "color_zones_outline" + +CONF_AVAILABLE_COLORS = [ + COLOR_CARPETS, + COLOR_CHARGER, + COLOR_CHARGER_OUTLINE, + COLOR_CLEANED_AREA, + COLOR_GOTO_PATH, + COLOR_GREY_WALL, + COLOR_IGNORED_OBSTACLE, + COLOR_IGNORED_OBSTACLE_WITH_PHOTO, + COLOR_MAP_INSIDE, + COLOR_MAP_OUTSIDE, + COLOR_MAP_WALL, + COLOR_MAP_WALL_V2, + COLOR_MOP_PATH, + COLOR_NEW_DISCOVERED_AREA, + COLOR_NO_CARPET_ZONES, + COLOR_NO_CARPET_ZONES_OUTLINE, + COLOR_NO_GO_ZONES, + COLOR_NO_GO_ZONES_OUTLINE, + COLOR_NO_MOPPING_ZONES, + COLOR_NO_MOPPING_ZONES_OUTLINE, + COLOR_OBSTACLE, + COLOR_OBSTACLE_WITH_PHOTO, + COLOR_PATH, + COLOR_PREDICTED_PATH, + COLOR_ROBO, + COLOR_ROBO_OUTLINE, + COLOR_ROOM_NAMES, + COLOR_SCAN, + COLOR_UNKNOWN, + COLOR_VIRTUAL_WALLS, + COLOR_ZONES, + COLOR_ZONES_OUTLINE, + COLOR_OBSTACLE_OUTLINE, +] + +COLOR_ROOM_1 = "color_room_1" +COLOR_ROOM_2 = "color_room_2" +COLOR_ROOM_3 = "color_room_3" +COLOR_ROOM_4 = "color_room_4" +COLOR_ROOM_5 = "color_room_5" +COLOR_ROOM_6 = "color_room_6" +COLOR_ROOM_7 = "color_room_7" +COLOR_ROOM_8 = "color_room_8" +COLOR_ROOM_9 = "color_room_9" +COLOR_ROOM_10 = "color_room_10" +COLOR_ROOM_11 = "color_room_11" +COLOR_ROOM_12 = "color_room_12" +COLOR_ROOM_13 = "color_room_13" +COLOR_ROOM_14 = "color_room_14" +COLOR_ROOM_15 = "color_room_15" +COLOR_ROOM_16 = "color_room_16" + +CONF_DEFAULT_ROOM_COLORS = [ + COLOR_ROOM_1, + COLOR_ROOM_2, + COLOR_ROOM_3, + COLOR_ROOM_4, + COLOR_ROOM_5, + COLOR_ROOM_6, + COLOR_ROOM_7, + COLOR_ROOM_8, + COLOR_ROOM_9, + COLOR_ROOM_10, + COLOR_ROOM_11, + COLOR_ROOM_12, + COLOR_ROOM_13, + COLOR_ROOM_14, + COLOR_ROOM_15, + COLOR_ROOM_16, +] + +ATTRIBUTE_CALIBRATION = "calibration_points" +ATTRIBUTE_CARPET_MAP = "carpet_map" +ATTRIBUTE_CHARGER = "charger" +ATTRIBUTE_CLEANED_ROOMS = "cleaned_rooms" +ATTRIBUTE_COUNTRY = "country" +ATTRIBUTE_GOTO = "goto" +ATTRIBUTE_GOTO_PATH = "goto_path" +ATTRIBUTE_GOTO_PREDICTED_PATH = "goto_predicted_path" +ATTRIBUTE_IGNORED_OBSTACLES = "ignored_obstacles" +ATTRIBUTE_IGNORED_OBSTACLES_WITH_PHOTO = "ignored_obstacles_with_photo" +ATTRIBUTE_IMAGE = "image" +ATTRIBUTE_IS_EMPTY = "is_empty" +ATTRIBUTE_MAP_NAME = "map_name" +ATTRIBUTE_MOP_PATH = "mop_path" +ATTRIBUTE_MAP_SAVED = "map_saved" +ATTRIBUTE_NO_CARPET_AREAS = "no_carpet_areas" +ATTRIBUTE_NO_GO_AREAS = "no_go_areas" +ATTRIBUTE_NO_MOPPING_AREAS = "no_mopping_areas" +ATTRIBUTE_OBSTACLES = "obstacles" +ATTRIBUTE_OBSTACLES_WITH_PHOTO = "obstacles_with_photo" +ATTRIBUTE_PATH = "path" +ATTRIBUTE_ROOMS = "rooms" +ATTRIBUTE_ROOM_NUMBERS = "room_numbers" +ATTRIBUTE_VACUUM_POSITION = "vacuum_position" +ATTRIBUTE_VACUUM_ROOM = "vacuum_room" +ATTRIBUTE_VACUUM_ROOM_NAME = "vacuum_room_name" +ATTRIBUTE_WALLS = "walls" +ATTRIBUTE_ZONES = "zones" +CONF_AVAILABLE_ATTRIBUTES = [ + ATTRIBUTE_CALIBRATION, + ATTRIBUTE_CARPET_MAP, + ATTRIBUTE_NO_CARPET_AREAS, + ATTRIBUTE_CHARGER, + ATTRIBUTE_CLEANED_ROOMS, + ATTRIBUTE_COUNTRY, + ATTRIBUTE_GOTO, + ATTRIBUTE_GOTO_PATH, + ATTRIBUTE_GOTO_PREDICTED_PATH, + ATTRIBUTE_IGNORED_OBSTACLES, + ATTRIBUTE_IGNORED_OBSTACLES_WITH_PHOTO, + ATTRIBUTE_IMAGE, + ATTRIBUTE_IS_EMPTY, + ATTRIBUTE_MAP_NAME, + ATTRIBUTE_MOP_PATH, + ATTRIBUTE_NO_GO_AREAS, + ATTRIBUTE_NO_MOPPING_AREAS, + ATTRIBUTE_OBSTACLES, + ATTRIBUTE_OBSTACLES_WITH_PHOTO, + ATTRIBUTE_PATH, + ATTRIBUTE_ROOMS, + ATTRIBUTE_ROOM_NUMBERS, + ATTRIBUTE_VACUUM_POSITION, + ATTRIBUTE_VACUUM_ROOM, + ATTRIBUTE_VACUUM_ROOM_NAME, + ATTRIBUTE_WALLS, + ATTRIBUTE_ZONES, +] + +ATTR_A = "a" +ATTR_ANGLE = "angle" +ATTR_CONFIDENCE_LEVEL = "confidence_level" +ATTR_DESCRIPTION = "description" +ATTR_HEIGHT = "height" +ATTR_MODEL = "model" +ATTR_NAME = "name" +ATTR_OFFSET_X = "offset_x" +ATTR_OFFSET_Y = "offset_y" +ATTR_PATH = "path" +ATTR_PHOTO_NAME = "photo_name" +ATTR_POINT_LENGTH = "point_length" +ATTR_POINT_SIZE = "point_size" +ATTR_ROTATION = "rotation" +ATTR_SCALE = "scale" +ATTR_SIZE = "size" +ATTR_TWO_FACTOR_AUTH = "url_2fa" +ATTR_TYPE = "type" +ATTR_USED_API = "used_api" +ATTR_WIDTH = "width" +ATTR_X = "x" +ATTR_X0 = "x0" +ATTR_X1 = "x1" +ATTR_X2 = "x2" +ATTR_X3 = "x3" +ATTR_Y = "y" +ATTR_Y0 = "y0" +ATTR_Y1 = "y1" +ATTR_Y2 = "y2" +ATTR_Y3 = "y3" + +MM = 50 diff --git a/roborock/map_parser/types.py b/roborock/map_parser/types.py new file mode 100644 index 00000000..6c00ec28 --- /dev/null +++ b/roborock/map_parser/types.py @@ -0,0 +1,9 @@ +from typing import Any, Dict, List, Tuple, Union + +Color = Union[Tuple[int, int, int], Tuple[int, int, int, int]] +Colors = Dict[str, Color] +Drawables = List[str] +Texts = List[Any] +Sizes = Dict[str, float] +ImageConfig = Dict[str, Any] +CalibrationPoints = List[Dict[str, Dict[str, Union[float, int]]]]