diff --git a/laser_detector.py b/laser_detector.py index 8e9a6db..24530fb 100644 --- a/laser_detector.py +++ b/laser_detector.py @@ -14,12 +14,16 @@ except ImportError: WIDTH = 640 HEIGHT = 480 THRESHOLD = 100 -RED_RATIO = 1 -SEARCH_RADIUS = 60 +RED_RATIO = 1.5 +SEARCH_RADIUS = 80 +TRACK_RADIUS = 30 +MIN_PIXELS = 3 +COARSE_STEP = 2 STABLE_COUNT = 2 +MAX_SKIP_FRAMES = 5 # Temporal smoothing -_EMA_ALPHA = 0.4 +_EMA_ALPHA = 0.35 _GATE_PX = 10 _FRAME_INTERVAL_MS = 50 @@ -83,44 +87,74 @@ def find_ellipse(img_cv, cx, cy, roi_r, th, ratio): return None +def is_red(r, g, b, th, ratio): + if r > th and r > g * ratio and r > b * ratio: + return True + if (r > 200 and g > 200 and b > 200 and r >= g and r >= b + and (r - g) > 10 and (r - b) > 10): + return True + return False + + def find_brightest_bytes(frame, cx, cy, roi_r, th, ratio): x1 = max(0, cx - roi_r) x2 = min(WIDTH, cx + roi_r) y1 = max(0, cy - roi_r) y2 = min(HEIGHT, cy + roi_r) data = frame.to_bytes() - rs, gs, bs = [], [], [] - xs, ys = [], [] - step = 2 - for y in range(y1, y2, step): - for x in range(x1, x2, step): + + best_score = 0 + best_x = (x1 + x2) // 2 + best_y = (y1 + y2) // 2 + found_any = False + for y in range(y1, y2, COARSE_STEP): + for x in range(x1, x2, COARSE_STEP): idx = (y * WIDTH + x) * 3 r = data[idx] g = data[idx + 1] b = data[idx + 2] - if (r > th and r > g * ratio and r > b * ratio) or \ - (r > 200 and g > 200 and b > 200 and r >= g and r >= b and (r - g) > 10 and (r - b) > 10): - rs.append(r) - gs.append(g) - bs.append(b) - xs.append(x) - ys.append(y) - if not rs: + if is_red(r, g, b, th, ratio): + score = r + g + b + dx = x - cx + dy = y - cy + dist_decay = max(0.5, 1.0 - ((dx * dx + dy * dy) ** 0.5 / roi_r) * 0.5) + score *= dist_decay + if score > best_score: + best_score = score + best_x = x + best_y = y + found_any = True + + if not found_any: return None - rs = np.array(rs, dtype=np.float64) - gs = np.array(gs, dtype=np.float64) - bs = np.array(bs, dtype=np.float64) - xs = np.array(xs, dtype=np.float64) - ys = np.array(ys, dtype=np.float64) - w = rs - np.maximum(gs, bs) - w = np.clip(w, 0, None) - w = w * w - total_w = w.sum() - if total_w < 1e-6: - return None - cx_f = (xs * w).sum() / total_w - cy_f = (ys * w).sum() / total_w - return (float(cx_f), float(cy_f)) + + sf = 4 + fx1 = max(x1, best_x - sf) + fx2 = min(x2, best_x + sf + 1) + fy1 = max(y1, best_y - sf) + fy2 = min(y2, best_y + sf + 1) + + sum_x = 0.0 + sum_y = 0.0 + total_w = 0.0 + count = 0 + for y in range(fy1, fy2): + for x in range(fx1, fx2): + idx = (y * WIDTH + x) * 3 + r = data[idx] + g = data[idx + 1] + b = data[idx + 2] + if is_red(r, g, b, th, ratio): + w = r + g + b + sum_x += x * w + sum_y += y * w + total_w += w + count += 1 + + if count < MIN_PIXELS: + return (float(best_x), float(best_y)) + + return (float(sum_x / total_w), float(sum_y / total_w)) def _ema_filter(pos, alpha=_EMA_ALPHA): @@ -151,6 +185,8 @@ def get_stable_laser_point(timeout_ms=15000, stable_count=STABLE_COUNT): stable = 0 start = time.ticks_ms() cx, cy = WIDTH // 2, HEIGHT // 2 + track_count = 0 + skip_count = 0 while True: if abs(time.ticks_diff(time.ticks_ms(), start)) > timeout_ms: _prev_smoothed = None @@ -159,19 +195,27 @@ def get_stable_laser_point(timeout_ms=15000, stable_count=STABLE_COUNT): if frame is None: time.sleep_ms(10) continue - pos_bright = find_brightest_bytes(frame, cx, cy, SEARCH_RADIUS, THRESHOLD, RED_RATIO) + + if track_count > 0 and _prev_smoothed is not None: + search_cx = int(_prev_smoothed[0]) + search_cy = int(_prev_smoothed[1]) + search_r = TRACK_RADIUS + else: + search_cx = cx + search_cy = cy + search_r = SEARCH_RADIUS + + pos_bright = find_brightest_bytes(frame, search_cx, search_cy, search_r, THRESHOLD, RED_RATIO) pos = pos_bright if _USE_CV: img_cv = image.image2cv(frame, False, False) - pos_ellipse = find_ellipse(img_cv, cx, cy, SEARCH_RADIUS, THRESHOLD, RED_RATIO) + pos_ellipse = find_ellipse(img_cv, search_cx, search_cy, search_r, THRESHOLD, RED_RATIO) if pos_ellipse is not None: pos = pos_ellipse + if pos is not None: - if not _gated(pos): - if logger_manager.logger: - logger_manager.logger.info(f"pos:{pos} gated,stable:{stable}") - time.sleep_ms(_FRAME_INTERVAL_MS) - continue + skip_count = 0 + track_count += 1 filtered = _ema_filter(pos) if last_raw is not None: dx = abs(filtered[0] - last_raw[0]) @@ -189,6 +233,16 @@ def get_stable_laser_point(timeout_ms=15000, stable_count=STABLE_COUNT): result = (int(filtered[0]), int(filtered[1])) _prev_smoothed = None return result + else: + skip_count += 1 + if logger_manager.logger: + logger_manager.logger.info(f"find_brightest_bytes None, skip={skip_count}, track={track_count}, search_center=({search_cx},{search_cy}), search_r={search_r}") + if skip_count > MAX_SKIP_FRAMES: + _prev_smoothed = None + track_count = 0 + stable = 0 + last_raw = None + time.sleep_ms(_FRAME_INTERVAL_MS) finally: _prev_smoothed = None