import cv2 import numpy as np import os import argparse # ------------------------------------------------------------ # Generate a video of blinking dots at the 10th roots of unity # ------------------------------------------------------------ def generate_target_video(output_path="target_roots.mp4", duration=10, fps=10, img_size=(640, 640)): """ Creates an MP4 video where bright dots appear at the 10th roots of unity. Dots blink randomly to aid detection. """ fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, img_size) # 10th roots of unity angles = np.linspace(0, 2*np.pi, 10, endpoint=False) radius = 0.4 * min(img_size) / 2 center = (img_size[0]//2, img_size[1]//2) root_points = [] for ang in angles: x = int(center[0] + radius * np.cos(ang)) y = int(center[1] + radius * np.sin(ang)) root_points.append((x, y)) total_frames = duration * fps for frame_idx in range(total_frames): img = np.zeros((img_size[1], img_size[0], 3), dtype=np.uint8) # Draw each dot with blinking intensity for i, (x, y) in enumerate(root_points): # Sinusoidal blinking intensity = int(128 + 127 * np.sin(2 * np.pi * 0.5 * frame_idx / fps + i)) cv2.circle(img, (x, y), 10, (intensity, intensity, intensity), -1) out.write(img) out.release() print(f"Generated video: {output_path} with {total_frames} frames") return output_path, root_points # ------------------------------------------------------------ # Detect dots in a video frame using blob detection # ------------------------------------------------------------ def detect_dots(frame, min_radius=4, max_radius=24): """ Returns list of (x, y) detected dot centers. Uses adaptive thresholding plus contour filtering. """ gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) blur = cv2.GaussianBlur(gray, (5, 5), 0) percentile_level = int(np.clip(np.percentile(blur, 99.2), 40, 245)) _, thresh_fixed = cv2.threshold(blur, percentile_level, 255, cv2.THRESH_BINARY) thresh_adaptive = cv2.adaptiveThreshold( blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 31, -8, ) thresh = cv2.bitwise_or(thresh_fixed, thresh_adaptive) kernel = np.ones((3, 3), np.uint8) thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel) thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel) contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) dots = [] for cnt in contours: area = cv2.contourArea(cnt) if area < 20: continue (x, y), radius = cv2.minEnclosingCircle(cnt) if min_radius <= radius <= max_radius: perimeter = cv2.arcLength(cnt, True) circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter > 1e-6 else 0.0 if circularity < 0.35: continue x_i, y_i = int(x), int(y) brightness = int(gray[max(0, y_i - 1):y_i + 2, max(0, x_i - 1):x_i + 2].mean()) dots.append((x_i, y_i, brightness, circularity)) dots.sort(key=lambda item: (item[2], item[3]), reverse=True) return [(x, y) for x, y, _, _ in dots[:20]] def deduplicate_points(points, min_distance=12): merged = [] for x, y in sorted(points): placed = False for idx, (mx, my, count) in enumerate(merged): if np.hypot(x - mx, y - my) < min_distance: new_count = count + 1 merged[idx] = ((mx * count + x) / new_count, (my * count + y) / new_count, new_count) placed = True break if not placed: merged.append((float(x), float(y), 1)) return [(int(round(x)), int(round(y)), count) for x, y, count in merged] def estimate_roots_from_points(points, frame_shape, expected_count=10): if not points: return [] center = np.mean(np.array(points, dtype=np.float32), axis=0) bins = [[] for _ in range(expected_count)] for point in points: angle = np.arctan2(point[1] - center[1], point[0] - center[0]) bin_idx = int(np.floor(((angle + np.pi) / (2 * np.pi)) * expected_count)) % expected_count bins[bin_idx].append(point) roots = [] for bucket in bins: if not bucket: continue bucket_arr = np.array(bucket, dtype=np.float32) roots.append(tuple(np.mean(bucket_arr, axis=0).astype(int))) return sort_roots_by_angle(roots, frame_shape) def detect_from_projection(projection): if projection.ndim == 2: projection = cv2.cvtColor(projection, cv2.COLOR_GRAY2BGR) return detect_dots(projection) def fit_circle_to_points(points): pts = np.array(points, dtype=np.float32) center = pts.mean(axis=0) radii = np.linalg.norm(pts - center, axis=1) radius = float(np.median(radii)) return center, radius def regularize_ring(points, frame_shape, expected_count=10): if len(points) < max(4, expected_count // 2): return [] center, radius = fit_circle_to_points(points) angles = np.sort(np.array([ np.arctan2(y - center[1], x - center[0]) for x, y in points ], dtype=np.float64)) slot_ids = np.arange(expected_count, dtype=np.float64) aligned_diffs = [] for idx, angle in enumerate(angles): slot = idx % expected_count aligned_diffs.append(np.exp(1j * (angle - 2 * np.pi * slot / expected_count))) phase = np.angle(np.mean(aligned_diffs)) if aligned_diffs else 0.0 roots = [] for slot in range(expected_count): theta = phase + 2 * np.pi * slot / expected_count x = int(round(center[0] + radius * np.cos(theta))) y = int(round(center[1] + radius * np.sin(theta))) roots.append((x, y)) return sort_roots_by_angle(roots, frame_shape) def best_circular_alignment(reference_points, candidate_points): if len(reference_points) != len(candidate_points): return candidate_points best_points = candidate_points best_error = float("inf") n = len(candidate_points) for shift in range(n): shifted = candidate_points[shift:] + candidate_points[:shift] error = np.mean([ np.linalg.norm(np.array(shifted[i]) - np.array(reference_points[i])) for i in range(n) ]) if error < best_error: best_error = error best_points = shifted return best_points def bootstrap_layout(video_path, expected_count=10, bootstrap_frames=40): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise RuntimeError(f"Cannot open video: {video_path}") pooled = [] accumulated = None max_projection = None frame_shape = None frame_count = 0 while frame_count < bootstrap_frames: ret, frame = cap.read() if not ret: break frame_shape = frame.shape gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if accumulated is None: accumulated = gray.astype(np.float32) max_projection = gray.copy() else: accumulated += gray.astype(np.float32) max_projection = np.maximum(max_projection, gray) pooled.extend(detect_dots(frame)) frame_count += 1 cap.release() if frame_shape is None or accumulated is None or max_projection is None: return [], None, None mean_frame = np.clip(accumulated / max(frame_count, 1), 0, 255).astype(np.uint8) pooled.extend(detect_from_projection(mean_frame)) pooled.extend(detect_from_projection(max_projection)) pooled = deduplicate_points(pooled) candidate_points = [(x, y) for x, y, _ in pooled] roots = estimate_roots_from_points(candidate_points, frame_shape, expected_count=expected_count) if len(roots) != expected_count: roots = regularize_ring(candidate_points, frame_shape, expected_count=expected_count) if len(roots) != expected_count: return [], frame_shape, None center, radius = fit_circle_to_points(roots) return roots, frame_shape, (center, radius) def refine_roots_from_expected(frame, expected_roots, search_radius=22): gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) refined = [] for x, y in expected_roots: x0 = max(0, x - search_radius) x1 = min(gray.shape[1], x + search_radius + 1) y0 = max(0, y - search_radius) y1 = min(gray.shape[0], y + search_radius + 1) patch = gray[y0:y1, x0:x1] if patch.size == 0: refined.append(None) continue local_max = int(patch.max()) if local_max < 8: refined.append(None) continue threshold = max(6, int(local_max * 0.45)) mask = patch >= threshold if not np.any(mask): refined.append(None) continue yy, xx = np.nonzero(mask) weights = patch[yy, xx].astype(np.float32) cx = x0 + np.average(xx, weights=weights) cy = y0 + np.average(yy, weights=weights) refined.append((int(round(cx)), int(round(cy)))) return refined def sort_roots_by_angle(points, frame_shape): center = (frame_shape[1] // 2, frame_shape[0] // 2) return sorted(points, key=lambda p: np.arctan2(p[1] - center[1], p[0] - center[0])) # ------------------------------------------------------------ # Main: run video root tracker # ------------------------------------------------------------ def main(): parser = argparse.ArgumentParser(description="Real‑time root extraction from video of blinking dots.") parser.add_argument("--video", type=str, default="", help="Path to existing video (if not provided, a synthetic one is generated).") parser.add_argument("--generate-only", action="store_true", help="Only generate the video and exit.") parser.add_argument("--fps", type=int, default=10, help="Frames per second for generation.") parser.add_argument("--duration", type=int, default=10, help="Duration in seconds for generated video.") parser.add_argument("--display", action="store_true", help="Show cv2.imshow preview window.") args = parser.parse_args() true_roots = None if args.video and os.path.exists(args.video): video_path = args.video else: # Generate synthetic video video_path, true_roots = generate_target_video(output_path="target_roots.mp4", duration=args.duration, fps=args.fps) if args.generate_only: print("Video generated. Exiting.") return inferred_roots, bootstrap_shape, circle = bootstrap_layout(video_path, expected_count=10) if len(inferred_roots) != 10: print("Could not infer a 10-point ring layout from the video.") return expected_roots = inferred_roots if circle is not None: center, radius = circle print( "Inferred ring:", f"center=({center[0]:.1f}, {center[1]:.1f})", f"radius={radius:.1f}", ) if true_roots is not None: true_roots = best_circular_alignment(expected_roots, true_roots) # Open video and process frame by frame cap = cv2.VideoCapture(video_path) if not cap.isOpened(): print("Cannot open video") return # For tracking, we maintain a moving average of detected root positions max_history = 5 estimated_roots = [] # list of full 10-root estimates all_detections = [] accumulated = None max_projection = None frame_shape = None root_tracks = [[] for _ in range(len(expected_roots))] display_enabled = args.display and ( bool(os.environ.get("DISPLAY")) or bool(os.environ.get("WAYLAND_DISPLAY")) ) frame_idx = 0 while True: ret, frame = cap.read() if not ret: break frame_shape = frame.shape gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if accumulated is None: accumulated = gray.astype(np.float32) max_projection = gray.copy() else: accumulated += gray.astype(np.float32) max_projection = np.maximum(max_projection, gray) # Guided local refinement is much more stable on blurred/noisy videos. detected = refine_roots_from_expected(frame, expected_roots) detected_points = [point for point in detected if point is not None] if len(detected_points) < 10: blob_detected = detect_dots(frame) if blob_detected: all_detections.append(sort_roots_by_angle(blob_detected, frame.shape)) for idx, point in enumerate(detected): if point is not None: root_tracks[idx].append(point) root_tracks[idx] = root_tracks[idx][-max_history:] tracked_estimate = [] for idx, history in enumerate(root_tracks): if history: tracked_estimate.append(tuple(np.mean(np.array(history), axis=0).astype(int))) else: tracked_estimate.append(expected_roots[idx]) if sum(bool(history) for history in root_tracks) == len(expected_roots): estimated_roots.append(tracked_estimate) estimated_roots = estimated_roots[-max_history:] if detected_points: all_detections.append(sort_roots_by_angle(detected_points, frame.shape)) # Compute current estimate as average of last few frames if estimated_roots: avg_roots = np.mean(estimated_roots, axis=0).astype(int) else: avg_roots = np.array(tracked_estimate, dtype=int) # Visualisation display = frame.copy() # Draw inferred root slots (red circles) for (x, y) in expected_roots: cv2.circle(display, (x, y), 8, (0, 0, 255), 2) # Draw detected dots (green) for point in detected_points: x, y = point cv2.circle(display, (x, y), 6, (0, 255, 0), -1) # Draw estimated average (blue) for (x, y) in avg_roots: cv2.circle(display, (x, y), 4, (255, 0, 0), -1) cv2.putText(display, f"Frame {frame_idx} | Detected: {len(detected_points)}/10", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2) if display_enabled: cv2.imshow("Root Tracker", display) if cv2.waitKey(1) & 0xFF == ord('q'): break frame_idx += 1 cap.release() if display_enabled: cv2.destroyAllWindows() if len(estimated_roots) < max_history and accumulated is not None and frame_shape is not None: projection_candidates = [] mean_frame = np.clip(accumulated / max(frame_idx, 1), 0, 255).astype(np.uint8) mean_guided = refine_roots_from_expected(cv2.cvtColor(mean_frame, cv2.COLOR_GRAY2BGR), expected_roots) max_guided = refine_roots_from_expected(cv2.cvtColor(max_projection, cv2.COLOR_GRAY2BGR), expected_roots) projection_candidates.extend([point for point in mean_guided if point is not None]) projection_candidates.extend([point for point in max_guided if point is not None]) projection_candidates.extend(detect_from_projection(mean_frame)) projection_candidates.extend(detect_from_projection(max_projection)) projection_candidates = deduplicate_points(projection_candidates) recovered = estimate_roots_from_points([(x, y) for x, y, _ in projection_candidates], frame_shape) if len(recovered) == 10: estimated_roots.append(recovered) if len(estimated_roots) < max_history and all_detections and frame_shape is not None: pooled = [point for detection in all_detections for point in detection] pooled = deduplicate_points(pooled) consensus = estimate_roots_from_points([(x, y) for x, y, _ in pooled], frame_shape) if len(consensus) == 10: estimated_roots.append(consensus) # Compute final accuracy: average distance between estimated and true roots if not estimated_roots and frame_shape is not None and all(root_tracks): estimated_roots.append([ tuple(np.mean(np.array(history), axis=0).astype(int)) for history in root_tracks ]) if estimated_roots and frame_shape is not None: last_avg = np.mean(estimated_roots, axis=0) slot_errors = [np.linalg.norm(np.array(est) - np.array(expected_roots[i])) for i, est in enumerate(last_avg)] print(f"\nFinal average distance to inferred slots: {np.mean(slot_errors):.2f} pixels") if true_roots is not None: true_errors = [np.linalg.norm(np.array(est) - np.array(true_roots[i])) for i, est in enumerate(last_avg)] print(f"Final average distance to true roots: {np.mean(true_errors):.2f} pixels") if len(estimated_roots) < max_history: print(f"Used {len(estimated_roots)} complete detections plus frame accumulation fallback.") print(f"Recovered {len(last_avg)}/10 root positions.") else: print("Not enough detections to compute accuracy.") if __name__ == "__main__": main()