Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 34 additions & 16 deletions roborock/map/b01_q10_map_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -40,6 +40,7 @@
decompose_grid,
)
from .map_parser import ParsedMapData
from .room_colors import adjacency_aware_room_colors

_MAP_FILE_FORMAT = "PNG"

Expand Down Expand Up @@ -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
Expand All @@ -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
61 changes: 58 additions & 3 deletions roborock/map/map_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -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 = {}
Expand All @@ -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
59 changes: 59 additions & 0 deletions roborock/map/room_colors.py
Original file line number Diff line number Diff line change
@@ -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
31 changes: 31 additions & 0 deletions tests/map/test_b01_q10_map_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
42 changes: 41 additions & 1 deletion tests/map/test_map_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Loading