diff --git a/settings.toml b/settings.toml
index d9e17b74008c3c753a63836210baa78f9555fa0d..2410c6a1dce07f9c1892bc237880c37600322cb1 100644
--- a/settings.toml
+++ b/settings.toml
@@ -1,11 +1,14 @@
[default]
CAMERA = 'RASPI_V2'
-OBJECTS_DETECTION_MODEL = "ssd_mobilenet_v2_roco_2018_03_29_20200314_015411_TWITCH_TEMP_733_IMGS_29595steps"
+OBJECTS_DETECTION_MODEL = "ssd_mobilenet_v2__210224_042137__Twitch2_Train_T470149066_T470150052_T470152289_T470153081_T470158483_1775_imgs__25297_steps"
ARMOR_DIGIT_MODEL = "20210117_145856_kd_cnn/distiled - temp 4.1e+01 - cnn - (32) - lr 7.8e-04 - drop 63.pkl"
-SERIAL_PORT = "/dev/ttyTHS0"
-BAUD_RATE = 230_400
+
+USE_UART = false
[development]
[production]
+SERIAL_PORT = "/dev/ttyTHS0"
+BAUD_RATE = 230_400
+USE_UART = true
\ No newline at end of file
diff --git a/src/polystar/communication/command.py b/src/polystar/communication/command.py
index f906f84b6b3cdd2f16d3a0336581d67ef7bd2eb7..a494ddd4d779c4f5f09d8b41094d174d07a29a33 100644
--- a/src/polystar/communication/command.py
+++ b/src/polystar/communication/command.py
@@ -15,3 +15,6 @@ class Command:
def make_target_command(target: SimpleTarget) -> Command:
return Command(id=2, data=b"Y" + bytes(target))
+
+
+NO_TARGET_COMMAND = Command(id=2, data=b"N\x00\x00\x00\x00\x00\x00")
diff --git a/src/polystar/communication/cs_link_abc.py b/src/polystar/communication/cs_link_abc.py
index fd31cf48e0de988fcef5850c14d921d051dea8a2..ad299ea0492c4d2804afd41a03780b6cfa3a77bc 100644
--- a/src/polystar/communication/cs_link_abc.py
+++ b/src/polystar/communication/cs_link_abc.py
@@ -1,7 +1,8 @@
from abc import ABC, abstractmethod
from typing import Iterable
-from polystar.communication.command import Command
+from polystar.communication.command import NO_TARGET_COMMAND, Command, make_target_command
+from polystar.target_pipeline.target_abc import SimpleTarget
class CSLinkABC(ABC):
@@ -11,3 +12,9 @@ class CSLinkABC(ABC):
@abstractmethod
def send_command(self, command: Command):
pass
+
+ def send_target(self, target: SimpleTarget):
+ self.send_command(make_target_command(target))
+
+ def send_no_target(self):
+ self.send_command(NO_TARGET_COMMAND)
diff --git a/src/polystar/communication/togglabe_cs_link.py b/src/polystar/communication/togglabe_cs_link.py
new file mode 100644
index 0000000000000000000000000000000000000000..2cc7a6d969919ddad2562f8480b8fd1f1e48a7ba
--- /dev/null
+++ b/src/polystar/communication/togglabe_cs_link.py
@@ -0,0 +1,16 @@
+from polystar.communication.command import Command
+from polystar.communication.cs_link_abc import CSLinkABC
+
+
+class TogglableCSLink(CSLinkABC):
+ def __init__(self, cs_link: CSLinkABC, is_on: bool):
+ self.is_on = is_on
+ self.cs_link = cs_link
+
+ def send_command(self, command: Command):
+ if not self.is_on:
+ return
+ self.cs_link.send_command(command)
+
+ def toggle(self):
+ self.is_on = not self.is_on
diff --git a/src/polystar/dependency_injection.py b/src/polystar/dependency_injection.py
index ab8bfdf374d9f88b20e76bb86d257b664b8860ea..6cb199adc8d5c980233015f2b9bbedc99e6ca0fe 100644
--- a/src/polystar/dependency_injection.py
+++ b/src/polystar/dependency_injection.py
@@ -22,7 +22,10 @@ from polystar.target_pipeline.detected_objects.detected_objects_factory import D
from polystar.target_pipeline.object_selectors.closest_object_selector import ClosestObjectSelector
from polystar.target_pipeline.object_selectors.object_selector_abc import ObjectSelectorABC
from polystar.target_pipeline.objects_detectors.objects_detector_abc import ObjectsDetectorABC
-from polystar.target_pipeline.objects_filters.confidence_object_filter import ConfidenceObjectsFilter
+from polystar.target_pipeline.objects_filters.confidence_object_filter import (
+ ConfidenceObjectsFilter,
+ RobotArmorConfidenceObjectsFilter,
+)
from polystar.target_pipeline.objects_filters.objects_filter_abc import ObjectsFilterABC
from polystar.target_pipeline.objects_linker.objects_linker_abs import ObjectsLinkerABC
from polystar.target_pipeline.objects_linker.simple_objects_linker import SimpleObjectsLinker
@@ -82,7 +85,7 @@ class CommonModule(Module):
@multiprovider
@singleton
def provide_objects_validators(self) -> List[ObjectsFilterABC]:
- return [ConfidenceObjectsFilter(0.6)]
+ return [RobotArmorConfidenceObjectsFilter(0.5)]
@provider
@singleton
@@ -97,9 +100,9 @@ class CommonModule(Module):
@provider
@singleton
def provide_cs_link(self) -> CSLinkABC:
- if self.settings.is_dev:
+ if not self.settings.USE_UART:
return Screen()
- return BoardA(Serial(settings.SERIAL_PORT))
+ return BoardA(Serial(settings.SERIAL_PORT, settings.BAUD_RATE))
@provider
@singleton
diff --git a/src/polystar/target_pipeline/detected_objects/detected_robot.py b/src/polystar/target_pipeline/detected_objects/detected_robot.py
index 97f72f46611f6477ff02c4c129d2812acc0ea607..13f5c363a136084ea3587cff3994a1c3a617f006 100644
--- a/src/polystar/target_pipeline/detected_objects/detected_robot.py
+++ b/src/polystar/target_pipeline/detected_objects/detected_robot.py
@@ -13,4 +13,4 @@ class DetectedRobot(DetectedROCOObject):
class FakeDetectedRobot(DetectedRobot):
def __init__(self, armor: DetectedArmor):
- super().__init__(type=ObjectType.CAR, box=armor.box, confidence=0, armors=[armor])
+ super().__init__(type=ObjectType.CAR, box=armor.box, confidence=armor.confidence, armors=[armor])
diff --git a/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py b/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py
index 1c406ce70c6cb904f8e7045936d202a8ef129179..4e2c0f42ecdd39601a217c6ef208f04057cc2866 100644
--- a/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py
+++ b/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py
@@ -42,7 +42,7 @@ class TRTModelObjectsDetector(ObjectsDetectorABC):
object_class_id=int(object_class_id),
)
for (_, object_class_id, score, xmin, ymin, xmax, ymax) in results
- if object_class_id >= 0
+ if object_class_id >= 0 and object_class_id == 4 # FIXME : remove me
],
image,
)
diff --git a/src/polystar/target_pipeline/objects_filters/confidence_object_filter.py b/src/polystar/target_pipeline/objects_filters/confidence_object_filter.py
index 075b5c2fc4ee4c71cbaae5aaef9f257fc6d2bb94..dc1bf9c609c3559af1dee49181ddccb0931064b0 100644
--- a/src/polystar/target_pipeline/objects_filters/confidence_object_filter.py
+++ b/src/polystar/target_pipeline/objects_filters/confidence_object_filter.py
@@ -1,5 +1,6 @@
from polystar.filters.filter_abc import FilterABC
from polystar.target_pipeline.detected_objects.detected_object import DetectedROCOObject
+from polystar.target_pipeline.detected_objects.detected_robot import DetectedRobot
class ConfidenceObjectsFilter(FilterABC[DetectedROCOObject]):
@@ -10,3 +11,11 @@ class ConfidenceObjectsFilter(FilterABC[DetectedROCOObject]):
def validate_single(self, obj: DetectedROCOObject) -> bool:
return obj.confidence >= self.confidence_threshold
+
+
+class RobotArmorConfidenceObjectsFilter(FilterABC[DetectedRobot]):
+ def __init__(self, armor_confidence_threshold: float):
+ self.armor_confidence_threshold = armor_confidence_threshold
+
+ def validate_single(self, robot: DetectedRobot) -> bool:
+ return any(armor.confidence >= self.armor_confidence_threshold for armor in robot.armors)
diff --git a/src/polystar/view/cv2_results_viewer.py b/src/polystar/view/cv2_results_viewer.py
index ef93a3d092a97346eba85c9ecb31eaea53f50b35..98594c241c1c625fdb25e9ecc3a0aa360f4e1fc2 100644
--- a/src/polystar/view/cv2_results_viewer.py
+++ b/src/polystar/view/cv2_results_viewer.py
@@ -1,3 +1,6 @@
+from contextlib import suppress
+from typing import Any, Callable, Dict
+
import cv2
import numpy as np
@@ -23,10 +26,12 @@ COLORS = [
(23, 190, 207),
] # seaborn.color_palette() * 255
+Callback = Callable[[], Any]
+
class CV2ResultViewer(ResultViewerABC):
- def __init__(self, name: str, delay: int = 1, end_keys: str = "q"):
- self.end_keys = [ord(c) for c in end_keys]
+ def __init__(self, name: str, delay: int = 1, end_key: str = "q", key_callbacks: Dict[str, Callback] = None):
+ self.keycode_callbacks = self._make_keycode_callbacks(end_key, key_callbacks or {})
self.delay = delay
self.name = name
self._current_image: Image = None
@@ -70,4 +75,13 @@ class CV2ResultViewer(ResultViewerABC):
def display(self):
cv2.imshow(self.name, self._current_image)
- self.finished = cv2.waitKey(self.delay) & 0xFF in self.end_keys
+ keycode = cv2.waitKey(self.delay) & 0xFF
+ with suppress(KeyError):
+ self.keycode_callbacks[keycode]()
+
+ def stop(self):
+ self.finished = True
+
+ def _make_keycode_callbacks(self, end_key: str, key_callbacks: Dict[str, Callback]) -> Dict[int, Callback]:
+ key_callbacks[end_key] = self.stop
+ return {ord(k): f for k, f in key_callbacks.items()}
diff --git a/src/research/armors/demos/__init__.py b/src/research/armors/demos/__init__.py
deleted file mode 100644
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000
diff --git a/src/research/armors/demos/demo_pipeline_camera.py b/src/research/armors/demos/demo_pipeline_camera.py
deleted file mode 100644
index 113e5e41c7da3dee0397baed6fc68b6848f0a76a..0000000000000000000000000000000000000000
--- a/src/research/armors/demos/demo_pipeline_camera.py
+++ /dev/null
@@ -1,31 +0,0 @@
-from injector import inject
-
-from polystar.dependency_injection import make_injector
-from polystar.frame_generators.frames_generator_abc import FrameGeneratorABC
-from polystar.target_pipeline.debug_pipeline import DebugTargetPipeline
-from polystar.target_pipeline.target_pipeline import NoTargetFoundException
-from polystar.utils.fps import FPS
-from polystar.view.cv2_results_viewer import CV2ResultViewer
-
-
-@inject
-def demo_pipeline_on_camera(pipeline: DebugTargetPipeline, webcam: FrameGeneratorABC):
- fps, pipeline_fps = FPS(), FPS()
- with CV2ResultViewer("TensorRT demo") as viewer:
- for image in webcam.generate():
- pipeline_fps.skip()
- try:
- pipeline.predict_target(image)
- except NoTargetFoundException:
- pass
- pipeline_fps.tick(), fps.tick()
- viewer.add_debug_info(pipeline.debug_info_)
- viewer.add_text(f"FPS: {fps:.1f} / {pipeline_fps:.1f}", 10, 10, (0, 0, 0))
- viewer.display()
- fps.skip()
- if viewer.finished:
- return
-
-
-if __name__ == "__main__":
- make_injector().call_with_injection(demo_pipeline_on_camera)
diff --git a/src/research/armors/demos/demo_pipeline.py b/src/research/scripts/demo_pipeline.py
similarity index 100%
rename from src/research/armors/demos/demo_pipeline.py
rename to src/research/scripts/demo_pipeline.py
diff --git a/src/research/scripts/demo_pipeline_camera.py b/src/research/scripts/demo_pipeline_camera.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2c79fd0c205ecca1e929a9ca05d54e4a9e8da9f
--- /dev/null
+++ b/src/research/scripts/demo_pipeline_camera.py
@@ -0,0 +1,53 @@
+from injector import inject
+
+from polystar.communication.cs_link_abc import CSLinkABC
+from polystar.communication.togglabe_cs_link import TogglableCSLink
+from polystar.dependency_injection import make_injector
+from polystar.frame_generators.frames_generator_abc import FrameGeneratorABC
+from polystar.models.image import Image
+from polystar.target_pipeline.debug_pipeline import DebugTargetPipeline
+from polystar.target_pipeline.target_pipeline import NoTargetFoundException
+from polystar.utils.fps import FPS
+from polystar.view.cv2_results_viewer import CV2ResultViewer
+
+
+class CameraPipelineDemo:
+ @inject
+ def __init__(self, pipeline: DebugTargetPipeline, webcam: FrameGeneratorABC, cs_link: CSLinkABC):
+ self.cs_link = TogglableCSLink(cs_link, is_on=False)
+ self.webcam = webcam
+ self.pipeline = pipeline
+ self.fps, self.pipeline_fps = FPS(), FPS()
+ self.persistence_last_detection = 0
+
+ def run(self):
+ with CV2ResultViewer("TensorRT demo", key_callbacks={" ": self.cs_link.toggle}) as viewer:
+ for image in self.webcam.generate():
+ self.pipeline_fps.skip()
+ self._detect(image)
+ self.pipeline_fps.tick(), self.fps.tick()
+ self._display(viewer)
+ self.fps.skip()
+ if viewer.finished:
+ return
+
+ def _detect(self, image: Image):
+ try:
+ target = self.pipeline.predict_target(image)
+ self.cs_link.send_target(target)
+ self.persistence_last_detection = 5
+ except NoTargetFoundException:
+ if self.persistence_last_detection:
+ self.persistence_last_detection -= 1
+ else:
+ self.cs_link.send_no_target()
+
+ def _display(self, viewer: CV2ResultViewer):
+ viewer.add_debug_info(self.pipeline.debug_info_)
+ viewer.add_text(f"FPS: {self.fps:.1f} / {self.pipeline_fps:.1f}", 10, 10, (0, 0, 0))
+ viewer.add_text("Communication: " + ("[ON]" if self.cs_link.is_on else "[OFF]"), 10, 30, (0, 0, 0))
+ viewer.display()
+
+
+if __name__ == "__main__":
+ make_injector().get(CameraPipelineDemo).run()