From 15b1ab6a1134a18f19285ddc6872bec2a20ee323 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 17 Oct 2026 22:03:08 +0000 Subject: [PATCH] refactor(map): move Q10 rasterization from parser into render module --- roborock/map/b01_q10_map_parser.py | 120 +-------------------------- roborock/map/b01_q10_render.py | 98 ++++++++++++++++++++-- tests/map/test_b01_q10_map_parser.py | 53 ------------ tests/map/test_b01_q10_render.py | 46 ++++++++++ tests/map/test_benchmarks.py | 38 +++++---- 5 files changed, 162 insertions(+), 193 deletions(-) diff --git a/roborock/map/b01_q10_map_parser.py b/roborock/map/b01_q10_map_parser.py index e3c7d61b..0efbbebc 100644 --- a/roborock/map/b01_q10_map_parser.py +++ b/roborock/map/b01_q10_map_parser.py @@ -20,17 +20,11 @@ https://github.com/v1b3c0d3x3r/roborock-qseries-map-bridge """ -import io import math import statistics import struct 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, Point - from roborock.data.b01_q10.b01_q10_containers import Q10RoborockPoint from roborock.data.code_mappings import RoborockEnum from roborock.data.containers import RoborockBase @@ -44,10 +38,6 @@ GridLayers, decompose_grid, ) -from .map_parser import ParsedMapData -from .room_colors import adjacency_aware_room_colors - -_MAP_FILE_FORMAT = "PNG" # Semantic raster classes, confirmed against real ss07 captures (rendered and # eyeballed): 243 is the background outside the home (~half the grid), 240 is @@ -772,115 +762,7 @@ def erased_packet(packet: "Q10MapPacket", cells: set[int]) -> "Q10MapPacket": @dataclass class B01Q10MapParserConfig: - """Configuration for the Q10 map parser.""" + """Configuration for the Q10 map renderer.""" map_scale: int = 4 """Scale factor for the rendered map image.""" - - -class B01Q10MapParser: - """Decoder/renderer for Q10 ``MAP_RESPONSE`` (protocol 301) payloads.""" - - def __init__(self, config: B01Q10MapParserConfig | None = None) -> None: - self._config = config or B01Q10MapParserConfig() - - @property - def config(self) -> B01Q10MapParserConfig: - """The parser configuration (image scale, ...).""" - return self._config - - def parse(self, payload: bytes) -> ParsedMapData: - """Parse a raw Q10 map packet into a rendered PNG + ``MapData``.""" - return self.parsed_from_packet(parse_map_packet(payload)) - - def parse_packet(self, packet: Q10MapPacket) -> ParsedMapData: - """Render an already-decoded packet. - - This preserves the public API used by the existing map trait. Callers - rendering a modified packet can use :meth:`parsed_from_packet` directly. - """ - return self.parsed_from_packet(packet) - - def map_data_from_packet(self, packet: Q10MapPacket) -> MapData: - """Build MapData with a rendered base image, without PNG serialization.""" - image = self._render(packet) - - map_data = MapData() - map_data.image = ImageData( - size=packet.width * packet.height, - top=0, - left=0, - height=packet.height, - width=packet.width, - image_config=ImageConfig(scale=self._config.map_scale), - data=image, - # ImageDimensions uses V1's bottom-up map convention before - # projecting into the top-down PNG. Q10 points are already - # top-down grid pixels, so this adapter cancels that standard flip - # and lets Q10 use the shared V1 ImageGenerator without moving - # overlays vertically. - img_transformation=lambda p: Point(p.x, packet.height - p.y - 1, p.a), - ) - room_names = {room.id: room.name for room in packet.rooms} - if room_names: - map_data.additional_parameters["room_names"] = room_names - - # Carpet cells are flat grid indices (y*width + x) in the same top-down - # pixel space as the rendered raster, so they line up with the image. - if packet.carpet_mask is not None: - map_data.carpet_map = {i for i, value in enumerate(packet.carpet_mask) if value} - - return map_data - - def parsed_from_packet(self, packet: Q10MapPacket) -> ParsedMapData: - """Render a (possibly erase-modified) packet into a PNG + ``MapData``.""" - map_data = self.map_data_from_packet(packet) - image_bytes = io.BytesIO() - if map_data.image is not None: - map_data.image.data.save(image_bytes, format=_MAP_FILE_FORMAT) - return ParsedMapData(image_content=image_bytes.getvalue(), map_data=map_data) - - def _render(self, packet: Q10MapPacket) -> Image.Image: - """Render the Q10 grid with the V1 map palette.""" - palette = _build_palette(packet.grid, packet.width) - flat_palette = [channel for color in palette for channel in color] - p_img = Image.frombytes("P", (packet.width, packet.height), packet.grid) - p_img.putpalette(flat_palette, rawmode="RGBA") - img = p_img.convert("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 - if scale > 1: - img = img.resize((packet.width * scale, packet.height * scale), resample=Image.Resampling.NEAREST) - return img - - -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] = 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/b01_q10_render.py b/roborock/map/b01_q10_render.py index 1aaaac62..1ee9ec00 100644 --- a/roborock/map/b01_q10_render.py +++ b/roborock/map/b01_q10_render.py @@ -1,7 +1,7 @@ """Compose a Q10 (B01/ss07) map into a single rendered result. -The :class:`~roborock.map.b01_q10_map_parser.B01Q10MapParser` turns wire bytes -into a :class:`~roborock.map.b01_q10_map_parser.Q10MapPacket`; this module +The :func:`~roborock.map.b01_q10_map_parser.parse_map_packet` function turns wire +bytes into a :class:`~roborock.map.b01_q10_map_parser.Q10MapPacket`; this module combines that map-protocol packet with the latest trace-protocol packet and DPS overlay snapshot. Calibration, path and position are derived from those source objects rather than managed independently by callers. @@ -18,9 +18,12 @@ from collections.abc import Sequence from dataclasses import dataclass +from PIL import Image +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 Area, MapData, Path, Point, Wall +from vacuum_map_parser_base.map_data import Area, ImageData, MapData, Path, Point, Wall from roborock.exceptions import RoborockException @@ -31,7 +34,9 @@ solve_calibration_with_origin, ) from .b01_q10_map_parser import ( - B01Q10MapParser, + _BACKGROUND_VALUE, + _UNSEGMENTED_FLOOR_VALUE, + _WALL_THRESHOLD, B01Q10MapParserConfig, Q10EraseZone, Q10HistoricalTracePacket, @@ -41,6 +46,7 @@ ) from .b01_q10_overlays import ZONE_TYPE_NO_GO, ZONE_TYPE_NO_MOP, Q10Zone from .map_parser import DEFAULT_DRAWABLES, MapParserConfig, _create_image_generator +from .room_colors import adjacency_aware_room_colors # Path-units-per-pixel candidates for calibration. A dense ss07 path lands a # best fit of 20.0 around the header origin -- ground-truthed June 2026 on the @@ -97,7 +103,6 @@ def render_q10_map( available trace and DPS overlays are projected and drawn in pixel space. Raises :class:`RoborockException` if map rendering fails. """ - parser = B01Q10MapParser(config) trace_calibration = solve_q10_calibration(packet, trace) vector_calibration = _vector_calibration(packet, trace_calibration) @@ -109,7 +114,7 @@ def render_q10_map( # areas disappear (as the app shows). render_packet = erased_packet(packet, cells) - map_data = parser.map_data_from_packet(render_packet) + map_data = _map_data_from_packet(render_packet, config=config) if map_data.image is None: raise RoborockException("Failed to render Q10 map image") @@ -335,3 +340,84 @@ def _draw_map_content( buffer = io.BytesIO() map_data.image.data.save(buffer, format="PNG") return buffer.getvalue() + + +def _map_data_from_packet( + packet: Q10MapPacket, + *, + config: B01Q10MapParserConfig, +) -> MapData: + """Build MapData with a rendered base image, without PNG serialization.""" + image = _render_base_grid(packet, scale=config.map_scale) + + map_data = MapData() + map_data.image = ImageData( + size=packet.width * packet.height, + top=0, + left=0, + height=packet.height, + width=packet.width, + image_config=ImageConfig(scale=config.map_scale), + data=image, + # ImageDimensions uses V1's bottom-up map convention before + # projecting into the top-down PNG. Q10 points are already + # top-down grid pixels, so this adapter cancels that standard flip + # and lets Q10 use the shared V1 ImageGenerator without moving + # overlays vertically. + img_transformation=lambda p: Point(p.x, packet.height - p.y - 1, p.a), + ) + room_names = {room.id: room.name for room in packet.rooms} + if room_names: + map_data.additional_parameters["room_names"] = room_names + + # Carpet cells are flat grid indices (y*width + x) in the same top-down + # pixel space as the rendered raster, so they line up with the image. + if packet.carpet_mask is not None: + map_data.carpet_map = {i for i, value in enumerate(packet.carpet_mask) if value} + + return map_data + + +def _render_base_grid(packet: Q10MapPacket, *, scale: int) -> Image.Image: + """Render the Q10 grid with the V1 map palette.""" + palette = _build_palette(packet.grid, packet.width) + flat_palette = [channel for color in palette for channel in color] + p_img = Image.frombytes("P", (packet.width, packet.height), packet.grid) + p_img.putpalette(flat_palette, rawmode="RGBA") + img = p_img.convert("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. + if scale > 1: + img = img.resize((packet.width * scale, packet.height * scale), resample=Image.Resampling.NEAREST) + return img + + +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] = wall + palette[_UNSEGMENTED_FLOOR_VALUE] = _opaque(colors.get_color(SupportedColor.MAP_INSIDE)) + palette[_BACKGROUND_VALUE] = outside + palette[0] = outside + return palette diff --git a/tests/map/test_b01_q10_map_parser.py b/tests/map/test_b01_q10_map_parser.py index e0585c71..bd7848c2 100644 --- a/tests/map/test_b01_q10_map_parser.py +++ b/tests/map/test_b01_q10_map_parser.py @@ -10,7 +10,6 @@ from roborock.exceptions import RoborockException from roborock.map.b01_grid_layers import LAYER_BACKGROUND, LAYER_FLOOR, LAYER_WALL from roborock.map.b01_q10_map_parser import ( - B01Q10MapParser, Q10MapPacketKind, Q10Point, Q10Room, @@ -170,37 +169,6 @@ 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 @@ -244,24 +212,6 @@ def test_room_pixel_count_matches_grid() -> None: assert room.pixel_count == packet.grid.count(room.pixel_value) -def test_parser_renders_png_and_room_names() -> None: - parsed = B01Q10MapParser().parse(_payload()) - assert parsed.image_content is not None - assert parsed.image_content[:8] == b"\x89PNG\r\n\x1a\n" # PNG magic - assert parsed.map_data is not None - assert parsed.map_data.additional_parameters["room_names"] == {2: "Living Room", 3: "Bedroom"} - - -def test_parse_packet_preserves_decoded_packet_api() -> None: - """The existing map trait can render a packet without re-decoding bytes.""" - packet = parse_map_packet(_payload()) - - parsed = B01Q10MapParser().parse_packet(packet) - - assert parsed.image_content is not None - assert parsed.map_data is not None - - def test_parse_rejects_non_map_packet() -> None: with pytest.raises(RoborockException, match="not a Q10 map packet"): parse_map_packet(b"\x02\x01" + b"\x00" * 40) @@ -462,9 +412,6 @@ def test_parse_carpet_mask_from_map_packet_tail() -> None: assert len(packet.erase_zones) == 1 # erase still parses assert packet.carpet_mask == bytes(carpet) - map_data = B01Q10MapParser().parsed_from_packet(packet).map_data - assert map_data is not None - assert map_data.carpet_map == {10, 20} # flat grid indices of the carpet cells def test_parse_map_packet_without_carpet() -> None: diff --git a/tests/map/test_b01_q10_render.py b/tests/map/test_b01_q10_render.py index fed6b1eb..131b9e86 100644 --- a/tests/map/test_b01_q10_render.py +++ b/tests/map/test_b01_q10_render.py @@ -21,6 +21,7 @@ Q10HeaderCalibration, Q10HistoricalTracePacket, Q10MapPacket, + Q10MapPacketKind, Q10Point, Q10TracePacket, parse_map_packet, @@ -36,6 +37,7 @@ Q10MapOverlays, _calibration_from_header_metadata, _erased_cells, + _map_data_from_packet, _place_docked_robot, _vector_calibration, render_q10_map, @@ -104,6 +106,50 @@ def test_render_base_map_without_calibration() -> None: assert image[:8] == b"\x89PNG\r\n\x1a\n" +def test_render_distinct_background_floor_and_walls() -> None: + """Background must not hide the wall layer by sharing its color.""" + packet = Q10MapPacket( + kind=Q10MapPacketKind.CURRENT, + map_id=1, + width=4, + height=1, + grid=bytes([243, 240, 249, 8]), + ) + rendered = _render(packet) + image = Image.open(io.BytesIO(rendered)) + scale = 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_render_gives_adjacent_rooms_distinct_palette_colors() -> None: + """Repeated V1 palette entries do not merge neighboring Q10 rooms.""" + packet = Q10MapPacket( + kind=Q10MapPacketKind.CURRENT, + map_id=1, + width=2, + height=1, + grid=bytes([2 * 4, 12 * 4]), + ) + rendered = _render(packet) + image = Image.open(io.BytesIO(rendered)) + scale = CONFIG.map_scale + assert image.getpixel((0, 0)) != image.getpixel((scale, 0)) + + +def test_map_data_from_packet_populates_room_names_and_carpet_map() -> None: + """Base MapData populates room_names and flat carpet_map indices.""" + carpet = bytearray(8 * 6) + carpet[10] = 4 + carpet[20] = 3 + packet = replace(_packet(), carpet_mask=bytes(carpet)) + map_data = _map_data_from_packet(packet, config=CONFIG) + assert map_data.additional_parameters["room_names"] == {2: "Living Room", 3: "Bedroom"} + assert map_data.carpet_map == {10, 20} + + def test_render_draws_path_and_position() -> None: """The map and trace derive calibration and draw the robot position.""" packet, trace = _calibrated_inputs() diff --git a/tests/map/test_benchmarks.py b/tests/map/test_benchmarks.py index 1c596b85..392decd4 100644 --- a/tests/map/test_benchmarks.py +++ b/tests/map/test_benchmarks.py @@ -31,7 +31,6 @@ from vacuum_map_parser_base.config.drawable import Drawable from roborock.map.b01_q10_map_parser import ( - B01Q10MapParser, B01Q10MapParserConfig, Q10MapPacket, Q10MapPacketKind, @@ -186,24 +185,25 @@ def test_benchmark_q10_map_packet_unpack(benchmark: BenchmarkFixture) -> None: def test_benchmark_q10_map_parse_and_render(benchmark: BenchmarkFixture) -> None: """Benchmark end-to-end Q10 small map packet decoding and PNG rendering.""" raw_data = _Q10_MAP.read_bytes() - parser = B01Q10MapParser() - result = benchmark(parser.parse, raw_data) + config = B01Q10MapParserConfig() + overlays = Q10MapOverlays() + result = benchmark(lambda: render_q10_map(parse_map_packet(raw_data), None, overlays, config=config)) _validate_image(result, (32, 24)) def test_benchmark_q10_map_full_scale_render_scale_1(benchmark: BenchmarkFixture) -> None: """Benchmark Q10 composite rendering on a realistic 200x200 4-room floorplan at scale 1.""" packet = _build_synthetic_q10_grid(200, 200) - parser = B01Q10MapParser(B01Q10MapParserConfig(map_scale=1)) - result = benchmark(parser.parse_packet, packet) + config = B01Q10MapParserConfig(map_scale=1) + result = benchmark(render_q10_map, packet, None, Q10MapOverlays(), config=config) _validate_image(result, (200, 200)) def test_benchmark_q10_map_full_scale_render_scale_4(benchmark: BenchmarkFixture) -> None: """Benchmark Q10 composite rendering on a realistic 200x200 4-room floorplan at scale 4.""" packet = _build_synthetic_q10_grid(200, 200) - parser = B01Q10MapParser(B01Q10MapParserConfig(map_scale=4)) - result = benchmark(parser.parse_packet, packet) + config = B01Q10MapParserConfig(map_scale=4) + result = benchmark(render_q10_map, packet, None, Q10MapOverlays(), config=config) _validate_image(result, (800, 800)) @@ -236,14 +236,13 @@ def _run_benchmarks(iterations: int, warmup: int, profile: bool) -> None: p_v1_s6_no_draw = MapParser(MapParserConfig(map_scale=4, drawables=[])) p_v1_s6_all_draw = MapParser(MapParserConfig(map_scale=4, drawables=list(Drawable))) - p_q10_small = B01Q10MapParser() - p_q10_scale1 = B01Q10MapParser(B01Q10MapParserConfig(map_scale=1)) - p_q10_scale4 = B01Q10MapParser(B01Q10MapParserConfig(map_scale=4)) + q10_config_scale1 = B01Q10MapParserConfig(map_scale=1) + q10_render_config = B01Q10MapParserConfig(map_scale=4) + q10_overlays = Q10MapOverlays() q10_trace_data = _Q10_TRACE.read_bytes() q10_map_packet = parse_map_packet(q10_data) q10_trace_packet = parse_trace_packet(q10_trace_data) - q10_render_config = B01Q10MapParserConfig(map_scale=4) benchmarks: list[tuple[str, Callable[[], Any]]] = [ ("V1 S5 Map (Scale 4, Default)", lambda: p_v1_s5.parse(v1_s5_data)), @@ -253,15 +252,24 @@ def _run_benchmarks(iterations: int, warmup: int, profile: bool) -> None: ("V1 S6 Map (No Drawables)", lambda: p_v1_s6_no_draw.parse(v1_s6_data)), ("V1 S6 Map (All Drawables)", lambda: p_v1_s6_all_draw.parse(v1_s6_data)), ("Q10 Wire Packet (Unpack only)", lambda: parse_map_packet(q10_data)), - ("Q10 Small Map (Scale 1)", lambda: p_q10_small.parse(q10_data)), - ("Q10 200x200 Map (Scale 1)", lambda: p_q10_scale1.parse_packet(q10_packet_200)), - ("Q10 200x200 Map (Scale 4)", lambda: p_q10_scale4.parse_packet(q10_packet_200)), + ( + "Q10 Small Map (Scale 4)", + lambda: render_q10_map(parse_map_packet(q10_data), None, q10_overlays, config=q10_render_config), + ), + ( + "Q10 200x200 Map (Scale 1)", + lambda: render_q10_map(q10_packet_200, None, q10_overlays, config=q10_config_scale1), + ), + ( + "Q10 200x200 Map (Scale 4)", + lambda: render_q10_map(q10_packet_200, None, q10_overlays, config=q10_render_config), + ), ( "Q10 Render with Trace", lambda: render_q10_map( q10_map_packet, trace=q10_trace_packet, - overlays=Q10MapOverlays(), + overlays=q10_overlays, config=q10_render_config, ), ),