From 23856c8c3635b1e6b29e07bf82ae80298cffd26c Mon Sep 17 00:00:00 2001 From: Vincent <2070309+tubededentifrice@users.noreply.github.com> Date: Sat, 25 Jul 2026 17:08:01 +0400 Subject: [PATCH] feat: align Q10 map colors with V1 --- roborock/map/b01_q10_map_parser.py | 50 +++++++++++++++-------- roborock/map/map_parser.py | 61 ++++++++++++++++++++++++++-- roborock/map/room_colors.py | 59 +++++++++++++++++++++++++++ tests/map/test_b01_q10_map_parser.py | 31 ++++++++++++++ tests/map/test_map_parser.py | 42 ++++++++++++++++++- 5 files changed, 223 insertions(+), 20 deletions(-) create mode 100644 roborock/map/room_colors.py diff --git a/roborock/map/b01_q10_map_parser.py b/roborock/map/b01_q10_map_parser.py index 062f85f32..df710a740 100644 --- a/roborock/map/b01_q10_map_parser.py +++ b/roborock/map/b01_q10_map_parser.py @@ -19,13 +19,13 @@ https://github.com/v1b3c0d3x3r/roborock-qseries-map-bridge """ -import colorsys import io import math import statistics from dataclasses import dataclass, field, replace from PIL import Image +from vacuum_map_parser_base.config.color import ColorsPalette, SupportedColor from vacuum_map_parser_base.config.image_config import ImageConfig from vacuum_map_parser_base.map_data import ImageData, MapData @@ -40,6 +40,7 @@ decompose_grid, ) from .map_parser import ParsedMapData +from .room_colors import adjacency_aware_room_colors _MAP_FILE_FORMAT = "PNG" @@ -655,12 +656,12 @@ def parsed_from_packet(self, packet: Q10MapPacket) -> ParsedMapData: return ParsedMapData(image_content=image_bytes.getvalue(), map_data=map_data) def _render(self, packet: Q10MapPacket) -> Image.Image: - """Render the Q10 grid: rooms get distinct colors, walls white, rest dark.""" - palette = _build_palette(packet.grid) - rgb = bytearray() + """Render the Q10 grid with the V1 map palette.""" + palette = _build_palette(packet.grid, packet.width) + rgba = bytearray() for value in packet.grid: - rgb.extend(palette[value]) - img = Image.frombytes("RGB", (packet.width, packet.height), bytes(rgb)) + rgba.extend(palette[value]) + img = Image.frombytes("RGBA", (packet.width, packet.height), bytes(rgba)) # The ss07 grid is stored top-down (row 0 = top of the home), so it is # rendered as-is -- unlike the V1/Q7 convention, no vertical flip. scale = self._config.map_scale @@ -669,15 +670,32 @@ def _render(self, packet: Q10MapPacket) -> Image.Image: return img -def _build_palette(grid: bytes) -> list[tuple[int, int, int]]: - """Map each grid value to an RGB color (rooms distinct, walls white).""" - palette: list[tuple[int, int, int]] = [(28, 30, 38)] * 256 # default: unknown/outside - room_values = sorted({v for v in set(grid) if 0 < v < _WALL_THRESHOLD}) - for index, value in enumerate(room_values): - hue = (index * 0.139) % 1.0 - r, g, b = colorsys.hsv_to_rgb(hue, 0.5, 0.95) - palette[value] = (int(r * 255), int(g * 255), int(b * 255)) +def _opaque(color: tuple[int, ...]) -> tuple[int, int, int, int]: + """Return a palette color as RGBA.""" + return (color[0], color[1], color[2], color[3] if len(color) == 4 else 255) + + +def _build_palette(grid: bytes, width: int) -> list[tuple[int, int, int, int]]: + """Map Q10 cells onto the same colors used by the V1 map renderer.""" + + def room_id(value: int) -> int | None: + return max(1, value // 4) if 0 < value < _WALL_THRESHOLD else None + + colors = ColorsPalette() + room_colors = adjacency_aware_room_colors( + grid, + width, + colors, + room_id, + ) + outside = (0, 0, 0, 0) + palette = [outside] * 256 + for value in {value for value in grid if 0 < value < _WALL_THRESHOLD}: + palette[value] = _opaque(room_colors[max(1, value // 4)]) + wall = _opaque(colors.get_color(SupportedColor.GREY_WALL)) for value in range(_WALL_THRESHOLD, 256): - palette[value] = (235, 235, 240) # walls / borders - palette[0] = (28, 30, 38) + palette[value] = wall + palette[_UNSEGMENTED_FLOOR_VALUE] = _opaque(colors.get_color(SupportedColor.MAP_INSIDE)) + palette[_BACKGROUND_VALUE] = outside + palette[0] = outside return palette diff --git a/roborock/map/map_parser.py b/roborock/map/map_parser.py index 44924a8a3..59c8c1c3e 100644 --- a/roborock/map/map_parser.py +++ b/roborock/map/map_parser.py @@ -2,6 +2,7 @@ import io import logging +import threading from dataclasses import dataclass, field from vacuum_map_parser_base.config.color import ColorsPalette, SupportedColor @@ -9,10 +10,13 @@ from vacuum_map_parser_base.config.image_config import ImageConfig from vacuum_map_parser_base.config.size import Size, Sizes from vacuum_map_parser_base.map_data import MapData +from vacuum_map_parser_roborock.image_parser import RoborockImageParser from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser from roborock.exceptions import RoborockException +from .room_colors import adjacency_aware_room_colors + _LOGGER = logging.getLogger(__name__) DEFAULT_DRAWABLES = { @@ -98,6 +102,49 @@ def parse(self, map_bytes: bytes) -> ParsedMapData | None: return ParsedMapData(image_content=img_byte_arr.getvalue(), map_data=parsed_map) +class _AdjacencyAwareRoborockImageParser(RoborockImageParser): + """Apply the shared adjacency color policy to V1 room cells.""" + + def __init__( + self, + palette: ColorsPalette, + image_config: ImageConfig, + *, + recolor_rooms: bool = True, + ) -> None: + super().__init__(palette, image_config) + self._room_palette = palette + self._base_room_colors = palette.cached_room_colors.copy() + self._recolor_rooms = recolor_rooms + self._palette_lock = threading.Lock() + + def parse( + self, + raw_data: bytes, + width: int, + height: int, + carpet_map: set[int] | None, + removed_map: set[int] | None = None, + ): + """Assign non-conflicting room colors before the V1 image pass.""" + with self._palette_lock: + self._room_palette.cached_room_colors.clear() + self._room_palette.cached_room_colors.update(self._base_room_colors) + + if self._recolor_rooms: + + def room_id(value: int) -> int | None: + if value in (self.MAP_OUTSIDE, self.MAP_WALL, self.MAP_INSIDE, self.MAP_SCAN): + return None + return self._get_room_number(value) if value & 0x07 == 0x07 else None + + room_colors = adjacency_aware_room_colors(raw_data, width, self._room_palette, room_id) + for number, color in room_colors.items(): + self._room_palette.cached_room_colors[number] = color + self._room_palette.cached_room_colors[str(number)] = color + return super().parse(raw_data, width, height, carpet_map, removed_map) + + def _create_map_data_parser(config: MapParserConfig) -> RoborockMapDataParser: """Create a RoborockMapDataParser based on the config entry.""" color_dicts = {} @@ -114,10 +161,18 @@ def _create_map_data_parser(config: MapParserConfig) -> RoborockMapDataParser: if not config.show_rooms: room_colors = {str(x): (0, 0, 0, 0) for x in range(1, 32)} - return RoborockMapDataParser( - ColorsPalette(color_dicts, room_colors), + palette = ColorsPalette(color_dicts, room_colors) + image_config = ImageConfig(scale=config.map_scale) + parser = RoborockMapDataParser( + palette, Sizes({k: v * config.map_scale for k, v in Sizes.SIZES.items() if k != Size.MOP_PATH_WIDTH}), config.drawables, - ImageConfig(scale=config.map_scale), + image_config, [], ) + parser._image_parser = _AdjacencyAwareRoborockImageParser( + palette, + image_config, + recolor_rooms=config.show_rooms, + ) + return parser diff --git a/roborock/map/room_colors.py b/roborock/map/room_colors.py new file mode 100644 index 000000000..7fa9c4014 --- /dev/null +++ b/roborock/map/room_colors.py @@ -0,0 +1,59 @@ +"""Deterministic room colors that keep adjacent segments distinguishable.""" + +from collections.abc import Callable, Sequence + +from vacuum_map_parser_base.config.color import Color, ColorsPalette + +RoomIdFromCell = Callable[[int], int | None] + + +def adjacency_aware_room_colors( + grid: Sequence[int], + width: int, + palette: ColorsPalette, + room_id_from_cell: RoomIdFromCell, +) -> dict[int, Color]: + """Return room colors, changing only adjacent same-color conflicts. + + Room IDs remain the stable preference, matching the existing V1 palette. + When two rooms sharing an edge resolve to the same RGB value, the + higher-numbered room receives the first palette color not already used by + one of its colored neighbors. + """ + if width <= 0: + return {} + + room_ids: set[int] = set() + neighbors: dict[int, set[int]] = {} + for index, value in enumerate(grid): + room_id = room_id_from_cell(value) + if room_id is None: + continue + room_ids.add(room_id) + neighbors.setdefault(room_id, set()) + + for neighbor_index in (index - 1 if index % width else -1, index - width): + if neighbor_index < 0: + continue + neighbor_id = room_id_from_cell(grid[neighbor_index]) + if neighbor_id is None or neighbor_id == room_id: + continue + neighbors[room_id].add(neighbor_id) + neighbors.setdefault(neighbor_id, set()).add(room_id) + + candidates: list[Color] = [] + for palette_id in map(int, ColorsPalette.ROOM_COLORS): + color = palette.get_room_color(palette_id) + if color not in candidates: + candidates.append(color) + + assigned: dict[int, Color] = {} + for room_id in sorted(room_ids): + preferred = palette.get_room_color(room_id) + neighbor_colors = {assigned[neighbor] for neighbor in neighbors[room_id] if neighbor in assigned} + assigned[room_id] = ( + preferred + if preferred not in neighbor_colors + else next((color for color in candidates if color not in neighbor_colors), preferred) + ) + return assigned diff --git a/tests/map/test_b01_q10_map_parser.py b/tests/map/test_b01_q10_map_parser.py index 9e815e29e..1e6277d33 100644 --- a/tests/map/test_b01_q10_map_parser.py +++ b/tests/map/test_b01_q10_map_parser.py @@ -159,6 +159,37 @@ def test_parse_map_packet_allows_zero_room_metadata() -> None: assert packet.rooms == [] +def test_parser_renders_distinct_background_floor_and_walls() -> None: + """Background must not hide the wall layer by sharing its color.""" + grid = bytes([243, 240, 249, 8]) + payload = _synthetic_map_payload(width=4, decoded_layout=grid + b"\x01\x00") + parser = B01Q10MapParser() + + parsed = parser.parse(payload) + + assert parsed.image_content is not None + image = Image.open(io.BytesIO(parsed.image_content)) + scale = parser.config.map_scale + assert image.getpixel((0, 0)) == (0, 0, 0, 0) + assert image.getpixel((scale, 0)) == (32, 115, 185, 255) + assert image.getpixel((scale * 2, 0)) == (93, 109, 126, 255) + assert image.getpixel((scale * 3, 0)) == (133, 193, 233, 255) + + +def test_parser_gives_adjacent_rooms_distinct_palette_colors() -> None: + """Repeated V1 palette entries do not merge neighboring Q10 rooms.""" + grid = bytes([2 * 4, 12 * 4]) + payload = _synthetic_map_payload(width=2, decoded_layout=grid + b"\x01\x00") + parser = B01Q10MapParser() + + parsed = parser.parse(payload) + + assert parsed.image_content is not None + image = Image.open(io.BytesIO(parsed.image_content)) + scale = parser.config.map_scale + assert image.getpixel((0, 0)) != image.getpixel((scale, 0)) + + def test_parse_map_packet_reads_header_height() -> None: """Width and height come straight from the u16be header fields.""" grid = bytes([8]) * 6 + bytes([12]) * 6 # two rooms, 4x3 grid diff --git a/tests/map/test_map_parser.py b/tests/map/test_map_parser.py index 4c6489727..8597360f4 100644 --- a/tests/map/test_map_parser.py +++ b/tests/map/test_map_parser.py @@ -3,9 +3,15 @@ from pathlib import Path import pytest +from vacuum_map_parser_base.config.color import ColorsPalette +from vacuum_map_parser_base.config.image_config import ImageConfig from roborock.exceptions import RoborockException -from roborock.map.map_parser import MapParser, MapParserConfig +from roborock.map.map_parser import ( + MapParser, + MapParserConfig, + _AdjacencyAwareRoborockImageParser, +) MAP_DATA_FILE = Path(__file__).parent / "raw_map_data" DEFAULT_MAP_CONFIG = MapParserConfig() @@ -19,4 +25,38 @@ def test_invalid_map_content(map_content: bytes): parser.parse(map_content) +def test_v1_parser_gives_adjacent_rooms_distinct_palette_colors() -> None: + """Repeated palette entries do not merge neighboring V1 rooms.""" + palette = ColorsPalette() + original_room_12 = palette.get_room_color(12) + image_parser = _AdjacencyAwareRoborockImageParser(palette, ImageConfig()) + raw_data = bytes([(2 << 3) | 7, (12 << 3) | 7]) + + image, _rooms = image_parser.parse(raw_data, 2, 1, None) + + assert image is not None + assert image.getpixel((0, 0)) != image.getpixel((1, 0)) + + isolated_image, _rooms = image_parser.parse(bytes([(12 << 3) | 7]), 1, 1, None) + + assert isolated_image is not None + assert isolated_image.getpixel((0, 0))[: len(original_room_12)] == original_room_12 + + +def test_v1_parser_keeps_adjacent_rooms_hidden_when_rooms_disabled() -> None: + """Adjacency conflict handling cannot override intentional transparency.""" + hidden_rooms = {str(room_id): (0, 0, 0, 0) for room_id in range(1, 32)} + parser = _AdjacencyAwareRoborockImageParser( + ColorsPalette({}, hidden_rooms), + ImageConfig(), + recolor_rooms=False, + ) + raw_data = bytes([(2 << 3) | 7, (12 << 3) | 7]) + + image, _rooms = parser.parse(raw_data, 2, 1, None) + + assert image is not None + assert [image.getpixel((x, 0)) for x in range(2)] == [(0, 0, 0, 0)] * 2 + + # We can add additional tests here in the future that actually parse valid map data