diff --git a/roborock/map/b01_map_parser.py b/roborock/map/b01_map_parser.py index b57912e4..40592de7 100644 --- a/roborock/map/b01_map_parser.py +++ b/roborock/map/b01_map_parser.py @@ -15,11 +15,38 @@ from roborock.exceptions import RoborockException from roborock.map.proto.b01_scmap_pb2 import RobotMap # type: ignore[attr-defined] +from .b01_grid_layers import ( + LAYER_BACKGROUND, + LAYER_FLOOR, + LAYER_WALL, + GridLayers, + decompose_grid, +) from .map_parser import ParsedMapData _MAP_FILE_FORMAT = "PNG" +# The Q7 occupancy grid encodes only these classes (no per-room segmentation in +# the raster -- room ids/names live in the protobuf metadata, not the pixels). +_Q7_WALL_VALUE = 127 +_Q7_FLOOR_VALUE = 128 +_Q7_RENDER_INTENSITY = { + LAYER_BACKGROUND: 0, + LAYER_WALL: 180, + LAYER_FLOOR: 255, +} + + +def classify_q7_cell(value: int) -> str: + """Map a Q7 SCMap grid cell value to a canonical layer class.""" + if value == _Q7_WALL_VALUE: + return LAYER_WALL + if value == _Q7_FLOOR_VALUE: + return LAYER_FLOOR + return LAYER_BACKGROUND # 0 = outside / unknown + + @dataclass class B01MapParserConfig: """Configuration for the B01/Q7 map parser.""" @@ -39,8 +66,9 @@ def parse(self, payload: bytes) -> ParsedMapData: parsed = _parse_scmap_payload(payload) size_x, size_y, grid = _extract_grid(parsed) room_names = _extract_room_names(parsed) + layers = decompose_grid(size_x, size_y, grid, [], classify_q7_cell) - image = _render_occupancy_image(grid, size_x=size_x, size_y=size_y, scale=self._config.map_scale) + image = _render_occupancy_image(layers, scale=self._config.map_scale) map_data = MapData() map_data.image = ImageData( @@ -102,23 +130,15 @@ def _extract_room_names(parsed: RobotMap) -> dict[int, str]: return room_names -def _render_occupancy_image(grid: bytes, *, size_x: int, size_y: int, scale: int) -> Image.Image: - """Render the B01 occupancy grid into a simple image.""" - - # The observed occupancy grid contains only: - # - 0: outside/unknown - # - 127: wall/obstacle - # - 128: floor/free - table = bytearray(range(256)) - table[0] = 0 - table[127] = 180 - table[128] = 255 - - mapped = grid.translate(bytes(table)) - img = Image.frombytes("L", (size_x, size_y), mapped) - img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM).convert("RGB") +def _render_occupancy_image(layers: GridLayers, *, scale: int) -> Image.Image: + """Render canonical Q7 grid classes into the composed map image.""" + mapped = bytes(_Q7_RENDER_INTENSITY[layers.cell_class(value)] for value in layers.grid) + img = Image.frombytes("L", (layers.width, layers.height), mapped) + if layers.flip: + img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM) + img = img.convert("RGB") if scale > 1: - img = img.resize((size_x * scale, size_y * scale), resample=Image.Resampling.NEAREST) + img = img.resize((layers.width * scale, layers.height * scale), resample=Image.Resampling.NEAREST) return img diff --git a/tests/map/test_b01_map_parser.py b/tests/map/test_b01_map_parser.py index 0829182e..803e9e4f 100644 --- a/tests/map/test_b01_map_parser.py +++ b/tests/map/test_b01_map_parser.py @@ -11,7 +11,13 @@ from PIL import Image from roborock.exceptions import RoborockException -from roborock.map.b01_map_parser import B01MapParser, _parse_scmap_payload +from roborock.map.b01_grid_layers import LAYER_BACKGROUND, LAYER_FLOOR, LAYER_WALL +from roborock.map.b01_map_parser import ( + B01MapParser, + B01MapParserConfig, + _parse_scmap_payload, + classify_q7_cell, +) from roborock.map.proto.b01_scmap_pb2 import RobotMap # type: ignore[attr-defined] from roborock.protocols.b01_q7_protocol import create_map_key, decode_map_payload @@ -126,6 +132,29 @@ def test_b01_scmap_parser_maps_observed_schema_fields() -> None: assert not parsed.roomDataInfo[1].HasField("roomName") +def test_classify_q7_cell() -> None: + assert classify_q7_cell(0) == LAYER_BACKGROUND + assert classify_q7_cell(127) == LAYER_WALL + assert classify_q7_cell(128) == LAYER_FLOOR + + +def test_b01_map_parser_renders_shared_q7_layer_classes() -> None: + """The parser keeps shared layer decomposition behind its public API.""" + payload = RobotMap() + payload.mapHead.sizeX = 2 + payload.mapHead.sizeY = 2 + payload.mapData.mapData = bytes([0, 127, 128, 128]) + + parsed = B01MapParser(B01MapParserConfig(map_scale=1)).parse(payload.SerializeToString()) + + assert parsed.image_content is not None + image = Image.open(io.BytesIO(parsed.image_content)) + assert image.getpixel((0, 0)) == (255, 255, 255) + assert image.getpixel((1, 0)) == (255, 255, 255) + assert image.getpixel((0, 1)) == (0, 0, 0) + assert image.getpixel((1, 1)) == (180, 180, 180) + + def test_b01_map_parser_rejects_invalid_payload() -> None: parser = B01MapParser() with pytest.raises(RoborockException, match="Failed to parse B01 SCMap"):