Spaces:
Sleeping
Sleeping
| """Pygame renderer: ghost blueprint, quality-coloured bricks, mm HUD. | |
| Visual language: | |
| - dashed light outlines = unfilled blueprint targets ("ghosts"); the next | |
| expected slot is highlighted; | |
| - placed bricks show a mortar-grey envelope band under a terracotta face; | |
| - once matched, the face is tinted by placement quality: green (within ±3mm), | |
| amber (close), red (poor); strays go dark; | |
| - matched bricks carry a small "+2.1" mm deviation label. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import numpy as np | |
| from atrium_sim.blueprint import BrickKind, Blueprint, brick_face | |
| from atrium_sim.constants import COURSE_MM, H_MAX, TOL_MM | |
| from atrium_sim.physics import BrickPose | |
| SCALE = 0.5 # px per mm | |
| MARGIN_MM = 150.0 | |
| HUD_H = 64 | |
| Y_TOP_MM = H_MAX + 170.0 # show sky so drops/topples are visible | |
| Y_BOT_MM = -30.0 | |
| BG = (24, 26, 32) | |
| GROUND = (70, 72, 78) | |
| GHOST = (110, 115, 125) | |
| NEXT_SLOT = (240, 200, 80) | |
| MORTAR = (105, 100, 95) | |
| TERRACOTTA = (178, 92, 62) | |
| STRAY = (90, 45, 40) | |
| HUD_TEXT = (225, 225, 220) | |
| LABEL = (200, 205, 215) | |
| ROBOT = (90, 170, 220) # mobile gantry body | |
| ROBOT_DARK = (60, 120, 165) | |
| ROBOT_TOOL = (240, 200, 80) # the descending tool / gripper | |
| CEMENT = (150, 148, 143) # curved cement heads (arches) - the "hard material" | |
| STONE = (178, 172, 160) # lintels / sills | |
| STONE_EDGE = (110, 108, 104) | |
| CENTERING = (120, 95, 60) # temporary timber former - visible only until struck | |
| VOUSSOIR_FACE = (178, 92, 62) # real structural arch voussoirs (dynamic bricks, not sensors) | |
| def _quality_color(d: float, in_tol: bool) -> tuple[int, int, int]: | |
| if in_tol: | |
| return (95, 180, 90) | |
| # amber -> red as error grows past tolerance | |
| t = min(1.0, (d - TOL_MM) / 15.0) | |
| return (int(200 + 30 * t), int(150 * (1 - t) + 30), 40) | |
| class PygameRenderer: | |
| def __init__(self, blueprint: Blueprint, render_mode: str | None): | |
| import pygame | |
| self.pygame = pygame | |
| self.blueprint = blueprint | |
| self.render_mode = render_mode | |
| # viewport top is size-aware: tall facades (a 40-course pier, a 13-course house with | |
| # arched heads) must fit, not just the old fixed H_MAX (6 courses). Leaves sky headroom. | |
| self.y_top = max(Y_TOP_MM, blueprint.n_courses * COURSE_MM + 170.0) | |
| self.w_px = int((blueprint.length + 2 * MARGIN_MM) * SCALE) | |
| self.h_px = HUD_H + int((self.y_top - Y_BOT_MM) * SCALE) | |
| pygame.font.init() | |
| self.font = pygame.font.Font(None, 16) | |
| self.hud_font = pygame.font.Font(None, 20) | |
| self.surface = pygame.Surface((self.w_px, self.h_px)) | |
| self.window = None | |
| self.clock = None | |
| if render_mode == "human": | |
| pygame.init() | |
| pygame.display.set_caption("atrium-sim") | |
| self.window = pygame.display.set_mode((self.w_px, self.h_px)) | |
| self.clock = pygame.time.Clock() | |
| # --- coordinate transform ------------------------------------------------- | |
| def _to_px(self, x: float, y: float) -> tuple[float, float]: | |
| return (x + MARGIN_MM) * SCALE, HUD_H + (self.y_top - y) * SCALE | |
| def _rect_corners(self, cx, cy, w, h, theta): | |
| c, s = math.cos(theta), math.sin(theta) | |
| pts = [] | |
| for dx, dy in ((-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2), (-w / 2, h / 2)): | |
| pts.append(self._to_px(cx + dx * c - dy * s, cy + dx * s + dy * c)) | |
| return pts | |
| def _poly_corners(self, cx, cy, theta, local_verts): | |
| """Generalizes _rect_corners to an ARBITRARY canonical polygon (e.g. a tapered arch | |
| voussoir wedge, not just a box): rotate `local_verts` by theta, translate to (cx, cy), | |
| return pixel coords - same contract as _rect_corners.""" | |
| c, s = math.cos(theta), math.sin(theta) | |
| return [self._to_px(cx + lx * c - ly * s, cy + lx * s + ly * c) for lx, ly in local_verts] | |
| def _dashed_rect(self, color, cx, cy, w, h, dash=6): | |
| pg = self.pygame | |
| corners = self._rect_corners(cx, cy, w, h, 0.0) | |
| for p1, p2 in zip(corners, corners[1:] + corners[:1]): | |
| length = math.hypot(p2[0] - p1[0], p2[1] - p1[1]) | |
| n = max(1, int(length / dash)) | |
| for i in range(0, n, 2): | |
| a = (p1[0] + (p2[0] - p1[0]) * i / n, p1[1] + (p2[1] - p1[1]) * i / n) | |
| b = ( | |
| p1[0] + (p2[0] - p1[0]) * min(i + 1, n) / n, | |
| p1[1] + (p2[1] - p1[1]) * min(i + 1, n) / n, | |
| ) | |
| pg.draw.line(self.surface, color, a, b, 1) | |
| # --- main draw -------------------------------------------------------------- | |
| def draw(self, poses: list[BrickPose], report, hud: dict, cursor: int, | |
| robot: tuple | None = None, hard_bodies: list | None = None) -> np.ndarray | None: | |
| pg = self.pygame | |
| surf = self.surface | |
| surf.fill(BG) | |
| # ground | |
| gx0, gy0 = self._to_px(-MARGIN_MM, 0) | |
| pg.draw.rect( | |
| surf, GROUND, pg.Rect(0, gy0, self.w_px, self.h_px - gy0) | |
| ) | |
| # mobile robot: shaded reach window (behind everything; the body is drawn | |
| # on top after the bricks so it reads clearly) | |
| if robot is not None: | |
| base_x, reach = robot[0], robot[1] | |
| lx, _ = self._to_px(base_x - reach, 0) | |
| rx, _ = self._to_px(base_x + reach, 0) | |
| band = pg.Surface((max(1, rx - lx), self.h_px - gy0), pg.SRCALPHA) | |
| band.fill((240, 200, 80, 20)) | |
| surf.blit(band, (lx, gy0)) | |
| # ghost blueprint + next expected slot | |
| matched_ids = {m.target_id for m in report.matches} | |
| open_in_cursor = [ | |
| t for t in self.blueprint.course_targets(cursor) if t.tid not in matched_ids | |
| ] | |
| next_tid = min((t.tid for t in sorted(open_in_cursor, key=lambda t: t.x)), default=None) | |
| for t in self.blueprint.targets: | |
| if t.tid in matched_ids: | |
| continue | |
| w, h = brick_face(t.kind) | |
| if t.tid == next_tid: | |
| pg.draw.polygon(surf, NEXT_SLOT, self._rect_corners(t.x, t.y, w, h, 0.0), 2) | |
| else: | |
| self._dashed_rect(GHOST, t.x, t.y, w, h) | |
| # static hard bodies: lintels/sills/cement heads (permanent) and the arch centering/ | |
| # skewback (centering is temporary - visible only until the ring closes and it's struck; | |
| # "voussoir" here is the older cosmetic lintel_soldier style's sensor fringe, kept | |
| # distinct from the real dynamic VOUSSOIR bricks drawn below) | |
| for _sid, kind, verts in (hard_bodies or []): | |
| pts = [self._to_px(x, y) for x, y in verts] | |
| if len(pts) < 3: | |
| continue | |
| if kind == "voussoir": | |
| fill, edge = TERRACOTTA, MORTAR | |
| elif kind == "centering": | |
| fill, edge = CENTERING, (80, 60, 35) | |
| elif kind == "skewback": | |
| fill, edge = STONE, STONE_EDGE | |
| elif kind == "cement": | |
| fill, edge = CEMENT, STONE_EDGE | |
| else: | |
| fill, edge = STONE, STONE_EDGE | |
| pg.draw.polygon(surf, fill, pts) | |
| pg.draw.polygon(surf, edge, pts, 2) | |
| # bricks | |
| match_by_brick = {m.brick_id: m for m in report.matches} | |
| stray_ids = set(report.stray_bricks) | |
| for p in poses: | |
| if p.kind == BrickKind.VOUSSOIR and p.verts: | |
| # a real structural arch wedge: arbitrary polygon, not a fixed (w, h) box - not | |
| # scored by the flat-wall audit (report never contains a VOUSSOIR match), so no | |
| # quality tint/label; rendered as its own distinct material instead. | |
| pts = self._poly_corners(p.x, p.y, p.theta, p.verts) | |
| if len(pts) >= 3: | |
| pg.draw.polygon(surf, VOUSSOIR_FACE, pts) | |
| pg.draw.polygon(surf, MORTAR, pts, 2) | |
| continue | |
| w, h = brick_face(p.kind) | |
| m = match_by_brick.get(p.brick_id) | |
| if m is not None: | |
| face = _quality_color(m.d, m.in_tol) | |
| elif p.brick_id in stray_ids: | |
| face = STRAY | |
| else: | |
| face = TERRACOTTA # mid-fall / not yet audited | |
| pg.draw.polygon(surf, MORTAR, self._rect_corners(p.x, p.y, w + 9, h + 9, p.theta)) | |
| pg.draw.polygon(surf, face, self._rect_corners(p.x, p.y, w, h, p.theta)) | |
| if m is not None: | |
| label = self.font.render(f"{m.dx:+.1f}", True, LABEL) | |
| lx, ly = self._to_px(p.x, p.y + h / 2 + 14) | |
| surf.blit(label, (lx - label.get_width() / 2, ly)) | |
| # mobile-robot body (drawn on top so it's clearly visible) | |
| if robot is not None: | |
| self._draw_robot_body(robot, poses) | |
| # HUD | |
| pg.draw.rect(surf, (18, 19, 24), pg.Rect(0, 0, self.w_px, HUD_H)) | |
| text = " ".join(f"{k}: {v}" for k, v in hud.items()) | |
| surf.blit(self.hud_font.render("atrium-sim", True, NEXT_SLOT), (10, 8)) | |
| surf.blit(self.hud_font.render(text, True, HUD_TEXT), (10, 34)) | |
| if self.render_mode == "human": | |
| self.window.blit(surf, (0, 0)) | |
| pg.event.pump() | |
| pg.display.flip() | |
| self.clock.tick(60) | |
| frame = np.transpose(pg.surfarray.array3d(surf), (1, 0, 2)) | |
| return frame | |
| def _draw_robot_body(self, robot: tuple, poses: list[BrickPose]) -> None: | |
| """A mobile gantry: wheeled chassis on the rail, a mast, a top beam spanning | |
| the reach, and a tool that descends from the beam to the brick being placed | |
| (so a high release visibly drops from up top).""" | |
| pg = self.pygame | |
| surf = self.surface | |
| base_x, reach = robot[0], robot[1] | |
| gantry_y = robot[2] if len(robot) > 2 and robot[2] else (H_MAX + 60.0) | |
| bx, gy0 = self._to_px(base_x, 0.0) | |
| _, beam_py = self._to_px(base_x, gantry_y) | |
| bx, gy0, beam_py = int(bx), int(gy0), int(beam_py) | |
| # top beam spanning the reach window | |
| half = int(reach * SCALE) | |
| pg.draw.line(surf, ROBOT, (bx - half, beam_py), (bx + half, beam_py), 5) | |
| # vertical mast up the middle | |
| pg.draw.line(surf, ROBOT, (bx, gy0 - 10), (bx, beam_py), 6) | |
| # wheeled chassis on the rail | |
| pg.draw.rect(surf, ROBOT_DARK, pg.Rect(bx - 26, gy0 - 15, 52, 15), border_radius=4) | |
| for wx in (bx - 15, bx + 15): | |
| pg.draw.circle(surf, (28, 30, 36), (wx, gy0), 6) | |
| pg.draw.circle(surf, ROBOT, (wx, gy0), 6, 2) | |
| # tool: descends from the beam to the current brick (the last-spawned pose), | |
| # clamped into the reach window — shows how high the release was | |
| if poses: | |
| p = poses[-1] | |
| tool_x = min(max(p.x, base_x - reach), base_x + reach) | |
| tpx, _ = self._to_px(tool_x, 0.0) | |
| _, tpy = self._to_px(0.0, p.y + brick_face(p.kind)[1] / 2) | |
| tpx, tpy = int(tpx), int(tpy) | |
| pg.draw.line(surf, ROBOT, (tpx, beam_py), (tpx, beam_py + 4), 8) # trolley on beam | |
| pg.draw.line(surf, ROBOT_TOOL, (tpx, beam_py), (tpx, tpy), 3) # descending tool | |
| pg.draw.circle(surf, ROBOT_TOOL, (tpx, tpy), 5) # gripper head | |
| def close(self): | |
| if self.window is not None: | |
| self.pygame.display.quit() | |
| self.window = None | |