feat: align Q10 map colors with V1 (#902)

* feat: align Q10 map colors with V1

* refactor: clarify room palette cache handling

Author: tubededentifrice

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

roborock/map/map_parser.py

@@ -2,17 +2,21 @@
 
 import io
 import logging
+import threading
 from dataclasses import dataclass, field
 
 from vacuum_map_parser_base.config.color import ColorsPalette, SupportedColor
 from vacuum_map_parser_base.config.drawable import Drawable
 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,52 @@ 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:
+            # cached_room_colors is a read-only property, so reset its dict in place.
+            cached_room_colors = self._room_palette.cached_room_colors
+            cached_room_colors.clear()
+            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():
+                    # ColorsPalette caches both forms for get_room_color(str | int).
+                    cached_room_colors[number] = color
+                    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 +164,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

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

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

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,39 @@ 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))
+    assert palette.get_room_color(12) == palette.get_room_color("12")
+
+    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