956 lines
26 KiB
Python
956 lines
26 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||
# SPDX-FileCopyrightText: 2023-2024 Wojtek Porczyk <woju@hackerspace.pl>
|
||
|
||
import datetime
|
||
import itertools
|
||
import sys
|
||
import traceback
|
||
|
||
import cv2
|
||
import funcparserlib.parser
|
||
import pygame
|
||
|
||
from . import (
|
||
assets,
|
||
badges,
|
||
camera,
|
||
detector as _detector,
|
||
parser,
|
||
utils,
|
||
)
|
||
from .theme import Colour
|
||
|
||
PYGAME_EVENT_UPDATE_CPU_LOAD = pygame.event.custom_type()
|
||
|
||
|
||
class Layout:
|
||
layout_width = 6
|
||
layout_height = 2
|
||
|
||
def __init__(self, screen):
|
||
self.screen = screen
|
||
self.screen_width, self.screen_height = self.screen.get_size()
|
||
assert self.screen_width >= 720
|
||
self.button_size = self.screen_width // self.layout_width
|
||
|
||
self.font = pygame.font.Font(
|
||
assets / 'fonts/SpaceGrotesk-Regular.otf', 36
|
||
)
|
||
self.font_mono = pygame.font.Font(
|
||
assets / 'fonts/Iosevka-Regular.ttf', 24
|
||
)
|
||
self.font_key = pygame.font.Font(
|
||
assets / 'fonts/Pilowlava-Regular.otf', 56
|
||
)
|
||
|
||
def place_button(
|
||
self, *, left=None, right=None, top=None, bottom=None, width=1, height=1
|
||
):
|
||
assert (left is not None) ^ (right is not None)
|
||
assert (top is not None) ^ (bottom is not None)
|
||
|
||
if left is not None:
|
||
x = left * self.button_size
|
||
if right is not None:
|
||
x = self.screen_width - (left + width) * self.button_size
|
||
if top is not None:
|
||
y = top * self.button_size
|
||
if bottom is not None:
|
||
y = self.screen_height - (bottom + height) * self.button_size
|
||
|
||
return {
|
||
'left': x,
|
||
'top': y,
|
||
'width': width * self.button_size,
|
||
'height': height * self.button_size,
|
||
}
|
||
|
||
|
||
class Control:
|
||
def __init__(self, game):
|
||
self.game = game
|
||
|
||
def handle_event(self, event):
|
||
pass
|
||
|
||
def get_blits(self):
|
||
raise NotImplementedError()
|
||
|
||
|
||
class Status(Control):
|
||
def get_blits(self):
|
||
surface = self.game.layout.font_mono.render(
|
||
f' FPS: {self.game.clock.get_fps():4.1f} Hz '
|
||
f' {datetime.datetime.now():%H:%M} ',
|
||
True,
|
||
Colour.ACCENT3,
|
||
Colour.BACKGROUND,
|
||
)
|
||
yield surface, (self.game.layout.screen_width - surface.get_width(), 0)
|
||
|
||
|
||
class Sonar(Control):
|
||
def get_blits(self):
|
||
if self.game.bot.last_range is None:
|
||
return
|
||
|
||
surface = self.game.layout.font_mono.render(
|
||
(
|
||
' \N{INFINITY} '
|
||
if self.game.bot.last_range == float('inf')
|
||
else f' {self.game.bot.last_range:3} cm '
|
||
),
|
||
True,
|
||
Colour.ACCENT2,
|
||
Colour.BACKGROUND,
|
||
)
|
||
|
||
yield surface, (
|
||
int(self.game.layout.button_size * 3.5 - surface.get_width() / 2),
|
||
self.game.layout.screen_height
|
||
- self.game.layout.button_size * 2
|
||
- surface.get_height()
|
||
- 5,
|
||
)
|
||
|
||
|
||
class Debug(Control):
|
||
def label(
|
||
self,
|
||
text,
|
||
*,
|
||
left,
|
||
top,
|
||
color=Colour.ACCENT3,
|
||
background=Colour.BACKGROUND,
|
||
):
|
||
surface = self.game.layout.font_mono.render(
|
||
text, True, color, background
|
||
)
|
||
return surface, (left, top * surface.get_height())
|
||
|
||
def get_blits(self):
|
||
yield self.label(
|
||
f'EXP:'
|
||
f' {self.game.metadata.ExposureTime:5} us'
|
||
f' (1/{1000000/self.game.metadata.ExposureTime:3.0f} s)',
|
||
top=2,
|
||
left=0,
|
||
)
|
||
yield self.label(
|
||
f'GAIN:'
|
||
f' A{self.game.metadata.AnalogueGain:.1f}'
|
||
f' D{self.game.metadata.DigitalGain:.1f}'
|
||
f' R{self.game.metadata.ColourGains[0]:.1f}'
|
||
f' B{self.game.metadata.ColourGains[1]:.1f}',
|
||
top=3,
|
||
left=0,
|
||
)
|
||
yield self.label(
|
||
f'ILLUM: {self.game.metadata.Lux:4.0f} lux', top=3, left=500
|
||
)
|
||
|
||
yield self.label(
|
||
f'TEMP: {self.game.metadata.SensorTemperature:2.0f} °C',
|
||
top=0,
|
||
left=100,
|
||
)
|
||
|
||
height, width, _ = self.game.frame.shape
|
||
|
||
yield self.label(f'W: {width} × H: {height}', top=1, left=0)
|
||
|
||
|
||
class RectControl(Control):
|
||
def __init__(self, *args, left, top, width, height, **kwds):
|
||
super().__init__(*args, **kwds)
|
||
self.rect = pygame.Rect(left, top, width, height)
|
||
self.uv = self.rect.move(-self.rect.x, -self.rect.y)
|
||
|
||
def get_blits(self):
|
||
surface = pygame.Surface(self.rect.size, pygame.SRCALPHA)
|
||
self.draw(surface)
|
||
yield surface, self.rect.topleft
|
||
|
||
def draw(self, surface):
|
||
raise NotImplementedError()
|
||
|
||
|
||
class BaseButton(RectControl):
|
||
key = None
|
||
|
||
def __init__(self, *args, key=None, on_press=None, **kwds):
|
||
super().__init__(*args, **kwds)
|
||
|
||
if key is not None:
|
||
self.key = key
|
||
if on_press is not None:
|
||
self.on_press = on_press
|
||
|
||
def on_press(self, pos): # pylint: disable=method-hidden
|
||
raise NotImplementedError()
|
||
|
||
def handle_event(self, event):
|
||
if (
|
||
event.type == pygame.MOUSEBUTTONDOWN
|
||
and event.button == 1
|
||
and self.rect.collidepoint(event.pos)
|
||
) or (event.type == pygame.KEYDOWN and event.key == self.key):
|
||
return self.on_press(event.pos)
|
||
return None
|
||
|
||
|
||
class Button(BaseButton):
|
||
# pylint: disable=abstract-method
|
||
label = None
|
||
|
||
def __init__(self, *args, on_press=None, label=None, **kwds):
|
||
super().__init__(*args, **kwds)
|
||
|
||
if label is not None:
|
||
self.label = label
|
||
|
||
def get_blits(self):
|
||
# for the interest of borders, we need to draw a bit bigger rectangle
|
||
rect = pygame.Rect(
|
||
self.rect.x - 1,
|
||
self.rect.y - 1,
|
||
self.rect.width + 2,
|
||
self.rect.height + 2,
|
||
)
|
||
|
||
if self.rect.left == 0:
|
||
rect.x = 0
|
||
rect.width -= 1
|
||
if self.rect.right == self.game.screen_width:
|
||
rect.width -= 1
|
||
if self.rect.bottom == self.game.screen_height:
|
||
rect.height -= 1
|
||
|
||
surface = pygame.Surface(rect.size)
|
||
self.draw(surface)
|
||
yield surface, rect.topleft
|
||
|
||
def draw_label(self, surface):
|
||
fontsurf = self.game.layout.font.render(
|
||
self.label, True, Colour.PRIMARY
|
||
)
|
||
fontrect = fontsurf.get_rect()
|
||
surface.blit(
|
||
fontsurf,
|
||
(
|
||
(self.rect.width - fontrect.width) // 2,
|
||
(self.rect.height - fontrect.height) // 2,
|
||
),
|
||
)
|
||
|
||
def draw_key(self, surface):
|
||
fontsurf = self.game.layout.font_key.render(
|
||
self.key, True, Colour.ACCENT5, Colour.BACKGROUND
|
||
)
|
||
fontrect = fontsurf.get_rect()
|
||
surface.blit(
|
||
fontsurf,
|
||
(
|
||
(self.rect.width - fontrect.width) // 2,
|
||
(self.rect.height - fontrect.height) // 2,
|
||
),
|
||
)
|
||
|
||
def draw(self, surface):
|
||
surface.fill(Colour.BACKGROUND)
|
||
pygame.draw.rect(surface, Colour.ACCENT3, surface.get_rect(), width=2)
|
||
|
||
if self.key is not None:
|
||
self.draw_key(surface)
|
||
|
||
self.draw_label(surface)
|
||
|
||
|
||
class UpButton(BaseButton):
|
||
# pylint: disable=abstract-method
|
||
def draw(self, surface):
|
||
surface.fill(Colour.BACKGROUND)
|
||
pygame.draw.lines(
|
||
surface,
|
||
Colour.ACCENT3,
|
||
False,
|
||
(
|
||
(self.rect.width * 0.2, self.rect.height - 1),
|
||
(0, self.rect.height - 1),
|
||
(0, 0),
|
||
(self.rect.width - 2, 0),
|
||
(self.rect.width - 2, self.rect.height - 1),
|
||
(self.rect.width * 0.8, self.rect.height - 1),
|
||
),
|
||
width=2,
|
||
)
|
||
|
||
fontsurf = self.game.layout.font.render(
|
||
'+', True, Colour.ACCENT1, Colour.BACKGROUND
|
||
)
|
||
fontrect = fontsurf.get_rect()
|
||
surface.blit(fontsurf, ((self.rect.width - fontrect.width) // 2, 2))
|
||
|
||
|
||
class DownButton(BaseButton):
|
||
# pylint: disable=abstract-method
|
||
def draw(self, surface):
|
||
surface.fill(Colour.BACKGROUND)
|
||
pygame.draw.lines(
|
||
surface,
|
||
Colour.ACCENT3,
|
||
False,
|
||
(
|
||
(self.rect.width * 0.2, 0),
|
||
(0, 0),
|
||
(0, self.rect.height - 2),
|
||
(self.rect.width - 2, self.rect.height - 2),
|
||
(self.rect.width - 2, 0),
|
||
(self.rect.width * 0.8, 0),
|
||
),
|
||
width=2,
|
||
)
|
||
|
||
fontsurf = self.game.layout.font.render(
|
||
'-', True, Colour.ACCENT1, Colour.BACKGROUND
|
||
)
|
||
fontrect = fontsurf.get_rect()
|
||
surface.blit(
|
||
fontsurf,
|
||
(
|
||
(self.rect.width - fontrect.width) // 2,
|
||
self.rect.height - fontrect.height - 2,
|
||
),
|
||
)
|
||
|
||
|
||
class Rocker(RectControl):
|
||
label = None
|
||
|
||
def __init__(self, *args, key_up=None, key_down=None, **kwds):
|
||
super().__init__(*args, **kwds)
|
||
self.up = UpButton(
|
||
self.game,
|
||
left=self.rect.left,
|
||
top=self.rect.top,
|
||
width=self.rect.width,
|
||
height=self.rect.centery - self.rect.top,
|
||
key=key_up,
|
||
on_press=self.on_press_up,
|
||
)
|
||
self.down = DownButton(
|
||
self.game,
|
||
left=self.rect.left,
|
||
top=self.rect.centery,
|
||
width=self.rect.width,
|
||
height=self.rect.height - self.up.rect.height,
|
||
key=key_down,
|
||
on_press=self.on_press_down,
|
||
)
|
||
|
||
def draw(self, surface):
|
||
fontsurf = self.game.layout.font.render(
|
||
self.label, True, Colour.PRIMARY, Colour.BACKGROUND
|
||
)
|
||
fontrect = fontsurf.get_rect()
|
||
surface.blit(
|
||
fontsurf,
|
||
(
|
||
self.uv.centerx - fontrect.width // 2,
|
||
self.uv.centery - fontrect.height - 2,
|
||
),
|
||
)
|
||
|
||
fontsurf = self.game.layout.font_mono.render(
|
||
self.get_value(), True, Colour.ACCENT2, Colour.BACKGROUND
|
||
)
|
||
fontrect = fontsurf.get_rect()
|
||
surface.blit(
|
||
fontsurf,
|
||
(self.uv.centerx - fontrect.width // 2, self.uv.centery + 2),
|
||
)
|
||
|
||
def get_value(self):
|
||
raise NotImplementedError()
|
||
|
||
def get_blits(self):
|
||
yield from self.up.get_blits()
|
||
yield from self.down.get_blits()
|
||
yield from super().get_blits()
|
||
|
||
def handle_event(self, event):
|
||
return self.up.handle_event(event) or self.down.handle_event(event)
|
||
|
||
def on_press_up(self, pos):
|
||
raise NotImplementedError()
|
||
|
||
def on_press_down(self, pos):
|
||
raise NotImplementedError()
|
||
|
||
|
||
class PWMRocker(Rocker):
|
||
label = 'PWM'
|
||
|
||
def get_value(self):
|
||
return f'{self.game.pwm.percent:3.0f} %'
|
||
|
||
def on_press_up(self, _pos):
|
||
self.game.pwm.duty_cycle = min(
|
||
1000000, self.game.pwm.duty_cycle + 100000
|
||
)
|
||
return True
|
||
|
||
def on_press_down(self, _pos):
|
||
self.game.pwm.duty_cycle = max(0, self.game.pwm.duty_cycle - 100000)
|
||
return True
|
||
|
||
|
||
class BlurRocker(Rocker):
|
||
label = 'BLUR'
|
||
|
||
def get_value(self):
|
||
return f'{self.game.detector.blur} px'
|
||
|
||
def on_press_up(self, _pos):
|
||
self.game.detector.blur += 1
|
||
return True
|
||
|
||
def on_press_down(self, _pos):
|
||
self.game.detector.blur = max(1, self.game.detector.blur - 1)
|
||
return True
|
||
|
||
|
||
class ExposureRocker(Rocker):
|
||
label = 'EXPO'
|
||
|
||
def get_value(self):
|
||
return (
|
||
'±0.0 EV'
|
||
if abs(self.game.camera.exposure) < 1e-5
|
||
else (
|
||
f'{self.game.camera.exposure:+3.1f}'
|
||
' E\N{LATIN SUBSCRIPT SMALL LETTER V}'
|
||
)
|
||
)
|
||
|
||
def on_press_up(self, _pos):
|
||
self.game.camera.exposure += 1 / 3
|
||
return True
|
||
|
||
def on_press_down(self, _pos):
|
||
self.game.camera.exposure -= 1 / 3
|
||
return True
|
||
|
||
|
||
class ResolutionRocker(Rocker):
|
||
label = 'RES'
|
||
|
||
def get_value(self):
|
||
return f'× {self.game.frame.shape[0] / self.game.screen_height:3.1f}'
|
||
|
||
def on_press_up(self, _pos):
|
||
self.game.resolution_index += 1
|
||
self.game.camera.resize(
|
||
(
|
||
self.game.screen_width * renard(self.game.resolution_index),
|
||
self.game.screen_height * renard(self.game.resolution_index),
|
||
)
|
||
)
|
||
|
||
def on_press_down(self, _pos):
|
||
self.game.resolution_index -= 1
|
||
self.game.camera.resize(
|
||
(
|
||
self.game.screen_width * renard(self.game.resolution_index),
|
||
self.game.screen_height * renard(self.game.resolution_index),
|
||
)
|
||
)
|
||
|
||
|
||
# def on_press_up(self, _pos):
|
||
# ratio = self.game.frame.shape[0] / self.game.screen_height
|
||
# ratio += 0.25
|
||
# self.game.camera.resize(
|
||
# (self.game.screen_width * ratio, self.game.screen_height * ratio))
|
||
# return True
|
||
|
||
# def on_press_down(self, _pos):
|
||
# ratio = self.game.frame.shape[0] / self.game.screen_height
|
||
# ratio -= 0.25
|
||
# self.game.camera.resize(
|
||
# (self.game.screen_width * ratio, self.game.screen_height * ratio))
|
||
# return True
|
||
|
||
|
||
class ButtonBack(Button):
|
||
label = 'BACK'
|
||
|
||
def on_press(self, _pos):
|
||
self.game.mode = self.game.mode_normal
|
||
return True
|
||
|
||
|
||
class ButtonCallibrate(Button):
|
||
label = 'ADJ'
|
||
|
||
def on_press(self, _pos):
|
||
self.game.mode = self.game.mode_callibrate
|
||
return True
|
||
|
||
|
||
class ButtonQuit(Button):
|
||
label = 'QUIT'
|
||
key = 'Q'
|
||
|
||
def on_press(self, _pos):
|
||
sys.exit()
|
||
|
||
|
||
class ButtonConnect(Button):
|
||
key = 'R'
|
||
|
||
def draw_label(self, surface):
|
||
if self.game.bot.connected:
|
||
self.label = 'CONNECTED'
|
||
elif self.game.bot.loop is not None:
|
||
self.label = 'CONNECTING'
|
||
else:
|
||
self.label = 'CONNECT'
|
||
|
||
return super().draw_label(surface)
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.bot.loop is None:
|
||
self.game.bot.connect()
|
||
return True
|
||
|
||
|
||
class ButtonSonar(Button):
|
||
label = 'SONAR'
|
||
key = 'E'
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.bot.connected:
|
||
self.game.bot.sonar_nowait()
|
||
return True
|
||
|
||
|
||
class ButtonForward(Button):
|
||
label = '↑'
|
||
key = 'W'
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.bot.connected:
|
||
self.game.bot.forward_nowait()
|
||
return True
|
||
|
||
|
||
class ButtonBackward(Button):
|
||
label = '↓'
|
||
key = 'S'
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.bot.connected:
|
||
self.game.bot.backward_nowait()
|
||
return True
|
||
|
||
|
||
class ButtonTurnLeft(Button):
|
||
label = '←'
|
||
key = 'A'
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.bot.connected:
|
||
self.game.bot.turn_left_nowait()
|
||
return True
|
||
|
||
|
||
class ButtonTurnRight(Button):
|
||
label = '→'
|
||
key = 'D'
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.bot.connected:
|
||
self.game.bot.turn_right_nowait()
|
||
return True
|
||
|
||
|
||
class ButtonCapture(Button):
|
||
label = 'CAPTURE'
|
||
key = 'SPACE'
|
||
|
||
def on_press(self, _pos):
|
||
if self.game.mode == self.game.mode_snapshot:
|
||
self.game.capture_frame()
|
||
else:
|
||
self.game.mode = self.game.mode_snapshot
|
||
return True
|
||
|
||
|
||
class ButtonRun(Button):
|
||
label = 'RUN'
|
||
key = 'R'
|
||
|
||
def on_press(self, _pos):
|
||
self.game.bot.execute_programme(self.game.current_ast)
|
||
return True
|
||
|
||
|
||
class ButtonStop(Button):
|
||
label = 'STOP'
|
||
key = 'X'
|
||
|
||
def on_press(self, _pos):
|
||
self.game.bot.stop_programme()
|
||
return True
|
||
|
||
|
||
class ButtonBadges(Button):
|
||
label = 'B'
|
||
key = 'B'
|
||
|
||
def on_press(self, _pos):
|
||
# TODO check how many badges are there, and then B switches to either
|
||
# mode_badges or just mode_badge with first one picked
|
||
self.game.current_badge = self.game.badges[0]
|
||
self.game.mode = self.game.mode_badge
|
||
|
||
|
||
# fmt: off
|
||
R20 = (
|
||
1.00, 1.12, 1.25, 1.40,
|
||
1.60, 1.80, 2.00, 2.24,
|
||
2.50, 2.80, 3.15, 3.55,
|
||
4.00, 4.50, 5.00, 5.60,
|
||
6.30, 7.10, 8.00, 9.00,
|
||
)
|
||
R10 = R20[::2]
|
||
R5 = R10[::2]
|
||
# fmt: on
|
||
|
||
|
||
def renard(index, series=R10):
|
||
magnitude, number = divmod(index, len(series))
|
||
return series[number] * 10**magnitude
|
||
|
||
|
||
class Game:
|
||
def __init__(self, screen, bot):
|
||
self.screen = screen
|
||
self.screen_width, self.screen_height = self.screen.get_size()
|
||
self.layout = Layout(self.screen)
|
||
|
||
self.clock = pygame.time.Clock()
|
||
|
||
self.mode = self.mode_normal
|
||
|
||
self.bot = bot
|
||
|
||
self.pwm = utils.PWM()
|
||
self.pwm.period = 1000000 # ns -> 1 kHz
|
||
self.pwm.duty_cycle = 0
|
||
self.pwm.enable = 1
|
||
|
||
self.cpuload = utils.CPULoad()
|
||
self.cpuload.update()
|
||
pygame.time.set_timer(PYGAME_EVENT_UPDATE_CPU_LOAD, 1000)
|
||
|
||
self.resolution_index = 2
|
||
self.camera = camera.Camera(
|
||
(
|
||
self.screen_width * renard(self.resolution_index),
|
||
self.screen_height * renard(self.resolution_index),
|
||
)
|
||
)
|
||
self.detector = _detector.Detector()
|
||
self.frame = None
|
||
self.metadata = None
|
||
|
||
self.current_programme = None
|
||
self.current_ast = None
|
||
|
||
self.badges = list(badges.Badge.load_from_assets())
|
||
self.current_badge = None
|
||
|
||
self.error = False
|
||
|
||
self.ctrl_pwm = PWMRocker(
|
||
self, **self.layout.place_button(left=1, bottom=0, height=2)
|
||
)
|
||
self.ctrl_blur = BlurRocker(
|
||
self, **self.layout.place_button(left=2, bottom=0, height=2)
|
||
)
|
||
self.ctrl_expo = ExposureRocker(
|
||
self, **self.layout.place_button(left=3, bottom=0, height=2)
|
||
)
|
||
self.ctrl_res = ResolutionRocker(
|
||
self, **self.layout.place_button(left=4, bottom=0, height=2)
|
||
)
|
||
|
||
self.ctrl_button_back = ButtonBack(
|
||
self, **self.layout.place_button(left=0, bottom=0)
|
||
)
|
||
|
||
self.ctrl_button_callibrate = ButtonCallibrate(
|
||
self, **self.layout.place_button(left=0, bottom=1)
|
||
)
|
||
self.ctrl_button_quit = ButtonQuit(
|
||
self, **self.layout.place_button(left=0, bottom=0)
|
||
)
|
||
self.ctrl_button_capture = ButtonCapture(
|
||
self, **self.layout.place_button(left=4, bottom=0, width=2)
|
||
)
|
||
self.ctrl_button_badges = ButtonBadges(
|
||
self, **self.layout.place_button(left=1, bottom=1)
|
||
)
|
||
|
||
self.ctrl_button_connect = ButtonConnect(
|
||
self, **self.layout.place_button(left=1, bottom=0, width=3)
|
||
)
|
||
self.ctrl_button_turn_left = ButtonTurnLeft(
|
||
self, **self.layout.place_button(left=1, bottom=0)
|
||
)
|
||
self.ctrl_button_forward = ButtonForward(
|
||
self, **self.layout.place_button(left=2, bottom=1)
|
||
)
|
||
self.ctrl_button_backward = ButtonBackward(
|
||
self, **self.layout.place_button(left=2, bottom=0)
|
||
)
|
||
self.ctrl_button_turn_right = ButtonTurnRight(
|
||
self, **self.layout.place_button(left=3, bottom=0)
|
||
)
|
||
self.ctrl_button_sonar = ButtonSonar(
|
||
self, **self.layout.place_button(left=3, bottom=1)
|
||
)
|
||
|
||
self.ctrl_button_run = ButtonRun(
|
||
self, **self.layout.place_button(left=4, bottom=1, width=2)
|
||
)
|
||
self.ctrl_button_stop = ButtonStop(
|
||
self,
|
||
**self.layout.place_button(left=0, bottom=0, width=6, height=2),
|
||
)
|
||
|
||
self.ctrl_status = Status(self)
|
||
self.ctrl_debug = Debug(self)
|
||
self.ctrl_sonar = Sonar(self)
|
||
|
||
def run(self):
|
||
with self.camera:
|
||
try:
|
||
while True:
|
||
self.mode()
|
||
pygame.display.flip()
|
||
self.clock.tick(50)
|
||
|
||
except Exception as exc:
|
||
# TODO make this configurable
|
||
if self.frame is not None:
|
||
now = f'{datetime.datetime.now():%Y%m%d-%H%M%S}'
|
||
cv2.imwrite(f'badframe-{now}.png', self.frame)
|
||
with open(
|
||
f'badframe-{now}.txt', 'w', encoding='utf-8'
|
||
) as file:
|
||
traceback.print_exception(exc, file=file)
|
||
raise
|
||
|
||
finally:
|
||
if self.bot.loop is not None:
|
||
self.bot.disconnect()
|
||
self.pwm.enable = 0
|
||
|
||
def capture_frame(self):
|
||
# NOTE about frames, arrays and shapes:
|
||
# picamera's make_array returns array of shape (height, width, 3)
|
||
# array[0][0] is top left, but the camera is mounted "upside down" in
|
||
# the box I'm using, therefore rotate 180°
|
||
self.frame, self.metadata = self.camera.capture_frame(rotate=2)
|
||
|
||
def handle_event(self, ui, event):
|
||
if event.type == pygame.QUIT or (
|
||
event.type == pygame.KEYDOWN and event.key == pygame.K_q
|
||
):
|
||
sys.exit()
|
||
|
||
if event.type == PYGAME_EVENT_UPDATE_CPU_LOAD:
|
||
self.cpuload.update()
|
||
return True
|
||
|
||
return any(ctrl.handle_event(event) for ctrl in ui)
|
||
|
||
def draw_ui(self, ui):
|
||
self.screen.blits(
|
||
list(itertools.chain.from_iterable(ctrl.get_blits() for ctrl in ui))
|
||
)
|
||
|
||
def compile_display(self):
|
||
display_frame = self.frame.copy()
|
||
|
||
detections = self.detector.detect_markers(self.frame)
|
||
analyser = _detector.Analyser(detections)
|
||
for detection in detections:
|
||
detection.debug_draw(display_frame)
|
||
try:
|
||
self.current_programme = list(
|
||
analyser.get_programme(debug_frame=display_frame)
|
||
)
|
||
self.current_ast = parser.parse(self.current_programme)
|
||
except (_detector.BoardError, funcparserlib.parser.NoParseError):
|
||
self.current_programme = None
|
||
self.current_ast = None
|
||
self.error = True
|
||
else:
|
||
self.error = False
|
||
|
||
# pygame's surface needs shape (width, height, 3), hence transpose
|
||
surface = pygame.pixelcopy.make_surface(
|
||
display_frame.transpose(1, 0, 2)
|
||
)
|
||
surface = pygame.transform.scale(surface, self.screen.get_size())
|
||
frame_width, frame_height = surface.get_size()
|
||
|
||
self.screen.blit(
|
||
surface,
|
||
(
|
||
(self.screen_width - frame_width) / 2,
|
||
(self.screen_height - frame_height) / 2,
|
||
),
|
||
)
|
||
|
||
def mode_normal(self):
|
||
ui = (
|
||
self.ctrl_status,
|
||
self.ctrl_button_callibrate,
|
||
self.ctrl_button_quit,
|
||
self.ctrl_button_badges,
|
||
self.ctrl_button_capture,
|
||
) + (
|
||
(
|
||
self.ctrl_button_forward,
|
||
self.ctrl_button_backward,
|
||
self.ctrl_button_turn_left,
|
||
self.ctrl_button_turn_right,
|
||
self.ctrl_button_sonar,
|
||
self.ctrl_sonar,
|
||
)
|
||
if self.bot.connected
|
||
else (self.ctrl_button_connect,)
|
||
)
|
||
|
||
self.capture_frame()
|
||
self.compile_display()
|
||
|
||
self.draw_ui(ui)
|
||
for event in pygame.event.get():
|
||
self.handle_event(ui, event)
|
||
|
||
def mode_snapshot(self):
|
||
ui = [
|
||
self.ctrl_status,
|
||
]
|
||
|
||
if self.bot.connected:
|
||
ui.append(self.ctrl_sonar)
|
||
|
||
if self.bot.current_task is not None:
|
||
ui.append(self.ctrl_button_stop)
|
||
else:
|
||
ui.extend(
|
||
(
|
||
self.ctrl_button_forward,
|
||
self.ctrl_button_backward,
|
||
self.ctrl_button_turn_left,
|
||
self.ctrl_button_turn_right,
|
||
self.ctrl_button_sonar,
|
||
self.ctrl_button_capture,
|
||
self.ctrl_button_badges,
|
||
self.ctrl_button_back,
|
||
)
|
||
)
|
||
if self.current_programme:
|
||
ui.append(self.ctrl_button_run)
|
||
else:
|
||
ui.extend(
|
||
(
|
||
self.ctrl_button_connect,
|
||
self.ctrl_button_capture,
|
||
self.ctrl_button_badges,
|
||
self.ctrl_button_back,
|
||
)
|
||
)
|
||
|
||
self.compile_display()
|
||
|
||
if self.current_programme is not None:
|
||
for i, line in enumerate(self.current_programme, 1):
|
||
line = ' '.join(token.value.text for token in line)
|
||
surface = self.layout.font_mono.render(
|
||
f'{i:2d} {line}', True, Colour.HIGHLIGHT, Colour.BACKGROUND
|
||
)
|
||
self.screen.blit(surface, (10, surface.get_height() * i))
|
||
|
||
self.draw_ui(ui)
|
||
for event in pygame.event.get():
|
||
self.handle_event(ui, event)
|
||
|
||
def mode_callibrate(self):
|
||
ui = (
|
||
self.ctrl_status,
|
||
self.ctrl_debug,
|
||
self.ctrl_button_back,
|
||
self.ctrl_pwm,
|
||
self.ctrl_blur,
|
||
self.ctrl_expo,
|
||
self.ctrl_res,
|
||
)
|
||
|
||
self.capture_frame()
|
||
self.compile_display()
|
||
|
||
self.draw_ui(ui)
|
||
for event in pygame.event.get():
|
||
self.handle_event(ui, event)
|
||
|
||
def mode_badge(self):
|
||
ui = (
|
||
self.ctrl_button_back,
|
||
) # fmt: skip
|
||
|
||
assert self.current_badge is not None
|
||
|
||
padding = 48
|
||
|
||
self.screen.fill((0, 0, 0))
|
||
|
||
qrcode_size = (
|
||
self.screen_height
|
||
- 3 * padding
|
||
- self.current_badge.image.get_height()
|
||
)
|
||
qrcode = pygame.transform.scale(
|
||
self.current_badge.qrcode, (qrcode_size, qrcode_size)
|
||
)
|
||
|
||
self.screen.blit(
|
||
qrcode, ((self.screen_width - qrcode_size) / 2, padding)
|
||
)
|
||
self.screen.blit(
|
||
self.current_badge.image,
|
||
(
|
||
(self.screen_width - self.current_badge.image.get_width()) / 2,
|
||
(
|
||
self.screen_height
|
||
- self.current_badge.image.get_height()
|
||
- padding
|
||
),
|
||
),
|
||
)
|
||
|
||
self.draw_ui(ui)
|
||
|
||
for event in pygame.event.get():
|
||
self.handle_event(ui, event)
|
||
|
||
|
||
# vim: tw=80 ts=4 sts=4 sw=4 et
|