fix: distinguish adjacent map rooms

Author: tubededentifrice

roborock/map/b01_q10_map_parser.py

@@ -40,6 +40,7 @@
     decompose_grid,
 )
 from .map_parser import ParsedMapData
+from .room_colors import adjacency_aware_room_colors
 
 _MAP_FILE_FORMAT = "PNG"
 
@@ -676,7 +677,7 @@ def parsed_from_packet(self, packet: Q10MapPacket) -> ParsedMapData:
 
     def _render(self, packet: Q10MapPacket) -> Image.Image:
         """Render the Q10 grid with the V1 map palette."""
-        palette = _build_palette(packet.grid)
+        palette = _build_palette(packet.grid, packet.width)
         rgba = bytearray()
         for value in packet.grid:
             rgba.extend(palette[value])
@@ -694,13 +695,23 @@ def _opaque(color: tuple[int, ...]) -> tuple[int, int, int, int]:
     return (color[0], color[1], color[2], color[3] if len(color) == 4 else 255)
 
 
-def _build_palette(grid: bytes) -> list[tuple[int, int, int, int]]:
+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(colors.get_room_color(max(1, value // 4)))
+        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] = wall

roborock/map/map_parser.py

@@ -10,10 +10,13 @@
 from vacuum_map_parser_base.config.size import Size, Sizes
 from vacuum_map_parser_base.image_generator import ImageGenerator
 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 = {
@@ -99,16 +102,50 @@ 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) -> None:
+        super().__init__(palette, image_config)
+        self._room_palette = palette
+        self._base_room_colors = palette.cached_room_colors.copy()
+
+    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."""
+        self._room_palette.cached_room_colors.clear()
+        self._room_palette.cached_room_colors.update(self._base_room_colors)
+
+        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."""
     palette, sizes, image_config = _create_rendering_components(config)
-    return RoborockMapDataParser(
+    parser = RoborockMapDataParser(
         palette,
         sizes,
         config.drawables,
         image_config,
         [],
     )
+    parser._image_parser = _AdjacencyAwareRoborockImageParser(palette, image_config)
+    return parser
 
 
 def create_image_generator(

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

@@ -176,6 +176,20 @@ def test_parser_renders_distinct_background_floor_and_walls() -> None:
     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,13 +3,16 @@
 from pathlib import Path
 
 import pytest
+from vacuum_map_parser_base.config.color import ColorsPalette
 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
 
 from roborock.exceptions import RoborockException
 from roborock.map.map_parser import (
     MapParser,
     MapParserConfig,
+    _AdjacencyAwareRoborockImageParser,
     _create_map_data_parser,
     create_image_generator,
 )
@@ -46,4 +49,22 @@ def test_shared_image_generator_matches_v1_rendering_components() -> None:
         assert shared._sizes.get_size(size) == v1._sizes.get_size(size)
 
 
+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
+
+
 # We can add additional tests here in the future that actually parse valid map data