diff --git a/scripts/example_bot.py b/scripts/example_bot.py new file mode 100644 index 0000000..c331a82 --- /dev/null +++ b/scripts/example_bot.py @@ -0,0 +1,40 @@ +from PyQt6.QtWidgets import QApplication + +from src.bot import Bot +from src.command import CarControl +from src.recorder import RecorderWindow +from src.snapshot import Snapshot + + +class ExampleBot(Bot): + def nn_infer(self, snapshot: Snapshot) -> list[tuple[CarControl, bool]]: + # Do smart NN inference here + return [(CarControl.FORWARD, True)] + + def on_snapshot_received(self, snapshot: Snapshot): + controls: list[tuple[CarControl, bool]] = self.nn_infer(snapshot) + for control, active in controls: + self.recorder.on_car_controlled(control, active) + + +def main(): + import sys + + def except_hook(cls, exception, traceback): + sys.__excepthook__(cls, exception, traceback) + + sys.excepthook = except_hook + + app: QApplication = QApplication(sys.argv) + recorder: RecorderWindow = RecorderWindow("localhost", 5000) + bot: ExampleBot = ExampleBot(recorder) + + app.aboutToQuit.connect(recorder.shutdown) + recorder.register_bot(bot) + recorder.show() + + app.exec() + + +if __name__ == "__main__": + main() diff --git a/src/bot.py b/src/bot.py new file mode 100644 index 0000000..85046c5 --- /dev/null +++ b/src/bot.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from src.snapshot import Snapshot + +if TYPE_CHECKING: + from src.recorder import RecorderWindow + + +class Bot: + def __init__(self, recorder: RecorderWindow): + self.recorder: RecorderWindow = recorder + + def on_snapshot_received(self, snapshot: Snapshot): + pass diff --git a/src/recorder.py b/src/recorder.py index 6392615..9c58c8e 100644 --- a/src/recorder.py +++ b/src/recorder.py @@ -9,6 +9,7 @@ from PyQt6.QtCore import QObject, QThread, QTimer, pyqtSignal, pyqtSlot from PyQt6.QtGui import QKeyEvent from PyQt6.QtWidgets import QMainWindow +from src.bot import Bot from src.command import ApplySnapshotCommand, CarControl, Command, ControlCommand, RecordingCommand, ResetCommand from src.record_file import RecordFile from src.recorder_ui import Ui_Recorder @@ -162,9 +163,11 @@ class RecorderWindow(Ui_Recorder, QMainWindow): self.recordDataButton.clicked.connect(self.toggle_record) self.resetButton.clicked.connect(self.rollback) + self.bot: Optional[Bot] = None self.autopiloting = False self.autopilotButton.clicked.connect(self.toggle_autopilot) + self.autopilotButton.setDisabled(True) self.saveRecordButton.clicked.connect(self.save_record) @@ -263,8 +266,15 @@ class RecorderWindow(Ui_Recorder, QMainWindow): self.snapshots.append(snapshot) self.nbrSnapshotSaved.setText(str(len(self.snapshots))) + if self.autopiloting and self.bot is not None: + self.bot.on_snapshot_received(snapshot) + def shutdown(self): self.close_signal.emit() def send_command(self, command: Command): self.send_signal.emit(command) + + def register_bot(self, bot: Bot): + self.bot = bot + self.autopilotButton.setDisabled(False)