119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
"""Core engine running the simulation, vision, physics, and rendering loops."""
|
|
|
|
import cv2
|
|
import pygame
|
|
import numpy as np
|
|
from typing import Optional, List, Tuple
|
|
from .config import Settings, CalibrationData
|
|
from .vision import VisionPipeline
|
|
from .physics import PhysicsWorld
|
|
from .pegman import PegMan
|
|
|
|
|
|
class SimulationEngine:
|
|
def __init__(self, settings: Settings, calibration: Optional[CalibrationData] = None, sim_mode: bool = False):
|
|
self.settings = settings
|
|
self.calibration = calibration
|
|
self.sim_mode = sim_mode
|
|
|
|
self.vision: Optional[VisionPipeline] = None
|
|
if not self.sim_mode:
|
|
self.vision = VisionPipeline(settings, calibration)
|
|
if not self.vision.start_camera():
|
|
print("[Engine] Could not open camera. Falling back to mouse sim mode.")
|
|
self.sim_mode = True
|
|
self.vision = None
|
|
|
|
self.physics = PhysicsWorld(settings)
|
|
self.pegman = PegMan(
|
|
self.physics,
|
|
settings,
|
|
start_pos=(settings.projector_width // 2, settings.projector_height - 100)
|
|
)
|
|
|
|
self.sim_polygons: List[np.ndarray] = []
|
|
self.latest_detected_polygons: List[np.ndarray] = []
|
|
self.latest_warped_frame: Optional[np.ndarray] = None
|
|
self.latest_raw_frame: Optional[np.ndarray] = None
|
|
self.latest_debug_mask: Optional[np.ndarray] = None
|
|
|
|
# Standalone projection window surface (if opened directly in pygame)
|
|
self.screen_surface = pygame.Surface((self.settings.projector_width, self.settings.projector_height))
|
|
|
|
def update_calibration(self, calibration: CalibrationData) -> None:
|
|
self.calibration = calibration
|
|
if self.vision:
|
|
self.vision.set_calibration_points(calibration.points)
|
|
|
|
def capture_background(self) -> None:
|
|
if self.vision:
|
|
self.vision.capture_reference_background()
|
|
|
|
def reset_pegman(self) -> None:
|
|
self.pegman.body.position = (self.settings.projector_width // 2, self.settings.projector_height - 100)
|
|
self.pegman.body.velocity = (0, 0)
|
|
|
|
def add_sim_polygon(self, pts: List[Tuple[int, int]]) -> None:
|
|
if len(pts) >= 3:
|
|
self.sim_polygons.append(np.array(pts, dtype=np.int32))
|
|
|
|
def clear_sim_polygons(self) -> None:
|
|
self.sim_polygons.clear()
|
|
|
|
def step(self, dt: float) -> None:
|
|
"""Runs one frame of computer vision, physics update, and agent decisions."""
|
|
dt = min(dt, 0.05)
|
|
|
|
polygons: List[np.ndarray] = []
|
|
|
|
if self.vision:
|
|
raw = self.vision.read_raw_frame()
|
|
if raw is not None:
|
|
self.latest_raw_frame = raw
|
|
warped = self.vision.warp_frame(raw)
|
|
self.latest_warped_frame = warped
|
|
polygons, mask = self.vision.extract_shadow_polygons(warped)
|
|
self.latest_debug_mask = mask
|
|
self.latest_detected_polygons = polygons
|
|
|
|
if self.sim_mode:
|
|
polygons = self.sim_polygons
|
|
self.latest_detected_polygons = polygons
|
|
|
|
# Update dynamic physics colliders
|
|
self.physics.update_shadow_geometry(polygons)
|
|
|
|
# Update agent and step physics space
|
|
self.pegman.update(dt)
|
|
self.physics.step(dt)
|
|
|
|
def render(self, surface: pygame.Surface, draw_debug_mask: bool = False) -> None:
|
|
"""Renders the projector scene to any given Pygame surface."""
|
|
# 1. Clean white canvas
|
|
surface.fill((255, 255, 255))
|
|
|
|
# 2. Border
|
|
w = self.settings.projector_width
|
|
h = self.settings.projector_height
|
|
pygame.draw.rect(surface, (40, 40, 40), (0, 0, w, h), 4)
|
|
|
|
# 3. Shadows / solids
|
|
for poly in self.latest_detected_polygons:
|
|
if len(poly) >= 3:
|
|
pts = [(int(p[0]), int(p[1])) for p in poly]
|
|
pygame.draw.polygon(surface, (230, 230, 230), pts)
|
|
pygame.draw.lines(surface, (100, 100, 100), True, pts, 2)
|
|
|
|
# 4. Peg Man
|
|
self.pegman.render(surface)
|
|
|
|
# 5. Optional debug mask
|
|
if draw_debug_mask and self.latest_debug_mask is not None:
|
|
mask_rgb = cv2.cvtColor(self.latest_debug_mask, cv2.COLOR_GRAY2RGB)
|
|
mask_surf = pygame.surfarray.make_surface(np.transpose(mask_rgb, (1, 0, 2)))
|
|
mask_surf.set_alpha(120)
|
|
surface.blit(mask_surf, (0, 0))
|
|
|
|
def close(self) -> None:
|
|
if self.vision:
|
|
self.vision.stop_camera()
|