diff --git a/src/polystar/dependency_injection.py b/src/polystar/dependency_injection.py
index 6cb199adc8d5c980233015f2b9bbedc99e6ca0fe..26061733bfb4a349806bc6e616836cb4f25b706a 100644
--- a/src/polystar/dependency_injection.py
+++ b/src/polystar/dependency_injection.py
@@ -9,7 +9,7 @@ from polystar.communication.board_a import BoardA
from polystar.communication.cs_link_abc import CSLinkABC
from polystar.communication.screen import Screen
from polystar.constants import LABEL_MAP_PATH
-from polystar.frame_generators.camera_frame_generator import RaspiV2CameraFrameGenerator, WebcamFrameGenerator
+from polystar.frame_generators.camera_frame_generator import CameraFrameGenerator, make_csi_camera_frame_generator
from polystar.frame_generators.frames_generator_abc import FrameGeneratorABC
from polystar.models.camera import Camera
from polystar.models.label_map import LabelMap
@@ -22,10 +22,7 @@ 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,
- RobotArmorConfidenceObjectsFilter,
-)
+from polystar.target_pipeline.objects_filters.confidence_object_filter import 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
@@ -78,7 +75,7 @@ class CommonModule(Module):
@singleton
def provide_armor_descriptors(self) -> List[ArmorsDescriptorABC]:
return [
- ArmorsColorDescriptor(ArmorColorPipeline.from_pipes([MeanChannels(), RedBlueComparisonClassifier()])),
+ # ArmorsColorDescriptor(ArmorColorPipeline.from_pipes([MeanChannels(), RedBlueComparisonClassifier()])),
# ArmorsDigitDescriptor(pkl_load(PIPELINES_DIR / "armor-digit" / settings.ARMOR_DIGIT_MODEL)),
]
@@ -110,8 +107,7 @@ class CommonModule(Module):
return SimpleObjectsLinker(min_percentage_intersection=0.8)
@provider
- @singleton
def provide_webcam(self) -> FrameGeneratorABC:
if self.settings.is_prod:
- return RaspiV2CameraFrameGenerator(1_280, 720)
- return WebcamFrameGenerator()
+ return make_csi_camera_frame_generator(1_280, 720)
+ return CameraFrameGenerator()
diff --git a/src/polystar/frame_generators/camera_frame_generator.py b/src/polystar/frame_generators/camera_frame_generator.py
index 237898bc5d0117d98df1907491d53af47a19d720..c2649fc2ded68134c4f5e11f9b39ed23d301b964 100644
--- a/src/polystar/frame_generators/camera_frame_generator.py
+++ b/src/polystar/frame_generators/camera_frame_generator.py
@@ -1,31 +1,56 @@
-from dataclasses import dataclass
-from typing import Any, Iterable
+from threading import Thread
+from typing import Any, Iterator
import cv2
-from polystar.frame_generators.cv2_frame_generator_abc import CV2FrameGeneratorABC
+from polystar.frame_generators.cv2_frame_generator_abc import CV2FrameGenerator
+from polystar.models.image import Image
-@dataclass
-class RaspiV2CameraFrameGenerator(CV2FrameGeneratorABC):
- width: int
- height: int
+class CameraFrameGenerator(CV2FrameGenerator):
+ def __init__(self, *capture_params: Any):
+ super().__init__(*capture_params)
+ self.camera_thread = CameraThread(super().__iter__())
- def _capture_params(self) -> Iterable[Any]:
- return (
- "nvarguscamerasrc ! "
- "video/x-raw(memory:NVMM), "
- f"width=(int){self.width}, height=(int){self.height}, "
- "format=(string)NV12, framerate=(fraction)60/1 ! "
- "nvvidconv flip-method=0 ! "
- f"video/x-raw, width=(int){self.width}, height=(int){self.height}, "
- "format=(string)BGRx ! "
- "videoconvert ! appsink",
- cv2.CAP_GSTREAMER,
- )
+ def __iter__(self):
+ self.camera_thread.start()
+ while True:
+ yield self.camera_thread.current_frame.copy()
+ def __del__(self):
+ self.camera_thread.stop()
-@dataclass
-class WebcamFrameGenerator(CV2FrameGeneratorABC):
- def _capture_params(self) -> Iterable[Any]:
- return (0,)
+
+class CameraThread(Thread):
+ def __init__(self, it: Iterator[Image]):
+ super().__init__()
+ self.it = it
+ self.running = True
+ self._get_next_frame()
+
+ def run(self):
+ while self.running:
+ self._get_next_frame()
+
+ def stop(self):
+ self.running = False
+
+ def _get_next_frame(self):
+ try:
+ self.current_frame = next(self.it)
+ except StopIteration:
+ self.running = False
+
+
+def make_csi_camera_frame_generator(width: int, height: int) -> CameraFrameGenerator:
+ return CameraFrameGenerator(
+ "nvarguscamerasrc ! "
+ "video/x-raw(memory:NVMM), "
+ f"width=(int){width}, height=(int){height}, "
+ "format=(string)NV12, framerate=60/1 ! "
+ "nvvidconv flip-method=0 ! "
+ f"video/x-raw, width=(int){width}, height=(int){height}, "
+ "format=(string)BGRx ! "
+ "videoconvert ! appsink drop=true sync=false",
+ cv2.CAP_GSTREAMER,
+ )
diff --git a/src/polystar/frame_generators/cv2_frame_generator_abc.py b/src/polystar/frame_generators/cv2_frame_generator_abc.py
index 6f6033988c5a81a6b0cd742be391cf7ece64e8f5..f191ea187582e6b267180a229d6c19dce298fb96 100644
--- a/src/polystar/frame_generators/cv2_frame_generator_abc.py
+++ b/src/polystar/frame_generators/cv2_frame_generator_abc.py
@@ -1,38 +1,33 @@
-from abc import ABC, abstractmethod
-from dataclasses import dataclass, field
-from typing import Any, Iterable
+from typing import Any, Iterable, Iterator
-import cv2
+from cv2 import CAP_PROP_BUFFERSIZE, VideoCapture
from polystar.frame_generators.frames_generator_abc import FrameGeneratorABC
from polystar.models.image import Image
-@dataclass
-class CV2FrameGeneratorABC(FrameGeneratorABC, ABC):
+class CV2FrameGenerator(FrameGeneratorABC):
+ def __init__(self, *capture_params: Any):
+ self.capture_params = capture_params or (0,)
- _cap: cv2.VideoCapture = field(init=False, repr=False)
+ def __iter__(self) -> Iterator[Image]:
+ return CV2Capture(self.capture_params)
- def __enter__(self):
- self._cap = cv2.VideoCapture(*self._capture_params())
- # self._cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
+
+class CV2Capture(Iterator[Image]):
+ def __init__(self, capture_params: Iterable):
+ self._cap = VideoCapture(*capture_params)
assert self._cap.isOpened()
- self._post_opening_operation()
+ self._cap.set(CAP_PROP_BUFFERSIZE, 0)
- def __exit__(self, exc_type, exc_val, exc_tb):
- self._cap.release()
+ def __next__(self) -> Image:
+ success, frame = self._cap.read()
- def generate(self) -> Iterable[Image]:
- with self:
- while 1:
- is_open, frame = self._cap.read()
- if not is_open:
- return
- yield frame
+ if success:
+ return frame
- @abstractmethod
- def _capture_params(self) -> Iterable[Any]:
- pass
+ raise StopIteration()
- def _post_opening_operation(self):
- pass
+ def __del__(self):
+ if self._cap.isOpened():
+ self._cap.release()
diff --git a/src/polystar/frame_generators/fps_video_frame_generator.py b/src/polystar/frame_generators/fps_video_frame_generator.py
index 7e45e0434a9028cfced229a34e8a5e805c3de996..83889e41597f2be4422c2ef32de01c77714ce750 100644
--- a/src/polystar/frame_generators/fps_video_frame_generator.py
+++ b/src/polystar/frame_generators/fps_video_frame_generator.py
@@ -1,19 +1,16 @@
-from dataclasses import dataclass
-from typing import Iterable
+from pathlib import Path
+from typing import Iterator, Optional
from polystar.frame_generators.video_frame_generator import VideoFrameGenerator
from polystar.models.image import Image
-@dataclass
class FPSVideoFrameGenerator(VideoFrameGenerator):
+ def __init__(self, video_path: Path, desired_fps: int, offset_seconds: Optional[int] = None):
+ super().__init__(video_path, offset_seconds)
+ self.frame_rate: int = self._video_fps // desired_fps
- desired_fps: int
-
- def __post_init__(self):
- self.frame_rate: int = self._video_fps // self.desired_fps
-
- def generate(self) -> Iterable[Image]:
- for i, frame in enumerate(super().generate(), -1):
+ def __iter__(self) -> Iterator[Image]:
+ for i, frame in enumerate(super().__iter__(), -1):
if not i % self.frame_rate:
yield frame
diff --git a/src/polystar/frame_generators/frames_generator_abc.py b/src/polystar/frame_generators/frames_generator_abc.py
index a6ec18af9dea2348be10230a0590b3e415853ff9..5fe8e1160eeea030a4f1c9ed096960c128ff1b7e 100644
--- a/src/polystar/frame_generators/frames_generator_abc.py
+++ b/src/polystar/frame_generators/frames_generator_abc.py
@@ -1,10 +1,10 @@
from abc import ABC, abstractmethod
-from typing import Iterable
+from typing import Iterable, Iterator
from polystar.models.image import Image
-class FrameGeneratorABC(ABC):
+class FrameGeneratorABC(ABC, Iterable[Image]):
@abstractmethod
- def generate(self) -> Iterable[Image]:
+ def __iter__(self) -> Iterator[Image]:
pass
diff --git a/src/polystar/frame_generators/video_frame_generator.py b/src/polystar/frame_generators/video_frame_generator.py
index 22701b54effe603b8f43d0089924db7929793147..1d9220b736ba128ddb62a8eebff1fb963dcb42f0 100644
--- a/src/polystar/frame_generators/video_frame_generator.py
+++ b/src/polystar/frame_generators/video_frame_generator.py
@@ -1,26 +1,24 @@
-from dataclasses import dataclass
from pathlib import Path
-from typing import Any, Iterable, Optional
+from typing import Iterable, Iterator, Optional
import ffmpeg
from cv2.cv2 import CAP_PROP_POS_FRAMES
from memoized_property import memoized_property
-from polystar.frame_generators.cv2_frame_generator_abc import CV2FrameGeneratorABC
+from polystar.frame_generators.cv2_frame_generator_abc import CV2Capture, CV2FrameGenerator
+from polystar.models.image import Image
-@dataclass
-class VideoFrameGenerator(CV2FrameGeneratorABC):
+class VideoFrameGenerator(CV2FrameGenerator):
+ def __init__(self, video_path: Path, offset_seconds: Optional[int] = None):
+ super().__init__(str(video_path))
+ self.offset_seconds = offset_seconds
+ self.video_path = video_path
- video_path: Path
- offset_seconds: Optional[int]
-
- def _capture_params(self) -> Iterable[Any]:
- return (str(self.video_path),)
-
- def _post_opening_operation(self):
+ def __iter__(self) -> Iterator[Image]:
if self.offset_seconds:
- self._cap.set(CAP_PROP_POS_FRAMES, self._video_fps * self.offset_seconds - 2)
+ return CV2CaptureWithOffset(self.capture_params, self._video_fps * self.offset_seconds - 2)
+ return CV2Capture(self.capture_params)
@memoized_property
def _video_fps(self) -> int:
@@ -30,3 +28,9 @@ class VideoFrameGenerator(CV2FrameGeneratorABC):
continue
return round(eval(stream_info["avg_frame_rate"]))
raise ValueError(f"No fps found for video {self.video_path.name}")
+
+
+class CV2CaptureWithOffset(CV2Capture):
+ def __init__(self, capture_params: Iterable, offset_frames: int):
+ super().__init__(capture_params)
+ self._cap.set(CAP_PROP_POS_FRAMES, offset_frames)
diff --git a/src/polystar/view/cv2_results_viewer.py b/src/polystar/view/cv2_results_viewer.py
index 98594c241c1c625fdb25e9ecc3a0aa360f4e1fc2..55db910d75a3fab8054dcc012c951de87f0e055d 100644
--- a/src/polystar/view/cv2_results_viewer.py
+++ b/src/polystar/view/cv2_results_viewer.py
@@ -35,11 +35,11 @@ class CV2ResultViewer(ResultViewerABC):
self.delay = delay
self.name = name
self._current_image: Image = None
- self.finished = False
super().__init__(COLORS)
def __exit__(self, exc_type, exc_val, exc_tb):
cv2.destroyWindow(self.name)
+ return exc_type is KeyboardInterrupt
def new(self, image: Image):
self.height, self.width, _ = image.shape
@@ -80,7 +80,7 @@ class CV2ResultViewer(ResultViewerABC):
self.keycode_callbacks[keycode]()
def stop(self):
- self.finished = True
+ raise KeyboardInterrupt()
def _make_keycode_callbacks(self, end_key: str, key_callbacks: Dict[str, Callback]) -> Dict[int, Callback]:
key_callbacks[end_key] = self.stop
diff --git a/src/polystar/view/results_viewer_abc.py b/src/polystar/view/results_viewer_abc.py
index c1ddfd6d0818d6ba2b3d1d20424dce36b4ef1429..ea78e95c28b5ca6b54bc4df2da3ca9b1e55f78e3 100644
--- a/src/polystar/view/results_viewer_abc.py
+++ b/src/polystar/view/results_viewer_abc.py
@@ -46,6 +46,10 @@ class ResultViewerABC(ABC):
for obj in objects:
self.add_object(obj, forced_color=forced_color)
+ def display_image(self, image: Image):
+ self.new(image)
+ self.display()
+
def display_image_with_objects(self, image: Image, objects: Iterable[ROCOObject]):
self.new(image)
self.add_objects(objects)
diff --git a/src/research/dataset/twitch/robots_views_extractor.py b/src/research/dataset/twitch/robots_views_extractor.py
index aa58f27bb3b6b1b29dbf234bfa63ac623c92c40b..057b1dadf65e89b2b527e94e426c5b8acc08566e 100644
--- a/src/research/dataset/twitch/robots_views_extractor.py
+++ b/src/research/dataset/twitch/robots_views_extractor.py
@@ -25,7 +25,7 @@ class RobotsViewExtractor:
def run(self):
self._progress_bar = tqdm(
- enumerate(self.frame_generator.generate(), 1 + self.OFFSET_SECONDS * self.FPS),
+ enumerate(self.frame_generator, 1 + self.OFFSET_SECONDS * self.FPS),
total=self._get_number_of_frames(),
desc=f"Extracting robots views from video {self.video_name}.mp4",
unit="frames",
diff --git a/src/research/scripts/demo_pipeline_camera.py b/src/research/scripts/demo_pipeline_camera.py
index c2c79fd0c205ecca1e929a9ca05d54e4a9e8da9f..a06d889db19977059c5b63ecc335befbcfe93e36 100644
--- a/src/research/scripts/demo_pipeline_camera.py
+++ b/src/research/scripts/demo_pipeline_camera.py
@@ -22,14 +22,12 @@ class CameraPipelineDemo:
def run(self):
with CV2ResultViewer("TensorRT demo", key_callbacks={" ": self.cs_link.toggle}) as viewer:
- for image in self.webcam.generate():
+ for image in self.webcam:
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:
diff --git a/src/research/scripts/display_camera.py b/src/research/scripts/display_camera.py
new file mode 100644
index 0000000000000000000000000000000000000000..d90e654bad26179b6a2159e2d194002bb1d3b2a1
--- /dev/null
+++ b/src/research/scripts/display_camera.py
@@ -0,0 +1,22 @@
+from injector import inject
+
+from polystar.dependency_injection import make_injector
+from polystar.frame_generators.frames_generator_abc import FrameGeneratorABC
+from polystar.utils.fps import FPS
+from polystar.view.cv2_results_viewer import CV2ResultViewer
+
+
+@inject
+def display_camera(webcam: FrameGeneratorABC):
+ fps = FPS()
+ fps_camera = FPS()
+ with CV2ResultViewer("Live Camera") as viewer:
+ for image in webcam:
+ viewer.new(image)
+ viewer.add_text(f"FPS: {fps.tick():.1f} / {fps_camera.tick():.1f}", 10, 10, (0, 0, 0))
+ viewer.display()
+ fps_camera.skip()
+
+
+if __name__ == "__main__":
+ make_injector().call_with_injection(display_camera)