From 1862b4bf73afafe7788df563f28967b62429ac19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20Sp=C3=B6rk?= Date: Fri, 6 Oct 2023 12:29:40 +0200 Subject: [PATCH] Refactor capture imutils (now cv) + gear Improve logging --- foosball/__main__.py | 13 +++-- foosball/capture/GearStream.py | 50 +++++++++++++++++++ .../{FileVideoStream.py => OpenCVStream.py} | 6 +-- foosball/capture/Stream.py | 17 ++++--- foosball/capture/filecapture.py | 27 ---------- foosball/capture/gearcapture.py | 49 ------------------ foosball/pipe/BaseProcess.py | 11 ++-- foosball/pipe/Pipe.py | 8 ++- foosball/tracking/ai.py | 2 +- 9 files changed, 81 insertions(+), 102 deletions(-) create mode 100644 foosball/capture/GearStream.py rename foosball/capture/{FileVideoStream.py => OpenCVStream.py} (81%) delete mode 100644 foosball/capture/filecapture.py delete mode 100644 foosball/capture/gearcapture.py diff --git a/foosball/__main__.py b/foosball/__main__.py index 8f0dcfa..0566614 100644 --- a/foosball/__main__.py +++ b/foosball/__main__.py @@ -45,7 +45,7 @@ def file_path(string): ap.add_argument("-yp", "--ypad", type=int, default=20, help="Vertical padding applied to ROI detected by aruco markers") ap.add_argument("-s", "--scale", type=float, default=0.4, help="Scale stream") -ap.add_argument("-cap", "--capture", choices=['gear', 'imutils'], default='gear', help="capture backend") +ap.add_argument("-cap", "--capture", choices=['cv', 'gear'], default='gear', help="capture backend") ap.add_argument("-d", "--display", choices=['cv', 'gear'], default='cv', help="display backend") ap.add_argument("-g", "--gpu", choices=['preprocess', 'tracker', 'render'], nargs='+', default=["render"], help="use GPU") kwargs = vars(ap.parse_args()) @@ -82,12 +82,11 @@ def usage_and_exit(): source = kwargs.get('file') or kwargs.get('cameraId') if kwargs.get('capture') == 'gear': - from .capture.gearcapture import GearCapture - cap = GearCapture(source, framerate=32, resolution=(1280, 720)) - elif kwargs.get('capture') == 'imutils': - from .capture.filecapture import FileVideoStream - - cap = FileVideoStream(source) + from .capture.GearStream import GearStream + cap = GearStream(source, framerate=32, resolution=(1280, 720)) + elif kwargs.get('capture') == 'cv': + from .capture.OpenCVStream import OpenCVStream + cap = OpenCVStream(source) else: usage_and_exit() print(kwargs) diff --git a/foosball/capture/GearStream.py b/foosball/capture/GearStream.py new file mode 100644 index 0000000..cf40356 --- /dev/null +++ b/foosball/capture/GearStream.py @@ -0,0 +1,50 @@ +import logging +import cv2 +from vidgear.gears import VideoGear + +from foosball.capture.Stream import Stream + + +class GearStream(Stream): + + def __init__(self, video=None, resolution=(640, 480), framerate=60, **kwargs): + # not skipping frames is crucial + # otherwise gear will not terminate as it's not forwarding sentinel, due to lack of explicit eos support + super().__init__(skip_frames=False, **kwargs) + self.logger = logging.getLogger(__name__) + self.eos = False + # if a video path was not supplied, grab the reference + # to the webcam + if video is None or type(video) == int: + options = { + "CAP_PROP_FRAME_WIDTH": resolution[0], + "CAP_PROP_FRAME_HEIGHT": resolution[1], + "CAP_PROP_FPS": framerate + } + self.gear = VideoGear(source=video or 0, logging=True, **options).start() + # otherwise, grab a reference to the video file + else: + options = { + # "CAP_PROP_FRAME_WIDTH": 800, # resolution 320x240 + # "CAP_PROP_FRAME_HEIGHT": 600, + # "CAP_PROP_FPS": 60, # framerate 60fps + } + self.gear = VideoGear(source=video, logging=True, **options).start() + self.logger.info(f"framerate = {self.gear.framerate}") + + def is_eos(self): + return self.eos + + def next_frame(self): + frame = self.gear.read() + self.eos = (frame is not None) + return self.eos, frame + + def close_capture(self): + self.gear.stop() + + def dim(self): + width = int(self.gear.stream.stream.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(self.gear.stream.stream.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + return [width, height] diff --git a/foosball/capture/FileVideoStream.py b/foosball/capture/OpenCVStream.py similarity index 81% rename from foosball/capture/FileVideoStream.py rename to foosball/capture/OpenCVStream.py index dfbd168..97a897d 100644 --- a/foosball/capture/FileVideoStream.py +++ b/foosball/capture/OpenCVStream.py @@ -4,11 +4,11 @@ from foosball.capture.Stream import Stream -class FileVideoStream(Stream): +class OpenCVStream(Stream): - def __init__(self, path, *args, **kwargs): + def __init__(self, source: str | int, *args, **kwargs): super().__init__(*args, **kwargs) - self.cap = cv2.VideoCapture(path) + self.cap = cv2.VideoCapture(source) self.total_frames = self.cap.get(cv2.CAP_PROP_FRAME_COUNT) def next_frame(self) -> (bool, np.array): diff --git a/foosball/capture/Stream.py b/foosball/capture/Stream.py index 5fd0011..1ef18d1 100644 --- a/foosball/capture/Stream.py +++ b/foosball/capture/Stream.py @@ -1,8 +1,10 @@ +import multiprocessing from abc import abstractmethod from multiprocessing import Queue from queue import Full from threading import Thread +from foosball.models import Frame from foosball.pipe.BaseProcess import Msg @@ -22,14 +24,14 @@ def __init__(self, maxsize=128, skip_frames=True, timeout=2, *args, **kwargs): self.Q = Queue(maxsize=maxsize) @abstractmethod - def is_eos(self): + def is_eos(self) -> bool: return False @abstractmethod - def next_frame(self): + def next_frame(self) -> Frame: pass - def read_frame(self): + def read_frame(self) -> (bool, Frame): # grab the current frame flag, frame = self.next_frame() @@ -47,10 +49,10 @@ def read_frame(self): return flag, frame @property - def output(self): + def output(self) -> multiprocessing.Queue: return self.Q - def send_frame(self, frame): + def send_frame(self, frame) -> None: msg = Msg(kwargs={'frame': frame}) if frame is not None else None while True: try: @@ -62,7 +64,7 @@ def send_frame(self, frame): if self.stopped: break - def run(self): + def run(self) -> None: while not self.stopped: # read the next frame from the file grabbed, frame = self.read_frame() @@ -77,7 +79,7 @@ def run(self): self.close_capture() @abstractmethod - def close_capture(self): + def close_capture(self) -> None: pass @abstractmethod @@ -86,4 +88,3 @@ def dim(self) -> [int, int]: def stop(self): self.stopped = True - return self diff --git a/foosball/capture/filecapture.py b/foosball/capture/filecapture.py deleted file mode 100644 index 2f7f926..0000000 --- a/foosball/capture/filecapture.py +++ /dev/null @@ -1,27 +0,0 @@ -from imutils.video import VideoStream - -from .FileVideoStream import FileVideoStream - - -class FileCapture: - - def __init__(self, video=None): - # if a video path was not supplied, grab the reference - # to the webcam - if video is None: - self.cap = VideoStream(src=1).start() - # otherwise, grab a reference to the video file - else: - # #vs = FileVideoStream(args['file']).start() - self.cap = FileVideoStream(video).start() - - self.is_file_capture = video is not None - - def next(self): - return self.cap.read() - - def stop(self): - self.cap.stop() - - def dim(self): - return self.cap.dim() diff --git a/foosball/capture/gearcapture.py b/foosball/capture/gearcapture.py deleted file mode 100644 index e76afc1..0000000 --- a/foosball/capture/gearcapture.py +++ /dev/null @@ -1,49 +0,0 @@ -import logging -import cv2 -from vidgear.gears import VideoGear - -logger = logging.getLogger(__name__) - -class GearCapture: - - def __init__(self, video=None, **kwargs): - self.source = 'file' if isinstance(video, str) else 'cam' - # if a video path was not supplied, grab the reference - # to the webcam - if video is None or type(video) == int: - resolution = kwargs.get('resolution', (640, 480)) - options = { - "CAP_PROP_FRAME_WIDTH": resolution[0], # resolution 320x240 - "CAP_PROP_FRAME_HEIGHT": resolution[1], - "CAP_PROP_FPS": kwargs.get('framerate'), # framerate 60fps - } - self.cap = VideoGear(source=video or 0, logging=True, **options).start() - # otherwise, grab a reference to the video file - else: - options = { - # "CAP_PROP_FRAME_WIDTH": 800, # resolution 320x240 - # "CAP_PROP_FRAME_HEIGHT": 600, - # "CAP_PROP_FPS": 60, # framerate 60fps - } - self.cap = VideoGear(source=video, logging=False, **options).start() - logger.info(f"framerate = {self.cap.framerate}") - - self.is_file_capture = video is not None - - def next(self): - return self.cap.read() - - def stop(self): - self.cap.stop() - - def stream(self): - if self.source == 'file': - return self.cap.cap.cap - else: - return self.cap.cap.cap - - def dim(self): - width = int(self.stream().get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(self.stream().get(cv2.CAP_PROP_FRAME_HEIGHT)) - - return [width, height] diff --git a/foosball/pipe/BaseProcess.py b/foosball/pipe/BaseProcess.py index 861c4f8..c6ae3a7 100644 --- a/foosball/pipe/BaseProcess.py +++ b/foosball/pipe/BaseProcess.py @@ -21,14 +21,14 @@ def __init__(self, args=None, kwargs=None): self.kwargs = kwargs self.args = args - class BaseProcess(multiprocessing.Process): - def __init__(self, *args, **kwargs): + def __init__(self, send_receive_timeout=0.5, *args, **kwargs): super().__init__(daemon=True, *args, **kwargs) self.args = args self.logger = logging.getLogger(kwargs.get('name') or __name__) self.kwargs = kwargs self.stop_event: multiprocessing.Event = multiprocessing.Event() + self.send_receive_timeout = send_receive_timeout def set_input(self, inq): self.inq = inq @@ -58,7 +58,7 @@ def stop(self): def send(self, msg: Msg): while True: try: - self.outq.put(msg, block=True, timeout=0.5) + self.outq.put(msg, block=True, timeout=self.send_receive_timeout) break except Full: print("Queue is full") @@ -68,7 +68,7 @@ def send(self, msg: Msg): def receive(self) -> Msg: while True: try: - return self.inq.get(block=True, timeout=0.5) + return self.inq.get(block=True, timeout=self.send_receive_timeout) except Empty: print("Queue is empty") if self.stop_event.is_set(): @@ -82,11 +82,12 @@ def run(self): try: msg = self.inq.get_nowait() if msg is SENTINEL: + self.logger.debug("received SENTINEL") self.send(SENTINEL) break out = self.process(msg) if out is None: - self.logger.debug("out is None") + self.logger.debug("sending SENTINEL") self.send(out) except Empty: pass diff --git a/foosball/pipe/Pipe.py b/foosball/pipe/Pipe.py index 735cbd6..aad17ed 100644 --- a/foosball/pipe/Pipe.py +++ b/foosball/pipe/Pipe.py @@ -44,6 +44,10 @@ def output(self): def input(self): return self.queues[0] + @property + def qsizes(self): + return [q.qsize() for q in self.queues] + def stop(self): self.logger.debug("Stopping pipe...") self.stream.stop() @@ -51,7 +55,7 @@ def stop(self): p.stop() # empty the last queue - self.logger.debug(f"Queue sizes: {' '.join([f'{q.qsize()}' for q in self.queues])}") + self.logger.debug(f"Queue sizes: {' '.join([str(s) for s in self.qsizes])}") self.logger.debug("draining queues...") # draining all queues for good for q in self.queues: @@ -59,5 +63,5 @@ def stop(self): self.logger.debug("joining...") for p in reversed(self.processes): p.join() - self.logger.debug(f"Queue sizes: {' '.join([f'{q.qsize()}' for q in self.queues])}") + self.logger.debug(f"Queue sizes: {' '.join([str(s) for s in self.qsizes])}") self.logger.debug("Stopped pipe") diff --git a/foosball/tracking/ai.py b/foosball/tracking/ai.py index 7517e26..14be524 100644 --- a/foosball/tracking/ai.py +++ b/foosball/tracking/ai.py @@ -90,7 +90,7 @@ def step_frame(): self.adjust_calibration() msg = self.tracking.output.get(block=False) if msg is None: - print("msg is None") + self.logger.debug("received SENTINEL") break f = msg.kwargs['result'] self.fps.stop()