feat: compose Q10 maps in a pure renderer

Author: tubededentifrice

roborock/map/b01_q10_render.py

@@ -0,0 +1,334 @@
+"""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
+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.
+
+It exists so the map trait stays about state management: the trait accumulates
+the pushed inputs and calls :func:`render_q10_map` once per change, holding the
+returned object rather than mutating a pile of derived fields itself. All the
+low-level pixel work (erase-zone blanking, world->pixel overlay placement, path
+drawing) and the calibration policy live here, next to the rest of the map code.
+"""
+
+import io
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from PIL import Image, ImageDraw
+from vacuum_map_parser_base.map_data import Area, MapData, Path, Point, Wall
+
+from roborock.exceptions import RoborockException
+
+from .b01_grid_layers import (
+    GridCalibration,
+    GridLayers,
+    solve_calibration,
+    solve_calibration_with_origin,
+)
+from .b01_q10_map_parser import (
+    B01Q10MapParser,
+    B01Q10MapParserConfig,
+    Q10EraseZone,
+    Q10MapPacket,
+    Q10Room,
+    Q10TracePacket,
+    erased_packet,
+)
+from .b01_q10_overlays import ZONE_TYPE_NO_GO, ZONE_TYPE_NO_MOP, Q10Zone
+
+# 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
+# R1: a corridor drive registered at 20 (matching the format author's
+# independent "20 path-units/px"), and the dock->corridor span lined up with the
+# ruler-measured 8.81 m corridor. With the header resolution=5 (50 mm/px grid)
+# that makes one path-unit exactly 50/20 = 2.5 mm -- so a path-unit is NOT a
+# millimetre (the open scale question). An earlier [10.0..18.0] range couldn't
+# reach 20 (it railed at the bound), biasing the fit. A dense cleaning path
+# selects the best fit within this bracket.
+_Q10_RESOLUTIONS = [step * 0.5 for step in range(24, 53)]  # 12.0 .. 26.0
+# A path needs enough shape to constrain a full (origin + resolution) fit; a few
+# points cannot.
+_MIN_CALIBRATION_POINTS = 20
+# When the grid-frame header supplies the origin, only the resolution is fit, so
+# a much shorter path suffices to confirm it (early in a clean, not just a dense
+# one). See :func:`solve_calibration_with_origin`.
+_MIN_HEADER_CALIBRATION_POINTS = 4
+
+
+@dataclass(frozen=True)
+class Q10MapOverlays:
+    """Latest decoded map-overlay values from the Q10 DPS stream."""
+
+    zones: Sequence[Q10Zone] = ()
+    virtual_walls: Sequence[Q10Zone] = ()
+
+
+@dataclass
+class Q10MapRender:
+    """The fully composed result of rendering a Q10 map packet.
+
+    Built by :func:`render_q10_map` from one map packet, trace packet and DPS
+    overlay snapshot, so every derived field is consistent with one set of
+    source inputs. Analogous to
+    :class:`~roborock.map.map_parser.ParsedMapData`, but also carrying the
+    separable :attr:`layers` and derived :attr:`calibration`.
+    """
+
+    image_content: bytes
+    """The rendered base map (PNG) with erase zones blanked, path not drawn."""
+
+    map_data: MapData
+    """Parsed map data: image metadata, room names, and -- once a calibration is
+    known -- the path / robot position / zones / walls placed in pixel space."""
+
+    layers: GridLayers
+    """Separable map layers (background / wall / floor / per-room) in grid-pixel
+    space, each renderable to a transparent PNG for frontend compositing."""
+
+    rooms: list[Q10Room]
+    """Rooms (segments) reported by the device, with ids and names."""
+
+    calibration: GridCalibration | None
+    """World<->pixel transform used to place the overlays, or ``None`` if no
+    calibration was available (the overlays are then absent from ``map_data``)."""
+
+
+def render_q10_map(
+    packet: Q10MapPacket,
+    trace: Q10TracePacket | None,
+    overlays: Q10MapOverlays,
+    *,
+    config: B01Q10MapParserConfig,
+) -> Q10MapRender:
+    """Compose the latest map, trace and DPS inputs into a render.
+
+    Calibration is derived from ``packet`` (layers + header calibration) and
+    ``trace`` (path points). Once calibrated, erase zones are blanked out of the
+    raster and trace/overlay data is projected into ``map_data`` pixel space.
+    Without a usable trace only the base raster is rendered. Raises
+    :class:`RoborockException` if map rendering fails.
+    """
+    parser = B01Q10MapParser(config)
+    layers = packet.layers
+    calibration = solve_q10_calibration(packet, trace)
+
+    render_packet = packet
+    if calibration is not None:
+        cells = _erased_cells(layers, packet.erase_zones, calibration)
+        if cells:
+            # Blank the erase-zone cells and re-derive the raster/layers from the
+            # modified packet so the phantom areas disappear (as the app shows).
+            render_packet = erased_packet(packet, cells)
+            layers = render_packet.layers
+
+    parsed = parser.parsed_from_packet(render_packet)
+    if parsed.image_content is None or parsed.map_data is None:
+        raise RoborockException("Failed to render Q10 map image")
+    map_data = parsed.map_data
+
+    if calibration is not None and trace is not None:
+        _place_trace(map_data, calibration, trace)
+        _place_overlays(map_data, calibration, overlays)
+
+    return Q10MapRender(
+        image_content=parsed.image_content,
+        map_data=map_data,
+        layers=layers,
+        rooms=packet.rooms,
+        calibration=calibration,
+    )
+
+
+def solve_q10_calibration(
+    packet: Q10MapPacket,
+    trace: Q10TracePacket | None,
+) -> GridCalibration | None:
+    """Derive world-to-pixel calibration from a map and its current trace.
+
+    When the map packet's grid-frame header carries a calibration origin (ss07),
+    only the resolution is fit -- around that fixed origin -- so a short path
+    suffices and the origin is exact rather than recovered by a slide. Otherwise
+    the full origin + resolution fit is used, which needs a reasonably dense
+    cleaning path. Returns ``None`` if the path is too short/featureless to fit.
+    """
+    if trace is None:
+        return None
+    points: list[tuple[float, float]] = [(point.x, point.y) for point in trace.points]
+    return _calibration_from_header(packet, points) or _calibration_from_fit(packet.layers, points)
+
+
+def _calibration_from_header(
+    packet: Q10MapPacket,
+    points: list[tuple[float, float]],
+) -> GridCalibration | None:
+    """Calibrate around the header-supplied origin (resolution fit to a path)."""
+    header_calibration = packet.header_calibration
+    if header_calibration is None or len(points) < _MIN_HEADER_CALIBRATION_POINTS:
+        return None
+    origin = header_calibration.origin_pixels()
+    if origin is None:  # keepalive frame -- no usable origin
+        return None
+    return solve_calibration_with_origin(packet.layers, points, origin, resolutions=_Q10_RESOLUTIONS)
+
+
+def _calibration_from_fit(layers: GridLayers, points: list[tuple[float, float]]) -> GridCalibration | None:
+    """Full origin + resolution fit; needs a reasonably dense path."""
+    if len(points) < _MIN_CALIBRATION_POINTS:
+        return None
+    return solve_calibration(layers, points, resolutions=_Q10_RESOLUTIONS)
+
+
+def _erased_cells(
+    layers: GridLayers,
+    erase_zones: Sequence[Q10EraseZone],
+    calibration: GridCalibration,
+) -> set[int]:
+    """Grid-cell indices covered by the erase zones (axis-aligned bbox fill)."""
+    if not erase_zones:
+        return set()
+    width, height = layers.width, layers.height
+    cells: set[int] = set()
+    for zone in erase_zones:
+        pixels = [calibration.world_to_pixel(x, y) for x, y in zone.vertices]
+        xs = [p[0] for p in pixels]
+        ys = [p[1] for p in pixels]
+        x0, x1 = int(min(xs)), int(max(xs))
+        y0, y1 = int(min(ys)), int(max(ys))
+        for py in range(max(0, y0), min(height, y1 + 1)):
+            for px in range(max(0, x0), min(width, x1 + 1)):
+                cells.add(py * width + px)
+    return cells
+
+
+def _place_trace(
+    map_data: MapData,
+    calibration: GridCalibration,
+    trace: Q10TracePacket,
+) -> None:
+    """Project trace path, position, heading and charger into pixel space.
+
+    Points are stored in grid-pixel space (origin top-left), matching the Q10's
+    top-down, un-flipped raster so they line up with the rendered image.
+    """
+    pixels = [Point(*calibration.world_to_pixel(point.x, point.y)) for point in trace.points]
+    map_data.path = Path(len(pixels), 1, 0, [pixels])
+    robot_position = trace.robot_position
+    if robot_position is not None:
+        px, py = calibration.world_to_pixel(robot_position.x, robot_position.y)
+        map_data.vacuum_position = Point(px, py, trace.heading)
+    if pixels:
+        map_data.charger = pixels[0]
+
+
+def _place_overlays(
+    map_data: MapData,
+    calibration: GridCalibration,
+    overlays: Q10MapOverlays,
+) -> None:
+    """Convert world-coordinate zones/walls into pixel-space ``MapData`` layers."""
+
+    def to_area(zone: Q10Zone) -> Area | None:
+        if len(zone.vertices) != 4:
+            return None  # MapData.Area is a quad
+        pts = [calibration.world_to_pixel(x, y) for x, y in zone.vertices]
+        return Area(pts[0][0], pts[0][1], pts[1][0], pts[1][1], pts[2][0], pts[2][1], pts[3][0], pts[3][1])
+
+    no_go = [area for zone in overlays.zones if zone.type == ZONE_TYPE_NO_GO and (area := to_area(zone))]
+    no_mop = [area for zone in overlays.zones if zone.type == ZONE_TYPE_NO_MOP and (area := to_area(zone))]
+    map_data.no_go_areas = no_go or None
+    map_data.no_mopping_areas = no_mop or None
+
+    walls: list[Wall] = []
+    for zone in overlays.virtual_walls:
+        if len(zone.vertices) >= 2:
+            (x0, y0), (x1, y1) = zone.vertices[0], zone.vertices[1]
+            p0 = calibration.world_to_pixel(x0, y0)
+            p1 = calibration.world_to_pixel(x1, y1)
+            walls.append(Wall(p0[0], p0[1], p1[0], p1[1]))
+    map_data.walls = walls or None
+
+
+def draw_path_on_map(
+    render: Q10MapRender,
+    *,
+    config: B01Q10MapParserConfig,
+    line_color: tuple[int, int, int, int] = (235, 64, 52, 255),
+    position_color: tuple[int, int, int, int] = (255, 211, 0, 255),
+) -> bytes:
+    """Draw the projected ``MapData`` content onto the base map PNG.
+
+    ``render`` must carry its derived calibration. Returns a fresh PNG; the base
+    raster in :attr:`Q10MapRender.image_content` is left untouched.
+    """
+    calibration = render.calibration
+    if calibration is None:
+        raise RoborockException("No calibration available; a cleaning path must be captured during a clean")
+
+    scale = config.map_scale
+    base = Image.open(io.BytesIO(render.image_content)).convert("RGBA")
+
+    def to_image(point: Point) -> tuple[float, float]:
+        return (point.x * scale, point.y * scale)
+
+    draw = ImageDraw.Draw(base, "RGBA")
+
+    # Erase zones are applied to the raster itself (cells blanked), so they are
+    # not drawn here -- the base image already reflects them.
+
+    # No-go (blue) and no-mop (magenta) zones beneath the path.
+    for areas, fill, outline in (
+        (render.map_data.no_go_areas or [], (0, 120, 255, 70), (0, 80, 200, 255)),
+        (render.map_data.no_mopping_areas or [], (255, 0, 200, 70), (200, 0, 160, 255)),
+    ):
+        for area in areas:
+            polygon = [
+                (area.x0 * scale, area.y0 * scale),
+                (area.x1 * scale, area.y1 * scale),
+                (area.x2 * scale, area.y2 * scale),
+                (area.x3 * scale, area.y3 * scale),
+            ]
+            draw.polygon(polygon, fill=fill, outline=outline)
+
+    # Virtual walls (line segments, not polygons) drawn over the zones.
+    for wall in render.map_data.walls or []:
+        draw.line(
+            [(wall.x0 * scale, wall.y0 * scale), (wall.x1 * scale, wall.y1 * scale)],
+            fill=(255, 64, 64, 255),
+            width=max(2, scale),
+        )
+
+    for path in render.map_data.path.path if render.map_data.path else []:
+        if len(path) >= 2:
+            draw.line([to_image(point) for point in path], fill=line_color, width=max(1, scale // 2))
+    if render.map_data.charger is not None:
+        dx, dy = to_image(render.map_data.charger)
+        draw.ellipse([dx - scale, dy - scale, dx + scale, dy + scale], outline=(40, 200, 40, 255), width=2)
+    robot_position = render.map_data.vacuum_position
+    if robot_position is not None:
+        cx, cy = to_image(robot_position)
+        radius = scale
+        draw.ellipse([cx - radius, cy - radius, cx + radius, cy + radius], fill=position_color)
+        robot_heading = robot_position.a
+        if robot_heading is not None:
+            # Heading is world-space degrees (0 = +x, +90 = +y). Map a unit
+            # world-space facing vector through the same transform (so the
+            # Y-flip/scale match the marker), then normalize to a fixed
+            # pixel-length tick so it reads at any calibration resolution.
+            angle = math.radians(robot_heading)
+            dx = math.cos(angle) / calibration.resolution
+            dy = -calibration.y_sign * math.sin(angle) / calibration.resolution
+            norm = math.hypot(dx, dy)
+            if norm > 0:
+                tick = 4 * radius
+                draw.line(
+                    [cx, cy, cx + dx / norm * tick, cy + dy / norm * tick],
+                    fill=position_color,
+                    width=max(1, scale // 2),
+                )
+    buffer = io.BytesIO()
+    base.save(buffer, format="PNG")
+    return buffer.getvalue()