RezaAdousi/ball_game
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
2e4c6964b21c4d7354ef718d99d640f94badea61
ball_game/detection.py
T
Reza Adousi 2e4c6964b2 Initial commit
2026-10-31 23:04:47 +03:30

1133 lines
34 KiB
Python
Raw Blame History

# detection.py
#
# Finds the white paper stickers stuck on the monitor, as ROTATED
# rectangles, in a stable, flicker-free way.
#
# Input: a BGR image of the monitor already warped into screen
# space (Calibration.transform_frame) at DETECTION_SCALE of the real
# screen resolution. Output: one dict per sticker, in real SCREEN
# pixels (see geometry.py for the fields/angle convention).
#
# How it decides, in order (each step exists because a real
# exhibition hall breaks the simpler version):
#
# 1. Find generous "maybe" blobs: anything LOCALLY brighter than
# its own immediate surroundings and not very colored, inside
# the calibration border (see imaging.py - brightness and color
# are measured in CIELAB, not HSV, specifically because that
# separates the two cleanly). "Locally brighter" rather than
# "brighter than one fixed number for the whole frame" is what
# lets a sticker sitting in a dimmer part of the hall still be
# found - a global floor would simply never see it, no matter
# how good every later step is. The area the game itself
# painted white for already-detected stickers is blanked out
# first, otherwise a removed sticker would be "seen" forever in
# its own painted white (see DETECTED_OBSTACLE_PAINT_*), and any
# small hole a glare reflection punches through a blob (or a
# thin glare streak that splits it in two) is patched before
# measuring it (see imaging.fill_holes).
#
# 2. Re-draw each blob's edge at the half-way lightness between
# that blob's own paper brightness and the dark screen beside
# it. A dim sticker in a shadowy corner and an identical one
# under a spotlight then come out the SAME size - a fixed
# brightness cutoff can't do that, and it is what made
# identical stickers seem to be "detected by force" some of the
# time and not others.
#
# 3. Measure the refined blob: best-fit rotated rectangle (true
# side lengths and angle), how well it fills that rectangle
# (rejects round glows), convexity, uniformity of brightness
# (rejects hot-spot reflections), average saturation (rejects
# skin/colored objects), and whether the ring of pixels just
# outside it is dark (a sticker sits on a black screen).
#
# 4. All stickers are identical, so once a couple have been seen
# the size is learned and a blob of a different size is
# rejected. Two or three stickers glued edge to edge form one
# bigger rectangle and are accepted as whole multiples.
#
# 5. Every sticker is tracked separately over time. NEW stickers
# must pass the strict tests on consecutive passes; stickers
# already established are kept under slightly relaxed tests, so
# a real sticker doesn't flicker on borderline frames, and they
# survive a hand in front of them for a few seconds. Motion is
# smoothed with a dead-band (no jitter when nothing moves) and
# catches up fast when a sticker really is moved.
import math
import time
from collections import deque
import cv2
import numpy as np
import imaging
from geometry import (
normalize_angle,
blend_angle,
obb_corners,
extract_obb,
paint_dimensions,
)
from config import (
SCREEN_WIDTH,
SCREEN_HEIGHT,
CALIBRATION_OUTER_BORDER_SIZE,
CALIBRATION_INNER_BORDER_SIZE,
DETECTION_SCALE,
MIN_OBSTACLE_AREA,
MAX_OBSTACLE_AREA,
OBSTACLE_ADAPTIVE_OFFSET,
OBSTACLE_MIN_ABSOLUTE_LIGHTNESS,
OBSTACLE_ADAPTIVE_BLOCK_SPAN,
OBSTACLE_ADAPTIVE_DEFAULT_SHORT_SIDE,
OBSTACLE_COARSE_CHROMA_MAX,
OBSTACLE_EDGE_LEVEL,
MIN_OBSTACLE_BRIGHTNESS,
MAX_OBSTACLE_CHROMA,
MIN_OBSTACLE_RECTANGULARITY,
MIN_OBSTACLE_SOLIDITY,
MAX_OBSTACLE_ASPECT_RATIO,
MAX_OBSTACLE_BRIGHTNESS_STD,
MIN_OBSTACLE_DARK_SURROUND_RATIO,
OBSTACLE_BORDER_MASK_MARGIN,
OBSTACLE_EDGE_MARGIN,
OBSTACLE_SIZE_TOLERANCE,
OBSTACLE_EXPECTED_SIZE,
OBSTACLE_ALLOW_MERGED_STICKERS,
OBSTACLE_SNAP_TO_EXPECTED_SIZE,
OBSTACLE_SIZE_FORGET_SECONDS,
OBSTACLE_RELAXED_FACTOR,
OBSTACLE_REQUIRED_STABLE_FRAMES,
OBSTACLE_MAX_MISSING_FRAMES,
OBSTACLE_PENDING_MAX_MISSING_FRAMES,
OBSTACLE_STALE_MISSES,
OBSTACLE_TEMPORAL_MATCH_DISTANCE,
OBSTACLE_MASK_OPEN_KERNEL,
OBSTACLE_MASK_CLOSE_KERNEL,
SHOW_DETECTED_OBSTACLES,
DETECTED_OBSTACLE_PAINT_INSET_RATIO,
)
# Extra margin (screen px) added around the painted patch when it is
# blanked out, to absorb the small misalignment between the pixels
# we painted and where the camera sees them.
_BLANK_MARGIN = 5.0
# Below this long/short ratio a sticker is "square enough" that its
# orientation only matters modulo 90 degrees.
_SQUARE_ASPECT = 1.15
class _Track:
"""One physical sticker followed over time."""
__slots__ = (
"id",
"cx",
"cy",
"angle",
"long",
"short",
"hits",
"misses",
"confirmed",
)
def __init__(self, track_id, candidate):
self.id = track_id
self.cx = candidate["cx"]
self.cy = candidate["cy"]
self.angle = candidate["angle"]
self.long = candidate["long"]
self.short = candidate["short"]
self.hits = 1
self.misses = 0
self.confirmed = False
class Detector:
def __init__(self):
self.scale = float(DETECTION_SCALE)
self.det_width = int(round(SCREEN_WIDTH * self.scale))
self.det_height = int(round(SCREEN_HEIGHT * self.scale))
self._s2 = self.scale * self.scale
# Kernels, sized in screen pixels then scaled (always odd).
self._open_kernel = self._kernel(OBSTACLE_MASK_OPEN_KERNEL)
self._close_kernel = self._kernel(OBSTACLE_MASK_CLOSE_KERNEL)
self._small_kernel = np.ones((3, 3), np.uint8)
# Region where stickers are searched: everything inside the
# calibration border (which is bright by design and would
# merge with a sticker touching it).
border = int(
round(
(
CALIBRATION_OUTER_BORDER_SIZE
+ CALIBRATION_INNER_BORDER_SIZE
+ OBSTACLE_BORDER_MASK_MARGIN
)
* self.scale
)
)
self._roi = np.zeros(
(self.det_height, self.det_width), np.uint8
)
self._roi[
border : self.det_height - border,
border : self.det_width - border,
] = 255
margin = int(round(OBSTACLE_EDGE_MARGIN * self.scale))
self._roi_bounds = (
border + margin,
border + margin,
self.det_width - border - margin,
self.det_height - border - margin,
)
self._pinned_size = None
if OBSTACLE_EXPECTED_SIZE:
long_side, short_side = OBSTACLE_EXPECTED_SIZE
self._pinned_size = (
float(max(long_side, short_side)),
float(min(long_side, short_side)),
)
# What "white" currently looks like to the camera, as a LAB
# (a*, b*) point - see Calibration.white_reference(). Starts
# at neutral (128, 128) and is kept updated by whoever drives
# this Detector (vision.py); NOT reset by reset() below,
# since it describes the room's current lighting, not
# anything about where the stickers are.
self._white_a = 128.0
self._white_b = 128.0
self.reset()
# ========================================================
# Public API
# ========================================================
def set_white_reference(self, reference_ab):
"""
Update what "white" currently looks like to the camera (see
Calibration.white_reference()). Falls back to LAB neutral
(128, 128) if None (not measured yet, or adaptive hue/white
learning is disabled) - the same as the old fixed behavior.
"""
if reference_ab is None:
self._white_a = 128.0
self._white_b = 128.0
else:
self._white_a, self._white_b = reference_ab
def reset(self):
self.tracks = []
self._next_id = 1
self._samples = deque(maxlen=30)
self._expected = self._pinned_size
self._established = self._pinned_size is not None
self._last_size_log = None
self._last_seen_any = time.monotonic()
self.obstacles = []
self._signature = None
# Last pass's raw candidates, for debugging/tuning.
self.last_candidates = []
def detect(self, frame):
"""
Run one detection pass on a screen-space frame at
DETECTION_SCALE and return the current list of confirmed
stickers (screen pixels). The SAME list object is returned
while nothing changed, so consumers can cheaply notice
changes with an identity check.
"""
if (
frame is None
or frame.ndim != 3
or frame.shape[0] != self.det_height
or frame.shape[1] != self.det_width
):
return self.obstacles
candidates = self._find_candidates(frame)
self.last_candidates = candidates
self._update_tracks(candidates)
self._forget_stale_size()
self._publish()
return self.obstacles
# ========================================================
# 1. Coarse blobs
# ========================================================
def _find_candidates(self, frame):
lightness, a_channel, b_channel = imaging.to_lab(frame)
lightness = cv2.GaussianBlur(lightness, (3, 3), 0)
chroma = np.clip(
imaging.chroma(
a_channel,
b_channel,
self._white_a,
self._white_b,
),
0,
255,
).astype(np.uint8)
# "Locally brighter than its own neighborhood" (not "above
# one fixed brightness for the whole frame") - see the
# module docstring and config.py's OBSTACLE_ADAPTIVE_* for
# why. A plain floor is still kept underneath it purely as a
# noise guard: without it, sensor noise inside a totally
# dark, sticker-free patch of screen can register as
# "locally brighter than its (equally dark) surroundings"
# even though nothing is actually there.
local = cv2.adaptiveThreshold(
lightness,
255,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY,
self._adaptive_block_size(),
-OBSTACLE_ADAPTIVE_OFFSET,
)
mask = local & cv2.inRange(
lightness,
OBSTACLE_MIN_ABSOLUTE_LIGHTNESS,
255,
)
mask &= cv2.inRange(
chroma,
0,
OBSTACLE_COARSE_CHROMA_MAX,
)
mask &= self._roi
blank = self._painted_blank_mask()
if blank is not None:
mask[blank > 0] = 0
mask = cv2.morphologyEx(
mask, cv2.MORPH_OPEN, self._open_kernel
)
mask = cv2.morphologyEx(
mask, cv2.MORPH_CLOSE, self._close_kernel
)
# Patches a glare highlight punched through the middle of a
# sticker, or a thin specular streak that splits one blob
# into two, without touching a real gap (e.g. a ball or a
# hand actually covering part of it), which stays open to
# the mask's edge and so is left alone.
mask = imaging.fill_holes(mask)
contours, _ = cv2.findContours(
mask,
cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE,
)
candidates = []
for contour in contours:
candidate = self._analyze(
contour,
lightness,
chroma,
blank,
)
if candidate is not None:
candidates.append(candidate)
return candidates
def _adaptive_block_size(self):
"""
Neighborhood size (detection-scale px, always odd) for the
local brightness threshold above. Tied to the sticker's own
short side (once known) so the neighborhood sampled around a
sticker is mostly black screen rather than the sticker's own
bright pixels, which would otherwise drag the local average
up and make the sticker fail to stand out against itself.
"""
if self._expected is not None:
short_side = min(self._expected)
else:
short_side = OBSTACLE_ADAPTIVE_DEFAULT_SHORT_SIDE
size = int(
round(
short_side
* self.scale
* OBSTACLE_ADAPTIVE_BLOCK_SPAN
)
)
size = max(15, size)
if size % 2 == 0:
size += 1
return size
def _painted_blank_mask(self):
"""
Mask of the exact patches the game painted white for the
stickers it is currently tracking. Those pixels are ignored
so that a REMOVED sticker's leftover painted white can never
be mistaken for a sticker (the camera is looking at the same
screen the game paints on).
"""
if not SHOW_DETECTED_OBSTACLES:
return None
# Keep blanking for two passes after painting stops, to
# cover the camera's own latency (it still shows the old
# painted patch for a moment).
confirmed = [
t
for t in self.tracks
if t.confirmed
and t.misses < OBSTACLE_STALE_MISSES + 2
]
if not confirmed:
return None
blank = np.zeros(
(self.det_height, self.det_width), np.uint8
)
for track in confirmed:
long_side, short_side = paint_dimensions(
track.long,
track.short,
DETECTED_OBSTACLE_PAINT_INSET_RATIO,
)
corners = obb_corners(
track.cx * self.scale,
track.cy * self.scale,
(long_side + 2 * _BLANK_MARGIN) * self.scale,
(short_side + 2 * _BLANK_MARGIN) * self.scale,
track.angle,
)
cv2.fillConvexPoly(
blank,
np.round(corners).astype(np.int32),
255,
)
return blank
# ========================================================
# 2 + 3. Refine one blob and measure it
# ========================================================
def _analyze(self, contour, lightness, chroma, blank):
area_px = cv2.contourArea(contour)
# Loose pre-gate on the coarse blob (the refined blob is
# gated exactly below).
loose_area = area_px / self._s2
if (
loose_area < MIN_OBSTACLE_AREA * 0.5
or loose_area > MAX_OBSTACLE_AREA * 1.5
):
return None
x, y, w, h = cv2.boundingRect(contour)
ring = max(3, int(round(0.15 * min(w, h))))
pad = ring + 3
x0 = max(0, x - pad)
y0 = max(0, y - pad)
x1 = min(self.det_width, x + w + pad)
y1 = min(self.det_height, y + h + pad)
crop_light = lightness[y0:y1, x0:x1]
crop_chroma = chroma[y0:y1, x0:x1]
crop_roi = self._roi[y0:y1, x0:x1]
crop_blank = (
blank[y0:y1, x0:x1] if blank is not None else None
)
shifted = contour - np.array([[[x0, y0]]], np.int32)
coarse_fill = np.zeros(crop_light.shape, np.uint8)
cv2.drawContours(coarse_fill, [shifted], -1, 255, -1)
inside = crop_light[coarse_fill > 0]
outside = crop_light[coarse_fill == 0]
if inside.size == 0 or outside.size < 8:
return None
plateau = float(np.percentile(inside, 75))
background = float(np.percentile(outside, 10))
# 2. Half-way brightness between this blob's own paper level
# and the dark screen beside it.
level = max(
background
+ OBSTACLE_EDGE_LEVEL * (plateau - background),
float(OBSTACLE_MIN_ABSOLUTE_LIGHTNESS),
)
refined = cv2.inRange(crop_light, level, 255)
refined &= cv2.inRange(
crop_chroma, 0, OBSTACLE_COARSE_CHROMA_MAX
)
refined &= crop_roi
if crop_blank is not None:
refined[crop_blank > 0] = 0
refined = cv2.morphologyEx(
refined, cv2.MORPH_OPEN, self._small_kernel
)
refined = cv2.morphologyEx(
refined, cv2.MORPH_CLOSE, self._small_kernel
)
found, _ = cv2.findContours(
refined,
cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE,
)
if not found:
return None
contour2 = max(found, key=cv2.contourArea)
area2_px = cv2.contourArea(contour2)
area = area2_px / self._s2
if area < MIN_OBSTACLE_AREA or area > MAX_OBSTACLE_AREA:
return None
# ---- 3. measurements ---------------------------------
rect = cv2.minAreaRect(contour2)
(rect_cx, rect_cy), (rect_w, rect_h), _raw_angle = rect
if rect_w <= 0 or rect_h <= 0:
return None
rectangularity = area2_px / (rect_w * rect_h)
# minAreaRect's own angle convention differs between
# OpenCV versions and jumps for near-square shapes, so the
# angle is measured from the box's actual long edge.
box = cv2.boxPoints(rect)
edge_a = box[1] - box[0]
edge_b = box[2] - box[1]
length_a = float(np.hypot(edge_a[0], edge_a[1]))
length_b = float(np.hypot(edge_b[0], edge_b[1]))
if length_a <= 0 or length_b <= 0:
return None
if length_a >= length_b:
long_edge, long_px, short_px = edge_a, length_a, length_b
else:
long_edge, long_px, short_px = edge_b, length_b, length_a
aspect = long_px / short_px
angle = math.degrees(
math.atan2(long_edge[1], long_edge[0])
)
angle = normalize_angle(
angle,
90.0 if aspect < _SQUARE_ASPECT else 180.0,
)
hull_area = cv2.contourArea(cv2.convexHull(contour2))
if hull_area <= 0:
return None
solidity = area2_px / hull_area
fill = np.zeros(crop_light.shape, np.uint8)
cv2.drawContours(fill, [contour2], -1, 255, -1)
inner = cv2.erode(fill, self._small_kernel)
if not inner.any():
inner = fill
mean_light = float(cv2.mean(crop_light, mask=inner)[0])
mean_chroma = float(cv2.mean(crop_chroma, mask=inner)[0])
_, std = cv2.meanStdDev(crop_light, mask=inner)
light_std = float(std[0][0])
# Ring just outside the blob must be dark (black screen).
ring_kernel = np.ones(
(2 * ring + 1, 2 * ring + 1), np.uint8
)
ring_mask = cv2.dilate(fill, ring_kernel) & ~fill
ring_mask &= crop_roi
if crop_blank is not None:
ring_mask[crop_blank > 0] = 0
ring_count = int(np.count_nonzero(ring_mask))
if ring_count >= 12:
dark = crop_light[ring_mask > 0] < (mean_light * 0.5)
dark_ratio = float(np.count_nonzero(dark)) / ring_count
else:
dark_ratio = 1.0
# Touching the edge of the search area means the sticker is
# (or may be) cut off - its true size is unknown.
bx, by, bw, bh = cv2.boundingRect(contour2)
left, top, right, bottom = self._roi_bounds
clipped = (
bx + x0 <= left
or by + y0 <= top
or bx + x0 + bw >= right
or by + y0 + bh >= bottom
)
# ---- back to real screen pixels ----------------------
inverse = 1.0 / self.scale
candidate = {
"cx": (rect_cx + x0) * inverse,
"cy": (rect_cy + y0) * inverse,
"angle": float(angle),
"long": long_px * inverse,
"short": short_px * inverse,
"area": float(area),
"aspect": float(aspect),
"rectangularity": float(rectangularity),
"solidity": float(solidity),
"brightness": mean_light,
"chroma": mean_chroma,
"std": light_std,
"dark_ratio": dark_ratio,
"clipped": clipped,
}
candidate["strict_ok"], candidate["snap"] = self._passes(
candidate, 1.0
)
candidate["relaxed_ok"], relaxed_snap = self._passes(
candidate, OBSTACLE_RELAXED_FACTOR
)
candidate["snap"] = candidate["snap"] or relaxed_snap
return candidate
# ========================================================
# Sticker tests (strict for new stickers, relaxed to keep
# established ones)
# ========================================================
def _passes(self, c, factor):
if c["clipped"]:
return False, False
if c["rectangularity"] < MIN_OBSTACLE_RECTANGULARITY * factor:
return False, False
if c["solidity"] < MIN_OBSTACLE_SOLIDITY * factor:
return False, False
if c["aspect"] > MAX_OBSTACLE_ASPECT_RATIO / factor:
return False, False
if c["brightness"] < MIN_OBSTACLE_BRIGHTNESS * factor:
return False, False
if c["chroma"] > MAX_OBSTACLE_CHROMA / factor:
return False, False
if c["std"] > MAX_OBSTACLE_BRIGHTNESS_STD / factor:
return False, False
if (
c["dark_ratio"]
< MIN_OBSTACLE_DARK_SURROUND_RATIO * factor
):
return False, False
return self._size_matches(
c["long"],
c["short"],
OBSTACLE_SIZE_TOLERANCE / factor,
)
def _size_matches(self, long_side, short_side, tolerance):
"""
(matches, is_single_sticker_size). With no learned size yet
anything passes; with a size learned from just one sighting
the check is deliberately loose (that one sighting could
itself have been a false positive).
"""
if self._expected is None:
return True, False
expected_long, expected_short = self._expected
if not self._established:
tolerance = max(tolerance, 0.45)
limit = 3 if OBSTACLE_ALLOW_MERGED_STICKERS else 1
for along in range(1, limit + 1):
for across in range(1, limit + 1):
if along * across > 6:
continue
target_a = along * expected_long
target_b = across * expected_short
target_long = max(target_a, target_b)
target_short = min(target_a, target_b)
if (
abs(long_side - target_long) / target_long
<= tolerance
and abs(short_side - target_short)
/ target_short
<= tolerance
):
single = along == 1 and across == 1
return True, single
return False, False
# ========================================================
# Learned sticker size
# ========================================================
def _learn_size(self, long_side, short_side):
if self._pinned_size is not None:
return
self._samples.append((long_side, short_side))
best = []
for a_long, a_short in self._samples:
members = [
(b_long, b_short)
for b_long, b_short in self._samples
if abs(b_long - a_long) / a_long <= 0.15
and abs(b_short - a_short) / a_short <= 0.15
]
if len(members) > len(best):
best = members
self._expected = (
float(np.median([m[0] for m in best])),
float(np.median([m[1] for m in best])),
)
self._established = len(best) >= 2
if self._established:
rounded = (
round(self._expected[0]),
round(self._expected[1]),
)
if (
self._last_size_log is None
or abs(rounded[0] - self._last_size_log[0]) > 4
or abs(rounded[1] - self._last_size_log[1]) > 4
):
self._last_size_log = rounded
print(
"[DETECT] sticker size learned: about "
f"{rounded[0]} x {rounded[1]} px. To pin it, set "
f"OBSTACLE_EXPECTED_SIZE = ({rounded[0]}, "
f"{rounded[1]}) in config.py."
)
def _forget_stale_size(self):
if self._pinned_size is not None:
return
now = time.monotonic()
if any(t.confirmed for t in self.tracks):
self._last_seen_any = now
return
if (
self._expected is not None
and now - self._last_seen_any
> OBSTACLE_SIZE_FORGET_SECONDS
):
self._samples.clear()
self._expected = None
self._established = False
self._last_size_log = None
# ========================================================
# 5. Tracking
# ========================================================
def _update_tracks(self, candidates):
pairs = []
for track_index, track in enumerate(self.tracks):
gate = min(
OBSTACLE_TEMPORAL_MATCH_DISTANCE,
max(25.0, 0.75 * track.short),
)
for cand_index, candidate in enumerate(candidates):
if not candidate["relaxed_ok"]:
continue
distance = math.hypot(
candidate["cx"] - track.cx,
candidate["cy"] - track.cy,
)
if distance <= gate:
pairs.append((distance, track_index, cand_index))
pairs.sort()
used_tracks = set()
used_candidates = set()
for _distance, track_index, cand_index in pairs:
if (
track_index in used_tracks
or cand_index in used_candidates
):
continue
used_tracks.add(track_index)
used_candidates.add(cand_index)
self._apply_hit(
self.tracks[track_index],
candidates[cand_index],
)
survivors = []
for track_index, track in enumerate(self.tracks):
if track_index not in used_tracks:
track.misses += 1
limit = (
OBSTACLE_MAX_MISSING_FRAMES
if track.confirmed
else OBSTACLE_PENDING_MAX_MISSING_FRAMES
)
if track.misses > limit:
continue
survivors.append(track)
self.tracks = survivors
# New stickers: strict tests only, and not on top of a
# sticker we already follow.
for cand_index, candidate in enumerate(candidates):
if (
cand_index in used_candidates
or not candidate["strict_ok"]
):
continue
overlaps = any(
math.hypot(
candidate["cx"] - t.cx,
candidate["cy"] - t.cy,
)
< 0.5 * min(t.short, candidate["short"])
for t in self.tracks
)
if overlaps:
continue
track = _Track(self._next_id, candidate)
self._next_id += 1
self.tracks.append(track)
self._maybe_confirm(track, candidate)
def _apply_hit(self, track, candidate):
track.misses = 0
track.hits += 1
long_side = candidate["long"]
short_side = candidate["short"]
if (
OBSTACLE_SNAP_TO_EXPECTED_SIZE
and self._established
and self._expected is not None
and candidate["snap"]
):
expected_long, expected_short = self._expected
if (
abs(long_side - expected_long) / expected_long <= 0.10
and abs(short_side - expected_short)
/ expected_short
<= 0.10
):
long_side, short_side = expected_long, expected_short
# Position: dead-band while nothing moves (no jitter for a
# ball resting on the sticker), fast catch-up when it does.
distance = math.hypot(
candidate["cx"] - track.cx,
candidate["cy"] - track.cy,
)
if distance > 0.8:
alpha = min(1.0, max(0.25, distance / 8.0))
track.cx += (candidate["cx"] - track.cx) * alpha
track.cy += (candidate["cy"] - track.cy) * alpha
period = (
90.0
if track.long / max(track.short, 1e-6) < _SQUARE_ASPECT
else 180.0
)
half = period / 2.0
angle_diff = abs(
((candidate["angle"] - track.angle + half) % period)
- half
)
if angle_diff > 0.6:
alpha = min(1.0, max(0.25, angle_diff / 4.0))
track.angle = blend_angle(
track.angle,
candidate["angle"],
alpha,
period,
)
track.long += (long_side - track.long) * 0.35
track.short += (short_side - track.short) * 0.35
self._maybe_confirm(track, candidate)
def _maybe_confirm(self, track, candidate):
if (
track.confirmed
or track.hits < OBSTACLE_REQUIRED_STABLE_FRAMES
):
return
track.confirmed = True
# Only single-sticker sightings teach the size (a merged
# pair would otherwise be learned as "the" sticker size).
if candidate["snap"] or self._expected is None:
self._learn_size(candidate["long"], candidate["short"])
# ========================================================
# Output
# ========================================================
def _publish(self):
confirmed = [t for t in self.tracks if t.confirmed]
signature = tuple(
(
t.id,
round(t.cx, 1),
round(t.cy, 1),
round(t.angle, 1),
round(t.long, 1),
round(t.short, 1),
t.misses < OBSTACLE_STALE_MISSES,
)
for t in confirmed
)
if signature == self._signature:
return
self._signature = signature
obstacles = []
for track in confirmed:
corners = obb_corners(
track.cx,
track.cy,
track.long,
track.short,
track.angle,
)
xs = [p[0] for p in corners]
ys = [p[1] for p in corners]
obstacles.append(
{
"id": track.id,
"center": (track.cx, track.cy),
"angle": track.angle,
"rect_long": track.long,
"rect_short": track.short,
# False once the sticker hasn't been seen for a
# few passes: rendering then stops painting it.
"painted": track.misses < OBSTACLE_STALE_MISSES,
"x": min(xs),
"y": min(ys),
"width": max(xs) - min(xs),
"height": max(ys) - min(ys),
}
)
self.obstacles = obstacles
# ========================================================
# Debug drawing
# ========================================================
def draw(self, frame, obstacles=None):
"""
Annotated copy of a screen-space detection frame: confirmed
stickers in green, rejected candidates in red (with the
first failing reason's numbers), for tuning on site.
"""
if obstacles is None:
obstacles = self.obstacles
output = frame.copy()
for candidate in self.last_candidates:
if candidate["relaxed_ok"]:
continue
corners = obb_corners(
candidate["cx"],
candidate["cy"],
candidate["long"],
candidate["short"],
candidate["angle"],
)
points = (
np.array(corners) * self.scale
).astype(np.int32).reshape((-1, 1, 2))
cv2.polylines(output, [points], True, (0, 0, 255), 1)
cv2.putText(
output,
f"r{candidate['rectangularity']:.2f} "
f"s{candidate['solidity']:.2f} "
f"d{candidate['dark_ratio']:.2f}",
(int(points[0][0][0]), int(points[0][0][1])),
cv2.FONT_HERSHEY_SIMPLEX,
0.4,
(0, 0, 255),
1,
cv2.LINE_AA,
)
for obstacle in obstacles:
obb = extract_obb(obstacle)
if obb is None:
continue
cx, cy, long_side, short_side, angle = obb
corners = obb_corners(cx, cy, long_side, short_side, angle)
points = (
np.array(corners) * self.scale
).astype(np.int32).reshape((-1, 1, 2))
cv2.polylines(output, [points], True, (0, 255, 0), 2)
cv2.circle(
output,
(int(cx * self.scale), int(cy * self.scale)),
4,
(0, 255, 0),
-1,
)
return output
# ========================================================
# Helpers
# ========================================================
def _kernel(self, size_screen_px):
size = max(1, int(round(size_screen_px * self.scale)))
if size % 2 == 0:
size += 1
return np.ones((size, size), np.uint8)
Reference in New Issue View Git Blame Copy Permalink