refactoring(data): Refactoring inter-process data passing WIP

foosball/models.py

@@ -1,8 +1,8 @@
 import collections
 from dataclasses import dataclass
-from datetime import datetime
+from typing import TypeVar
 from enum import Enum
-from typing import Optional, Union
+from typing import Optional, Union, Generic
 
 import cv2
 import numpy as np
@@ -138,29 +138,32 @@ def to_string(self):
         return " - ".join([i.to_string() for i in self.infos])
 
 
+R = TypeVar('R')
+
+
 @dataclass
-class TrackResult:
-    frame: CPUFrame
-    goals: Optional[Goals]
-    ball_track: Track
-    ball: Blob
-    info: Info
+class Result(Generic[R]):
+    info: InfoLog
+    data: R
 
 
 @dataclass
-class AnalyzeResult:
+class TrackerResultData:
     frame: CPUFrame
-    score: Score
     goals: Optional[Goals]
     ball_track: Track
     ball: Blob
-    info: Info
+
+
+TrackerResult = Result[TrackerResultData]
 
 
 @dataclass
-class PreprocessResult:
+class PreprocessorResultData:
     original: CPUFrame
     preprocessed: Optional[CPUFrame]
     homography_matrix: Optional[np.ndarray]  # 3x3 matrix used to warp the image and project points
     goals: Optional[Goals]
-    info: InfoLog
+
+
+PreprocessorResult = Result[PreprocessorResultData]

foosball/pipe/BaseProcess.py

@@ -1,20 +1,35 @@
 import abc
-import dataclasses
 import logging
 import multiprocessing
 import traceback
 import datetime as dt
 from queue import Empty, Full
+from dataclasses import dataclass
 
+from foosball.models import InfoLog, Result, R
 from foosball.pipe.Pipe import clear, SENTINEL
 
 
-@dataclasses.dataclass
+# TODO: check why merging into one Msg is having a huge impact on FPS
+@dataclass
 class Msg:
     args: list[any]
     kwargs: dict
+    info: InfoLog = None
     timestamp: dt.datetime = dt.datetime.now()
 
+    def add(self, name: str, data: any, info=InfoLog([])):
+        self.kwargs[name] = data
+        if self.info is not None:
+            self.info.concat(info)
+        else:
+            self.info = InfoLog([])
+
+    def remove(self, name) -> Result[R]:
+        return self.kwargs.pop(name)
+
+
+
     def __init__(self, args=None, kwargs=None, timestamp=dt.datetime.now()):
         if kwargs is None:
             kwargs = dict()

foosball/tracking/__init__.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from const import GPU, CalibrationMode
-from .analyze import Analyzer
+from .analyzer.analyze import Analyzer
 from .preprocess import PreProcessor
 from .render import Renderer
 from .tracker import Tracker

foosball/tracking/ai.py

@@ -106,8 +106,8 @@ def step_frame():
                         self.logger.debug("received SENTINEL")
                         break
                     self.fps.update()
-                    frame = msg.kwargs['result']
-                    info: InfoLog = msg.kwargs['info']
+                    frame = msg.kwargs.get('Renderer', None)
+                    info: InfoLog = msg.info
                     self.fps.stop()
                     fps = int(self.fps.fps())
                     if not self.headless:

foosball/tracking/analyze.py → foosball/tracking/analyzer/ScoreAnalyzer.py

@@ -3,25 +3,21 @@
 from typing import Optional
 import datetime as dt
 
-from .. import hooks
-from ..hooks import generate_goal_webhook
-from ..models import Team, Goals, Score, AnalyzeResult, Track, Info, Verbosity
-from ..pipe.BaseProcess import BaseProcess, Msg
-from ..utils import contains
+from . import AbstractAnalyzer, ScoreAnalyzerResult, ScoreAnalyzerResultData
+from ...models import Team, Goals, Score, Track, Info, Verbosity, TrackerResult, InfoLog
+from ...utils import contains
 
 
-class Analyzer(BaseProcess):
+class ScoreAnalyzer(AbstractAnalyzer):
     def close(self):
         pass
 
-    def __init__(self, audio: bool = False, webhook: bool = False, goal_grace_period_sec: float = 1.0, *args, **kwargs):
-        super().__init__(name="Analyzer")
+    def __init__(self, goal_grace_period_sec: float = 1.0, *args, **kwargs):
+        super().__init__(name="ScoreAnalyzer")
         self.kwargs = kwargs
         self.goal_grace_period_sec = goal_grace_period_sec
         self.score = Score()
         self.score_reset = multiprocessing.Event()
-        self.audio = audio
-        self.webhook = webhook
         self.last_track_sighting: dt.datetime | None = None
         self.last_track: Optional[Track] = None
         self.goal_candidate = None
@@ -46,22 +42,14 @@ def goal_shot(self, goals: Goals, track: Track) -> Optional[Team]:
             self.logger.error(f"Error {e}")
         return None
 
-    def call_hooks(self, team: Team) -> None:
-        if self.audio:
-            hooks.play_random_sound('goal')
-        if self.webhook:
-            hooks.webhook(generate_goal_webhook(team))
-
-    def process(self, msg: Msg) -> Msg:
-        track_result = msg.kwargs['result']
-        goals = track_result.goals
-        ball = track_result.ball
-        track = track_result.ball_track
-        frame = track_result.frame
-        info = track_result.info
+    def analyze(self, track_result: TrackerResult, timestamp: dt.datetime) -> ScoreAnalyzerResult:
+        goals = track_result.data.goals
+        track = track_result.data.ball_track
+        info = InfoLog([])
+        team_scored = None
         try:
             self.check_reset_score()
-            now = msg.timestamp  # take frame timestamp as now instead of dt.datetime.now (to prevent drift due to pushing/dragging pipeline)
+            now = timestamp  # take frame timestamp as now instead of dt.datetime.now (to prevent drift due to pushing/dragging pipeline)
             no_track_sighting_in_grace_period = (now - self.last_track_sighting).total_seconds() >= self.goal_grace_period_sec if self.last_track_sighting is not None else None
             if not self.is_track_empty(track):
                 # track is not empty, so we save our state and remove a potential goal (which was wrongly tracked)
@@ -74,6 +62,7 @@ def process(self, msg: Msg) -> Msg:
                 # whatever happens first
                 if self.goal_candidate is not None and self.last_track_sighting is not None and no_track_sighting_in_grace_period:
                     self.count_goal(self.goal_candidate)
+                    team_scored = self.goal_candidate
                     self.goal_candidate = None
                 elif self.goal_candidate is None:
                     # if track is empty, and we have no current goal candidate, check if there is one
@@ -83,7 +72,7 @@ def process(self, msg: Msg) -> Msg:
             traceback.print_exc()
         self.last_track = track
         info.append(Info(verbosity=Verbosity.INFO, title="Score", value=self.score.to_string()))
-        return Msg(timestamp=msg.timestamp, kwargs={"result": AnalyzeResult(score=self.score, ball=ball, goals=goals, frame=frame, info=info, ball_track=track)})
+        return ScoreAnalyzerResult(data=ScoreAnalyzerResultData(score=self.score, team_scored=team_scored), info=info)
 
     def check_reset_score(self):
         if self.score_reset.is_set():
@@ -92,9 +81,6 @@ def check_reset_score(self):
 
     def count_goal(self, team: Team):
         self.score.inc(team)
-        if team is not None:
-            self.logger.info(f"GOAL Team:{team} - {self.score.red} : {self.score.blue}")
-            self.call_hooks(team)
 
     def reset_score(self):
-        self.score_reset.set()
+        self.score_reset.set()

foosball/tracking/analyzer/__init__.py

@@ -0,0 +1,25 @@
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime
+
+from foosball.models import TrackerResult, Score, Team, Result
+
+@dataclass
+class ScoreAnalyzerResultData:
+    score: Score
+    team_scored: Team
+
+
+ScoreAnalyzerResult = Result[ScoreAnalyzerResultData]
+
+
+class AbstractAnalyzer(ABC):
+
+    def __init__(self, name: str = "UnknownAnalyzer", **kwargs):
+        self.name = name
+        self.logger = logging.getLogger(name)
+
+    @abstractmethod
+    def analyze(self, track_result: TrackerResult, timestamp: datetime) -> dict:
+        pass

foosball/tracking/analyzer/analyze.py

@@ -0,0 +1,47 @@
+import traceback
+
+from .ScoreAnalyzer import ScoreAnalyzer
+from ... import hooks
+from ...hooks import generate_goal_webhook
+from ...models import Team, TrackerResult
+from ...pipe.BaseProcess import BaseProcess, Msg
+
+
+class Analyzer(BaseProcess):
+    def close(self):
+        pass
+
+    def __init__(self, audio: bool = False, webhook: bool = False, goal_grace_period_sec: float = 1.0, *args, **kwargs):
+        super().__init__(name="Analyzer")
+        self.kwargs = kwargs
+        self.analyzers = [ScoreAnalyzer(goal_grace_period_sec, args, kwargs)]
+        # TODO: catch up here
+        self.effects = []
+        self.audio = audio
+        self.webhook = webhook
+
+    def call_hooks(self, team: Team) -> None:
+        if self.audio:
+            hooks.play_random_sound('goal')
+        if self.webhook:
+            hooks.webhook(generate_goal_webhook(team))
+
+    def process(self, msg: Msg) -> Msg:
+        track_result: TrackerResult = msg.kwargs['Tracker']
+        data = track_result.data
+
+        for a in self.analyzers:
+            try:
+                ar = a.analyze(track_result, msg.timestamp)
+                msg.add(a.name, ar, info=ar.info)
+            except Exception as e:
+                self.logger.error("Error in Analyzer - analyzers ", e)
+                traceback.print_exc()
+        # TODO: catch up here
+        for e in self.effects:
+            try:
+                e.invoke(track_result, msg.timestamp)
+            except Exception as e:
+                self.logger.error("Error in Analyzer - effects ", e)
+                traceback.print_exc()
+        return msg

foosball/tracking/effects/__init__.py

@@ -0,0 +1,8 @@
+from abc import abstractmethod, ABC
+
+
+class Effect(ABC):
+
+    @abstractmethod
+    def invoke(self, *args, **kwargs):
+        pass

foosball/tracking/preprocess.py

@@ -10,8 +10,8 @@
 from ..arUcos import calibration
 from ..arUcos.models import Aruco
 from ..detectors.color import GoalDetector, GoalColorConfig
-from ..models import Frame, PreprocessResult, Point, Rect, Blob, Goals, FrameDimensions, ScaleDirection, \
-    InfoLog, Info, Verbosity
+from ..models import Frame, PreprocessorResult, Point, Rect, Blob, Goals, FrameDimensions, ScaleDirection, \
+    InfoLog, Info, Verbosity, PreprocessorResultData
 from ..pipe.BaseProcess import BaseProcess, Msg
 from ..pipe.Pipe import clear
 from ..utils import ensure_cpu, generate_processor_switches, relative_change, scale
@@ -44,7 +44,7 @@ def __init__(self, dims: FrameDimensions, goal_detector: GoalDetector, headless=
                  redetect_markers_frames: int = 60, aruco_dictionary=cv2.aruco.DICT_4X4_1000,
                  aruco_params=cv2.aruco.DetectorParameters(), xpad: int = 50, ypad: int = 20,
                  goal_change_threshold: float = 0.10, useGPU: bool = False, calibrationMode=None, verbose=False, **kwargs):
-        super().__init__(name="Preprocess")
+        super().__init__(name="Preprocessor")
         self.dims = dims
         self.goal_change_threshold = goal_change_threshold
         self.redetect_markers_frame_threshold = redetect_markers_frames
@@ -147,7 +147,11 @@ def process(self, msg: Msg) -> Msg:
         except Exception as e:
             self.logger.error(f"Error in preprocess {e}")
             traceback.print_exc()
-        return Msg(timestamp=msg.timestamp, kwargs={"result": PreprocessResult(self.iproc(frame), self.iproc(preprocessed), self.homography_matrix, self.goals, info)})
+        msg.add("Preprocessor", PreprocessorResult(
+            data=PreprocessorResultData(self.iproc(frame), self.iproc(preprocessed), self.homography_matrix, self.goals),
+            info=info
+        ))
+        return msg
 
     @staticmethod
     def corners2pt(corners) -> [int, int]:

foosball/tracking/render.py

@@ -98,7 +98,6 @@ def r_track(frame, ball_track, scale) -> None:
         thickness = max(1, int(int(np.sqrt(ball_track.maxlen / float(i + 1)) * 2) * scale))
         cv2.line(frame, ball_track[i - 1], ball_track[i], (b, g, r), thickness)
 
-
 class Renderer(BaseProcess):
     def close(self):
         pass
@@ -111,16 +110,17 @@ def __init__(self, dims: FrameDimensions, headless=False, useGPU: bool = False,
         [self.proc, self.iproc] = generate_processor_switches(useGPU)
 
     def process(self, msg: Msg) -> Msg:
-        analyze_result = msg.kwargs['result']
-        info: InfoLog = analyze_result.info
+        goal_analyzer = msg.kwargs["ScoreAnalyzer"]
+        tracker = msg.kwargs["Tracker"]
+        info: InfoLog = msg.info
         try:
             if not self.headless:
-                shape = analyze_result.frame.shape
-                f = self.proc(analyze_result.frame)
-                ball = analyze_result.ball
-                goals = analyze_result.goals
-                track = analyze_result.ball_track
-                score = analyze_result.score
+                shape = tracker.data.frame.shape
+                f = self.proc(tracker.data.frame)
+                ball = tracker.data.ball
+                goals = tracker.data.goals
+                track = tracker.data.ball_track
+                score = goal_analyzer.data.score
 
                 if ball is not None:
                     r_ball(f, ball)
@@ -131,11 +131,8 @@ def process(self, msg: Msg) -> Msg:
                 r_score(f, score, text_scale=1, thickness=4)
                 if self.infoVerbosity is not None:
                     r_info(f, shape, info.filter(self.infoVerbosity), text_scale=0.5, thickness=1)
-                return Msg(kwargs={"result": self.iproc(f), 'info': info})
-            else:
-                return Msg(kwargs={"result": None, 'info': info})
-
+                msg.add("Renderer", self.iproc(f), info=InfoLog([]))
         except Exception as e:
             logger.error(f"Error in renderer {e}")
             traceback.print_exc()
-        return Msg(analyze_result)
+        return msg