Skip to content

Commit

Permalink
enhancement(zmq): Add zmq hook WIP
Browse files Browse the repository at this point in the history
Add viewbox to Framestats
  • Loading branch information
DarwinsBuddy committed Nov 23, 2024
1 parent d0ed4cd commit 501e458
Show file tree
Hide file tree
Showing 6 changed files with 19 additions and 8 deletions.
2 changes: 2 additions & 0 deletions foosball/tracking/analyzer/ScoreAnalyzer.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ def analyze(self, msg: Msg, timestamp: dt.datetime) -> [ScoreAnalyzerResult, [In
goals = msg.data["Tracker"].goals
track = msg.data["Tracker"].ball_track
ball = msg.data["Tracker"].ball
viewbox = msg.data["Tracker"].viewbox
team_scored = None
try:
self.check_reset_score()
Expand Down Expand Up @@ -103,6 +104,7 @@ def analyze(self, msg: Msg, timestamp: dt.datetime) -> [ScoreAnalyzerResult, [In
score=self.score,
ball=ball,
dims=self.dims,
viewbox=viewbox,
timestamp=timestamp.strftime("%Y-%m-%dT%H:%M:%S")
)
self.frame_hooks.invoke(stats.to_dict())
Expand Down
3 changes: 3 additions & 0 deletions foosball/tracking/analyzer/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,15 @@

from foosball.models import Goals, Track, Blob, Score, FrameDimensions


@dataclass
class FrameStats:
goals: Goals | None
track: Track | None
ball: Blob | None
score: Score | None
dims: FrameDimensions | None
viewbox: list | None
timestamp: str | None

def to_dict(self) -> dict:
Expand All @@ -18,5 +20,6 @@ def to_dict(self) -> dict:
"ball": self.ball.to_json() if self.ball else None,
"score": self.score.to_json() if self.score else None,
"dims": self.dims.to_json() if self.dims else None,
"viewbox": self.viewbox if self.viewbox else None,
"timestamp": self.timestamp
}
11 changes: 7 additions & 4 deletions foosball/tracking/preprocess/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ def __init__(self, dims: FrameDimensions, goal_detector: GoalDetector, headless=
self.calibration_out = Queue() if self.goals_calibration else None
self.config_in = Queue() if self.goals_calibration else None
self.goal_detector = goal_detector
self.viewbox = None

def config_input(self, config: GoalColorConfig) -> None:
if self.goals_calibration:
Expand Down Expand Up @@ -119,8 +120,10 @@ def process(self, msg: Msg) -> Msg:

self.frames_since_last_marker_detection = (self.frames_since_last_marker_detection + 1) % self.redetect_markers_frame_threshold
if len(self.markers) == 4:

# crop and warp
preprocessed, self.homography_matrix = self.four_point_transform(frame, self.markers)
preprocessed, self.homography_matrix, marker_pts = self.four_point_transform(frame, self.markers)
self.viewbox = [pt.tolist() for pt in marker_pts]
if trigger_marker_detection:
# detect goals anew
goals_detection_result = self.goal_detector.detect(preprocessed)
Expand Down Expand Up @@ -150,7 +153,7 @@ def process(self, msg: Msg) -> Msg:
traceback.print_exc()
# Not passing original msg due to performance impact (copying whole frames, etc.)
return Msg(info=info, data={
"Preprocessor": PreprocessorResult(original=self.iproc(frame), preprocessed=self.iproc(preprocessed), homography_matrix=self.homography_matrix, goals=self.goals)
"Preprocessor": PreprocessorResult(original=self.iproc(frame), preprocessed=self.iproc(preprocessed), homography_matrix=self.homography_matrix, goals=self.goals, viewbox=self.viewbox)
})

@staticmethod
Expand Down Expand Up @@ -181,7 +184,7 @@ def order_points(pts: np.ndarray) -> np.ndarray:
# return the ordered coordinates
return rect

def four_point_transform(self, frame: Frame, markers: list[Aruco]) -> tuple[Frame, [int, int]]:
def four_point_transform(self, frame: Frame, markers: list[Aruco]) -> tuple[Frame, [int, int], [int, int, int, int]]:
pts = np.array([self.corners2pt(marker.corners) for marker in markers])
# obtain a consistent order of the points and unpack them
# individually
Expand Down Expand Up @@ -226,7 +229,7 @@ def four_point_transform(self, frame: Frame, markers: list[Aruco]) -> tuple[Fram
# since we want to draw on the original image (later on), which is still originally not warped,
# we want to have a function set to project from/onto the warped/un-warped version of the frame
# for future reference. so we return the warped image and the used homography matrix
return warped, homography_matrix
return warped, homography_matrix, src_pts


def project_point(pt: Point, homography_matrix: np.array, mode: WarpMode):
Expand Down
1 change: 1 addition & 0 deletions foosball/tracking/preprocess/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,3 +12,4 @@ class PreprocessorResult:
preprocessed: Optional[Frame]
homography_matrix: Optional[np.ndarray] # 3x3 matrix used to warp the image and project points
goals: Optional[Goals]
viewbox: Optional[list]
9 changes: 5 additions & 4 deletions foosball/tracking/tracker/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ def close(self) -> None:
clear(self.calibration_out)
self.calibration_out.close()

def update_ball_track(self, detected_ball: Blob) -> Track:
def update_ball_track(self, detected_ball: Blob) -> Track | None:
if detected_ball is not None:
center = detected_ball.center
# update the points queue (track history)
Expand All @@ -47,7 +47,7 @@ def update_ball_track(self, detected_ball: Blob) -> Track:
self.ball_track.appendleft(None)
return self.ball_track

def get_info(self, ball_track: Track) -> [Info]:
def get_info(self, ball_track: Track | None) -> [Info]:
info = [
Info(verbosity=Verbosity.DEBUG, title="Track length", value=f"{str(sum([1 for p in ball_track or [] if p is not None])).rjust(2, ' ')}"),
Info(verbosity=Verbosity.TRACE, title="Calibration", value=f"{self.calibrationMode if self.calibrationMode is not None else 'off'}"),
Expand Down Expand Up @@ -75,6 +75,7 @@ def process(self, msg: Msg) -> Msg:
ball = None
goals = data.goals
ball_track = None
viewbox = data.viewbox
tracker_info = []
try:
if not self.off:
Expand Down Expand Up @@ -105,6 +106,6 @@ def process(self, msg: Msg) -> Msg:
traceback.print_exc()
# Not passing original msg due to performance impact (copying whole frames, etc.)
if not self.verbose:
return Msg(info=tracker_info, data={"Tracker": TrackerResult(frame=data.original, goals=goals, ball_track=ball_track, ball=ball)})
return Msg(info=tracker_info, data={"Tracker": TrackerResult(frame=data.original, goals=goals, ball_track=ball_track, ball=ball, viewbox=viewbox)})
else:
return Msg(info=tracker_info, data={"Tracker": TrackerResult(frame=data.preprocessed, goals=goals, ball_track=ball_track, ball=ball)})
return Msg(info=tracker_info, data={"Tracker": TrackerResult(frame=data.preprocessed, goals=goals, ball_track=ball_track, ball=ball, viewbox=viewbox)})
1 change: 1 addition & 0 deletions foosball/tracking/tracker/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,4 @@ class TrackerResult:
goals: Goals | None
ball_track: Track | None
ball: Blob | None
viewbox: list | None

0 comments on commit 501e458

Please sign in to comment.