From 9d38b8c26e2a51355b8bb22a413404b87a93e2bf Mon Sep 17 00:00:00 2001 From: Vincent <2070309+tubededentifrice@users.noreply.github.com> Date: Sun, 19 Jul 2026 20:18:52 +0400 Subject: [PATCH] refactor(map): address Q10 decoding review feedback --- roborock/map/b01_grid_layers.py | 48 ++++++++++++++++++---------- roborock/map/b01_q10_map_parser.py | 14 ++++---- roborock/map/b01_q10_overlays.py | 27 ++++++++-------- tests/map/test_b01_grid_layers.py | 40 ++--------------------- tests/map/test_b01_q10_map_parser.py | 31 ++++++++++++++++++ tests/map/test_b01_q10_overlays.py | 22 +++++++------ 6 files changed, 97 insertions(+), 85 deletions(-) diff --git a/roborock/map/b01_grid_layers.py b/roborock/map/b01_grid_layers.py index 46026f17..bcc0f457 100644 --- a/roborock/map/b01_grid_layers.py +++ b/roborock/map/b01_grid_layers.py @@ -15,6 +15,7 @@ class (background / wall / per-room floor / ...). This module turns such a grid import io from collections.abc import Callable, Iterable from dataclasses import dataclass, field +from enum import IntEnum from math import ceil from PIL import Image @@ -28,6 +29,14 @@ class (background / wall / per-room floor / ...). This module turns such a grid _PNG = "PNG" +class _CalibrationCell(IntEnum): + """Cell categories used while scoring calibration candidates.""" + + OTHER = 0 + FLOOR = 1 + BLOCKED = 2 + + @dataclass class RoomLayer: """A single room (segment) and where its pixels sit in the grid.""" @@ -140,6 +149,21 @@ def pixel_to_world(self, px: float, py: float) -> tuple[float, float]: return ((px - self.origin_x) * self.resolution, self.y_sign * (self.origin_y - py) * self.resolution) +def _calibration_cells(layers: GridLayers) -> bytes: + """Classify the flat, row-major grid for calibration scoring.""" + cells = bytearray() + for value in layers.grid: + layer = layers.cell_class(value) + if layer == LAYER_FLOOR: + cell = _CalibrationCell.FLOOR + elif layer in (LAYER_WALL, LAYER_BACKGROUND): + cell = _CalibrationCell.BLOCKED + else: + cell = _CalibrationCell.OTHER + cells.append(cell) + return bytes(cells) + + def solve_calibration( layers: GridLayers, points: list[tuple[float, float]], @@ -162,11 +186,7 @@ def solve_calibration( if not points: return None w, h = layers.width, layers.height - classify = layers.classifier - # 1 = floor, 2 = wall/background (blocked), 0 = other. Index by cell. - klass = bytes( - 1 if (c := classify(v)) == LAYER_FLOOR else 2 if c in (LAYER_WALL, LAYER_BACKGROUND) else 0 for v in layers.grid - ) + cells = _calibration_cells(layers) best: tuple[float, GridCalibration] | None = None for resolution in resolutions: @@ -186,10 +206,10 @@ def solve_calibration( blocked = 0 for px_f, py_f in pts: cell = int(oy - py_f) * w + int(px_f + ox) - k = klass[cell] - if k == 1: + cell_type = cells[cell] + if cell_type == _CalibrationCell.FLOOR: on_floor += 1 - elif k == 2: + elif cell_type == _CalibrationCell.BLOCKED: blocked += 1 score = on_floor - 1.5 * blocked if best is None or score > best[0]: @@ -223,11 +243,7 @@ def solve_calibration_with_origin( return None w, h = layers.width, layers.height ox, oy = origin - classify = layers.classifier - # 1 = floor, 2 = wall/background (blocked), 0 = other. Index by cell. - klass = bytes( - 1 if (c := classify(v)) == LAYER_FLOOR else 2 if c in (LAYER_WALL, LAYER_BACKGROUND) else 0 for v in layers.grid - ) + cells = _calibration_cells(layers) best: tuple[float, GridCalibration] | None = None for resolution in resolutions: @@ -242,10 +258,10 @@ def solve_calibration_with_origin( if not (0 <= px < w and 0 <= py < h): blocked += 1 continue - k = klass[py * w + px] - if k == 1: + cell_type = cells[py * w + px] + if cell_type == _CalibrationCell.FLOOR: on_floor += 1 - elif k == 2: + elif cell_type == _CalibrationCell.BLOCKED: blocked += 1 score = on_floor - 1.5 * blocked if best is None or score > best[0]: diff --git a/roborock/map/b01_q10_map_parser.py b/roborock/map/b01_q10_map_parser.py index a3f7d1e7..062f85f3 100644 --- a/roborock/map/b01_q10_map_parser.py +++ b/roborock/map/b01_q10_map_parser.py @@ -64,13 +64,6 @@ def classify_q10_cell(value: int) -> str: return LAYER_FLOOR -def decompose_layers(packet: "Q10MapPacket") -> GridLayers: - """Split a parsed Q10 map packet into separable grid-pixel layers.""" - rooms = [(room.id, room.name, room.pixel_value, room.pixel_count) for room in packet.rooms] - # The ss07 grid is stored top-down (row 0 = top), so no display flip is applied. - return decompose_grid(packet.width, packet.height, packet.grid, rooms, classify_q10_cell, flip=False) - - MAP_PACKET_MARKER = b"\x01\x01" TRACE_PACKET_MARKER = b"\x02\x01" @@ -202,6 +195,13 @@ class Q10MapPacket: the same (top-down) pixel space as :attr:`grid`, where a non-zero cell is carpet (the value is the carpet kind). ``None`` if the packet carried none.""" + @property + def layers(self) -> GridLayers: + """Split the occupancy grid into separable grid-pixel layers.""" + rooms = [(room.id, room.name, room.pixel_value, room.pixel_count) for room in self.rooms] + # The ss07 grid is stored top-down (row 0 = top), so no display flip is applied. + return decompose_grid(self.width, self.height, self.grid, rooms, classify_q10_cell, flip=False) + @dataclass class Q10Point: diff --git a/roborock/map/b01_q10_overlays.py b/roborock/map/b01_q10_overlays.py index 34fcf71c..45b9bd80 100644 --- a/roborock/map/b01_q10_overlays.py +++ b/roborock/map/b01_q10_overlays.py @@ -27,6 +27,7 @@ """ import base64 +import binascii from dataclasses import dataclass, field _DEFAULT_RECORD_SIZE = 38 # 2-byte record header + up to 9 (x, y) int16 pairs @@ -40,25 +41,23 @@ class Q10Zone: vertices: list[tuple[int, int]] = field(default_factory=list) -def _as_bytes(data: bytes | str | None) -> bytes: +def _decode_blob(data: str | None) -> bytes: if data is None: return b"" - if isinstance(data, bytes): - return data try: return base64.b64decode(data + "=" * (-len(data) % 4)) - except (ValueError, base64.binascii.Error): # type: ignore[attr-defined] + except (ValueError, binascii.Error): return b"" -def parse_zone_blob(data: bytes | str | None) -> list[Q10Zone]: +def parse_zone_blob(data: str | None) -> list[Q10Zone]: """Decode a Q10 zone/wall overlay blob into a list of :class:`Q10Zone`. - Accepts the raw bytes or the base64 string straight from the data point. - Returns ``[]`` for empty/absent/unparsable blobs (the device sends a single - ``0x00`` byte when there are none). + Accepts the base64 string straight from the data point. Returns ``[]`` for + empty/absent/unparsable blobs (the device sends a single ``0x00`` byte when + there are none). """ - raw = _as_bytes(data) + raw = _decode_blob(data) if len(raw) < 2: return [] count = raw[1] @@ -94,7 +93,7 @@ def parse_zone_blob(data: bytes | str | None) -> list[Q10Zone]: _WALL_RECORD_SIZE = 8 # two (x, y) int16-BE endpoints -def parse_virtual_wall_blob(data: bytes | str | None) -> list[Q10Zone]: +def parse_virtual_wall_blob(data: str | None) -> list[Q10Zone]: """Decode a Q10 virtual-wall overlay blob (``dpVirtualWallUp`` 57). Virtual walls use a *different framing* from the restricted-zone DPs handled @@ -109,9 +108,9 @@ def parse_virtual_wall_blob(data: bytes | str | None) -> list[Q10Zone]: can place them onto the map through the same :class:`~roborock.map.b01_grid_layers.GridCalibration` as the zones. - Accepts raw bytes or the base64 string straight from the data point. Returns - ``[]`` for empty/absent/unparsable blobs (the device sends a single ``0x00`` - byte -- base64 ``AA==`` -- when there are none). + Accepts the base64 string straight from the data point. Returns ``[]`` for + empty/absent/unparsable blobs (the device sends a single ``0x00`` byte -- + base64 ``AA==`` -- when there are none). The axis order was confirmed against the app: a horizontal wall drawn below a room reads back with x varying and y constant (and the wide RDC no-go zone @@ -119,7 +118,7 @@ def parse_virtual_wall_blob(data: bytes | str | None) -> list[Q10Zone]: the wall axes to ``(y, x)`` -- following a misreading of PR #850's notes -- which placed every wall transposed 90 degrees from where it was drawn. """ - raw = _as_bytes(data) + raw = _decode_blob(data) if len(raw) < 1: return [] count = raw[0] diff --git a/tests/map/test_b01_grid_layers.py b/tests/map/test_b01_grid_layers.py index 52954203..85ab14f7 100644 --- a/tests/map/test_b01_grid_layers.py +++ b/tests/map/test_b01_grid_layers.py @@ -1,7 +1,6 @@ -"""Tests for the device-agnostic grid->layers decomposition + Q10 classifier.""" +"""Tests for the device-agnostic grid-to-layers decomposition.""" import io -from pathlib import Path import pytest from PIL import Image @@ -15,28 +14,6 @@ solve_calibration, solve_calibration_with_origin, ) -from roborock.map.b01_q10_map_parser import ( - classify_q10_cell, - decompose_layers, - parse_map_packet, -) - -FIXTURE = Path(__file__).resolve().parent / "testdata" / "b01_q10_map.bin" - - -@pytest.mark.parametrize( - ("value", "expected"), - [ - (0, "unknown"), - (8, LAYER_FLOOR), - (12, LAYER_FLOOR), - (240, LAYER_FLOOR), - (243, LAYER_BACKGROUND), - (249, LAYER_WALL), - ], -) -def test_classify_q10_cell(value: int, expected: str) -> None: - assert classify_q10_cell(value) == expected def test_decompose_grid_generic_classifier_and_bbox() -> None: @@ -97,21 +74,8 @@ def test_render_scale_upsamples() -> None: assert Image.open(io.BytesIO(png)).size == (6, 3) -def test_decompose_layers_on_q10_fixture() -> None: - """The Q10 synthetic fixture splits into floor + per-room layers.""" - layers = decompose_layers(parse_map_packet(FIXTURE.read_bytes())) - assert layers.class_counts.get(LAYER_FLOOR) == 26 - names = {room.id: room.name for room in layers.rooms} - assert names == {2: "Living Room", 3: "Bedroom"} - # Each room renders to a valid PNG and only its own pixels are opaque. - living = layers.render_room(2, (255, 0, 0, 255)) - img = Image.open(io.BytesIO(living)) - opaque = sum(1 for *_rgb, a in img.getdata() if a > 0) - assert opaque == next(r.pixel_count for r in layers.rooms if r.id == 2) - - def test_render_room_unknown_id_raises() -> None: - layers = decompose_layers(parse_map_packet(FIXTURE.read_bytes())) + layers = decompose_grid(1, 1, b"\x01", [(1, "Room", 1, 1)], lambda _: LAYER_FLOOR) with pytest.raises(KeyError): layers.render_room(999, (0, 0, 0, 255)) diff --git a/tests/map/test_b01_q10_map_parser.py b/tests/map/test_b01_q10_map_parser.py index 43095c57..9e815e29 100644 --- a/tests/map/test_b01_q10_map_parser.py +++ b/tests/map/test_b01_q10_map_parser.py @@ -1,13 +1,17 @@ """Tests for the Roborock Q10 (B01/ss07) map parser.""" +import io from pathlib import Path import pytest +from PIL import Image 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, Q10Room, + classify_q10_cell, is_map_packet, is_trace_packet, lz4_block_decompress, @@ -118,6 +122,33 @@ def test_parse_map_packet() -> None: assert [(r.id, r.raw_name) for r in packet.rooms] == [(2, "rr_living_room"), (3, "bedroom")] +@pytest.mark.parametrize( + ("value", "expected"), + [ + (0, "unknown"), + (8, LAYER_FLOOR), + (12, LAYER_FLOOR), + (240, LAYER_FLOOR), + (243, LAYER_BACKGROUND), + (249, LAYER_WALL), + ], +) +def test_classify_q10_cell(value: int, expected: str) -> None: + assert classify_q10_cell(value) == expected + + +def test_packet_layers_decompose_q10_fixture() -> None: + """The Q10 synthetic fixture splits into floor + per-room layers.""" + layers = parse_map_packet(_payload()).layers + assert layers.class_counts.get(LAYER_FLOOR) == 26 + assert {room.id: room.name for room in layers.rooms} == {2: "Living Room", 3: "Bedroom"} + + living = layers.render_room(2, (255, 0, 0, 255)) + image = Image.open(io.BytesIO(living)) + opaque = sum(1 for *_rgb, alpha in image.getdata() if alpha > 0) + assert opaque == next(room.pixel_count for room in layers.rooms if room.id == 2) + + def test_parse_map_packet_allows_zero_room_metadata() -> None: """A map can be present before the robot has room segmentation records.""" grid = bytes([240, 240, 249, 243, 240, 240]) diff --git a/tests/map/test_b01_q10_overlays.py b/tests/map/test_b01_q10_overlays.py index 1d13eb62..9ed8481c 100644 --- a/tests/map/test_b01_q10_overlays.py +++ b/tests/map/test_b01_q10_overlays.py @@ -24,6 +24,10 @@ def _rect(zone_type: int, corners: list[tuple[int, int]]) -> bytes: return out +def _encoded(blob: bytes) -> str: + return base64.b64encode(blob).decode() + + def test_zone_type_constants() -> None: """ss07 + ioBroker: 0 no-go, 1 virtual wall, 2 no-mop, 3 threshold.""" assert (ZONE_TYPE_NO_GO, ZONE_TYPE_VIRTUAL_WALL, ZONE_TYPE_NO_MOP, ZONE_TYPE_THRESHOLD) == (0, 1, 2, 3) @@ -33,14 +37,14 @@ def test_parse_zone_blob_distinguishes_no_mop_and_threshold() -> None: """A no-mop (2) and a door-threshold (3) zone keep distinct types.""" no_mop = _rect(ZONE_TYPE_NO_MOP, [(0, 0), (10, 0), (10, 10), (0, 10)]) threshold = _rect(ZONE_TYPE_THRESHOLD, [(20, 20), (30, 20), (30, 22), (20, 22)]) - zones = parse_zone_blob(_blob(1, [no_mop, threshold])) + zones = parse_zone_blob(_encoded(_blob(1, [no_mop, threshold]))) assert [z.type for z in zones] == [ZONE_TYPE_NO_MOP, ZONE_TYPE_THRESHOLD] def test_parse_zone_blob_two_typed_rectangles() -> None: rect_a = _rect(ZONE_TYPE_NO_GO, [(0, 0), (10, 0), (10, 10), (0, 10)]) rect_b = _rect(ZONE_TYPE_NO_MOP, [(-5, -5), (5, -5), (5, 5), (-5, 5)]) - zones = parse_zone_blob(_blob(1, [rect_a, rect_b])) + zones = parse_zone_blob(_encoded(_blob(1, [rect_a, rect_b]))) assert [z.type for z in zones] == [ZONE_TYPE_NO_GO, ZONE_TYPE_NO_MOP] assert zones[0].vertices == [(0, 0), (10, 0), (10, 10), (0, 10)] assert zones[1].vertices == [(-5, -5), (5, -5), (5, 5), (-5, 5)] # signed coords @@ -48,29 +52,28 @@ def test_parse_zone_blob_two_typed_rectangles() -> None: def test_parse_zone_blob_accepts_base64() -> None: blob = _blob(1, [_rect(ZONE_TYPE_NO_GO, [(1, 2), (3, 4), (5, 6), (7, 8)])]) - zones = parse_zone_blob(base64.b64encode(blob).decode()) + zones = parse_zone_blob(_encoded(blob)) assert len(zones) == 1 and zones[0].vertices[2] == (5, 6) def test_parse_zone_blob_empty_variants() -> None: assert parse_zone_blob(None) == [] - assert parse_zone_blob(b"\x00") == [] # device's "no zones" sentinel assert parse_zone_blob("AA==") == [] # base64 of 0x00 - assert parse_zone_blob(bytes([1, 0, 0])) == [] # version=1, count=0 + assert parse_zone_blob("AQAA") == [] # version=1, count=0 def test_parse_zone_blob_skips_malformed_record() -> None: # vertex_count claims 9 verts (needs 38 bytes) but record is only 18 -> skipped. bad = bytes([ZONE_TYPE_NO_GO, 9]) + b"\x00" * 16 good = _rect(ZONE_TYPE_NO_GO, [(1, 1), (2, 2), (3, 3), (4, 4)]) - zones = parse_zone_blob(_blob(1, [bad, good])) + zones = parse_zone_blob(_encoded(_blob(1, [bad, good]))) assert len(zones) == 1 and zones[0].vertices[0] == (1, 1) def test_parse_zone_blob_real_record_size_inferred() -> None: """Record size is inferred from total/count (real device uses 38).""" rect = _rect(ZONE_TYPE_NO_GO, [(100, 200), (300, 200), (300, 50), (100, 50)]) - zones = parse_zone_blob(_blob(1, [rect], record_size=38)) + zones = parse_zone_blob(_encoded(_blob(1, [rect], record_size=38))) assert len(zones) == 1 and zones[0].vertices[0] == (100, 200) @@ -148,7 +151,6 @@ def test_parse_virtual_wall_blob_real_rdc_two_walls() -> None: def test_parse_virtual_wall_blob_empty_variants() -> None: assert parse_virtual_wall_blob(None) == [] - assert parse_virtual_wall_blob(b"\x00") == [] # device's "no walls" sentinel assert parse_virtual_wall_blob("AA==") == [] # base64 of 0x00 @@ -156,14 +158,14 @@ def test_parse_virtual_wall_blob_multiple_walls() -> None: """Two walls back-to-back; each is a separate 8-byte (x, y) record.""" wall_a = bytes.fromhex("000a0014001e0028") # (x,y)=(10,20)->(30,40) wall_b = bytes.fromhex("fffb0005fff6000a") # (x,y)=(-5,5)->(-10,10) - walls = parse_virtual_wall_blob(bytes([2]) + wall_a + wall_b) + walls = parse_virtual_wall_blob(_encoded(bytes([2]) + wall_a + wall_b)) assert [w.vertices for w in walls] == [[(10, 20), (30, 40)], [(-5, 5), (-10, 10)]] def test_parse_virtual_wall_blob_truncated_record_dropped() -> None: """A trailing record shorter than 8 bytes is dropped, not misread.""" blob = bytes([2]) + bytes([0x00, 0x0A, 0x00, 0x14, 0x00, 0x1E, 0x00, 0x28]) + b"\x00\x00" - walls = parse_virtual_wall_blob(blob) + walls = parse_virtual_wall_blob(_encoded(blob)) assert [w.vertices for w in walls] == [[(10, 20), (30, 40)]]