#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import math import time from pathlib import Path VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi", ".mkv", ".m4v"} def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="UAV video multi-object detection and tracking.") parser.add_argument("--input", default="test_video", help="Input video file or directory.") parser.add_argument("--output", default="outputs", help="Output directory.") parser.add_argument("--model", default="yolov8n.pt", help="YOLO model path, e.g. yolov8n.pt or custom.pt.") parser.add_argument("--imgsz", type=int, default=1280, help="YOLO inference image size.") parser.add_argument("--conf", type=float, default=0.25, help="Detection confidence threshold.") parser.add_argument("--det-iou", type=float, default=0.45, help="YOLO NMS IoU threshold.") parser.add_argument("--display-conf", type=float, default=0.25, help="Minimum tracked score to draw and write.") parser.add_argument("--display-min-hits", type=int, default=1, help="Visible frames required before drawing/writing a track.") parser.add_argument("--rescue-detections", action="store_true", help="Show high-confidence detections missed by the tracker.") parser.add_argument("--rescue-conf", type=float, default=0.45, help="Minimum score for tracker-missed detection rescue.") parser.add_argument("--rescue-max-missed", type=int, default=20, help="Frames to keep a rescue track alive while unmatched.") parser.add_argument("--track-nms-iou", type=float, default=0.55, help="Suppress duplicate visible tracks above this IoU.") parser.add_argument("--track-nms-center", type=float, default=0.65, help="Suppress duplicate tracks with very close centers.") parser.add_argument("--single-target", action="store_true", help="Keep one visible target and stabilize it as one ID.") parser.add_argument("--single-target-id", type=int, default=1, help="ID shown when --single-target is enabled.") parser.add_argument("--stitch-ids", action="store_true", help="Merge non-overlapping tracklets when motion continuity is plausible.") parser.add_argument("--stitch-buffer", type=int, default=180, help="Frames to keep disappeared IDs available for stitching.") parser.add_argument("--stitch-min-gap", type=int, default=4, help="Minimum absent frames before a new raw track may reuse an old ID.") parser.add_argument( "--tracker", choices=("bytetrack", "motion", "iou"), default="motion", help="Tracking backend. bytetrack is recommended for production-style ID continuity.", ) parser.add_argument("--track-iou", type=float, default=0.3, help="Tracker IoU matching threshold.") parser.add_argument("--max-missed", type=int, default=30, help="Frames to keep a lost track alive.") parser.add_argument("--reconnect-buffer", type=int, default=60, help="Frames after removal when a lost track may be reconnected.") parser.add_argument("--max-match-gap", type=int, default=60, help="Maximum missed frames still allowed for direct matching.") parser.add_argument("--min-hits", type=int, default=2, help="Matched frames required before showing a new track.") parser.add_argument("--new-track-conf", type=float, default=0.20, help="Minimum score allowed to start a new motion track.") parser.add_argument("--allow-near-new-tracks", action="store_true", help="Allow new tracks near existing tracks; useful for dense multi-UAV scenes.") parser.add_argument("--smooth-alpha", type=float, default=0.75, help="Motion tracker box smoothing weight for current detection.") parser.add_argument("--trail-length", type=int, default=80, help="Trajectory history length per target.") parser.add_argument("--bt-high", type=float, default=0.25, help="ByteTrack first association confidence threshold.") parser.add_argument("--bt-low", type=float, default=0.05, help="ByteTrack second association confidence threshold.") parser.add_argument("--bt-new", type=float, default=0.20, help="ByteTrack new-track confidence threshold.") parser.add_argument("--bt-match", type=float, default=0.88, help="ByteTrack matching threshold.") parser.add_argument("--logo", default="ykz.png", help="Logo image path to overlay at top-right. Empty disables it.") parser.add_argument( "--classes", default="0", help="Comma-separated class IDs to keep. Default 0 keeps only the custom drone class.", ) parser.add_argument("--display-class-name", default="drone", help="Class name shown in video labels and CSV.") parser.add_argument("--device", default=None, help="Ultralytics device string, e.g. cpu, 0, 0,1.") parser.add_argument("--limit-frames", type=int, default=0, help="Optional debug limit per video.") parser.add_argument("--show", action="store_true", help="Preview result while processing.") return parser.parse_args() def collect_videos(input_path: Path) -> list[Path]: if input_path.is_file(): return [input_path] videos = [p for p in sorted(input_path.iterdir()) if p.suffix.lower() in VIDEO_EXTENSIONS] if not videos: raise FileNotFoundError(f"No videos found under {input_path}") return videos def parse_classes(value: str | None) -> list[int] | None: if not value: return None return [int(item.strip()) for item in value.split(",") if item.strip()] def filter_visible_tracks(tracks, min_score: float, nms_iou: float, center_ratio: float): from src.tracker import box_iou candidates = [track for track in tracks if track.score >= min_score] candidates.sort(key=lambda track: (track.hits, track.score), reverse=True) kept = [] for track in candidates: duplicate = False tx, ty = track.center tw, th = track.size tdiag = max(1.0, (tw * tw + th * th) ** 0.5) for selected in kept: sx, sy = selected.center sw, sh = selected.size sdiag = max(1.0, (sw * sw + sh * sh) ** 0.5) center_distance = ((tx - sx) ** 2 + (ty - sy) ** 2) ** 0.5 if box_iou(track.xyxy, selected.xyxy) >= nms_iou: duplicate = True break if center_distance <= min(tdiag, sdiag) * center_ratio: duplicate = True break if not duplicate: kept.append(track) return kept def select_single_target_track(tracks, state: dict, target_id: int, trail_length: int): if not tracks: return [] if "center" in state: last_x, last_y = state["center"] selected = min( tracks, key=lambda track: ((track.center[0] - last_x) ** 2 + (track.center[1] - last_y) ** 2) ** 0.5 - track.score * 18.0 - min(track.hits, 12) * 1.5, ) else: selected = max(tracks, key=lambda track: (track.hits, track.score)) state["center"] = selected.center history = state.setdefault("history", []) history.append(selected.center) if len(history) > trail_length: del history[:-trail_length] selected.track_id = target_id selected.history = list(history) return [selected] def stitch_track_ids( tracks, state: dict, frame_index: int, frame_shape, buffer_frames: int, min_gap: int, trail_length: int, ): if not tracks: return [] height, width = frame_shape[:2] frame_diag = (width * width + height * height) ** 0.5 raw_to_canonical = state.setdefault("raw_to_canonical", {}) canonical_state = state.setdefault("canonical_state", {}) next_id = state.setdefault("next_id", 1) active_ids = set() stitched = [] ordered_tracks = sorted( tracks, key=lambda item: (item.track_id not in raw_to_canonical, -item.hits, -item.score), ) for track in ordered_tracks: raw_id = track.track_id cx, cy = track.center tw, th = track.size tdiag = max(1.0, (tw * tw + th * th) ** 0.5) if raw_id in raw_to_canonical: canonical_id = raw_to_canonical[raw_id] else: best_id = None best_cost = 1e6 for candidate_id, info in canonical_state.items(): if candidate_id in active_ids: continue gap = frame_index - info["frame"] if gap < min_gap or gap > buffer_frames: continue px = info["center"][0] + info["velocity"][0] * min(gap, 45) py = info["center"][1] + info["velocity"][1] * min(gap, 45) distance = ((cx - px) ** 2 + (cy - py) ** 2) ** 0.5 old_w, old_h = info["size"] size_cost = abs(math.log(tw / max(1.0, old_w))) + abs(math.log(th / max(1.0, old_h))) gate = max(60.0, min(frame_diag * 0.08, tdiag * (6.0 + gap * 0.08))) if distance > gate or size_cost > 1.35: continue cost = distance / gate + 0.18 * min(size_cost, 1.0) - track.score * 0.08 if cost < best_cost: best_cost = cost best_id = candidate_id if best_id is not None and best_cost < 0.92: canonical_id = best_id else: canonical_id = next_id next_id += 1 raw_to_canonical[raw_id] = canonical_id state["next_id"] = next_id if canonical_id in active_ids: canonical_id = next_id next_id += 1 raw_to_canonical[raw_id] = canonical_id state["next_id"] = next_id previous = canonical_state.get(canonical_id) if previous: gap = max(1, frame_index - previous["frame"]) velocity = ((cx - previous["center"][0]) / gap, (cy - previous["center"][1]) / gap) history = previous.get("history", []) else: velocity = (0.0, 0.0) history = [] history = history + [(cx, cy)] if len(history) > trail_length: history = history[-trail_length:] canonical_state[canonical_id] = { "frame": frame_index, "center": (cx, cy), "velocity": velocity, "size": (tw, th), "history": history, } track.track_id = canonical_id track.history = list(history) active_ids.add(canonical_id) stitched.append(track) state["canonical_state"] = { cid: info for cid, info in canonical_state.items() if frame_index - info["frame"] <= buffer_frames } return stitched def confirm_visible_tracks(tracks, state: dict, min_hits: int): if min_hits <= 1: return tracks counts = state.setdefault("counts", {}) confirmed = [] current_ids = {track.track_id for track in tracks} for track in tracks: counts[track.track_id] = counts.get(track.track_id, 0) + 1 if counts[track.track_id] >= min_hits: confirmed.append(track) state["counts"] = {track_id: count for track_id, count in counts.items() if track_id in current_ids or count >= min_hits} return confirmed def rescue_missing_detections(detections, visible_tracks, state: dict, min_score: float, max_missed: int): import numpy as np from src.tracker import Track, box_iou rescue_tracks = state.setdefault("tracks", []) if "next_id" not in state: state["next_id"] = max([track.track_id for track in visible_tracks] + [0]) + 1 for track in rescue_tracks: track.missed += 1 if track.history: cx, cy = track.center vx, vy = track.velocity track.xyxy = track.xyxy + np.array([vx, vy, vx, vy], dtype=float) track.history.append((int(cx + vx), int(cy + vy))) track.history = track.history[-80:] def near_track(detection, track, center_scale=1.1): dcx = float((detection.xyxy[0] + detection.xyxy[2]) / 2.0) dcy = float((detection.xyxy[1] + detection.xyxy[3]) / 2.0) tx, ty = track.center tw, th = track.size dw = max(1.0, float(detection.xyxy[2] - detection.xyxy[0])) dh = max(1.0, float(detection.xyxy[3] - detection.xyxy[1])) gate = max(36.0, min((tw * tw + th * th) ** 0.5, (dw * dw + dh * dh) ** 0.5) * center_scale) return ((dcx - tx) ** 2 + (dcy - ty) ** 2) ** 0.5 <= gate or box_iou(detection.xyxy, track.xyxy) >= 0.20 for detection in detections: if detection.score < min_score: continue if any(near_track(detection, track) for track in visible_tracks): continue best = None best_distance = 1e9 dcx = float((detection.xyxy[0] + detection.xyxy[2]) / 2.0) dcy = float((detection.xyxy[1] + detection.xyxy[3]) / 2.0) for track in rescue_tracks: if track.missed > max_missed: continue tx, ty = track.center distance = ((dcx - tx) ** 2 + (dcy - ty) ** 2) ** 0.5 if distance < best_distance and near_track(detection, track, center_scale=2.5): best = track best_distance = distance if best is None: best = Track( track_id=state["next_id"], xyxy=detection.xyxy.copy(), score=detection.score, class_id=detection.class_id, class_name=detection.class_name, ) state["next_id"] += 1 rescue_tracks.append(best) else: old_cx, old_cy = best.center best.xyxy = detection.xyxy.copy() best.score = detection.score best.missed = 0 best.hits += 1 best.velocity = (0.6 * best.velocity[0] + 0.4 * (dcx - old_cx), 0.6 * best.velocity[1] + 0.4 * (dcy - old_cy)) best.history.append(best.center) best.history = best.history[-80:] rescue_tracks = [track for track in rescue_tracks if track.missed <= max_missed] state["tracks"] = rescue_tracks visible_ids = {track.track_id for track in visible_tracks} additions = [ track for track in rescue_tracks if track.missed == 0 and track.track_id not in visible_ids and not any(near_track(DetectionLike(track), vt) for vt in visible_tracks) ] return visible_tracks + additions class DetectionLike: def __init__(self, track): self.xyxy = track.xyxy def process_video( video_path: Path, output_dir: Path, detector, args: argparse.Namespace, ) -> None: import cv2 from src.tracker import ByteTrackTracker, IouTracker, MotionTracker from src.visualization import draw_tracks, load_logo, overlay_logo cap = cv2.VideoCapture(str(video_path)) if not cap.isOpened(): raise RuntimeError(f"Cannot open video: {video_path}") fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) output_video = output_dir / f"{video_path.stem}_tracked.mp4" output_csv = output_dir / f"{video_path.stem}_tracks.csv" writer = cv2.VideoWriter( str(output_video), cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height), ) if not writer.isOpened(): raise RuntimeError(f"Cannot create output video: {output_video}") logo = load_logo(args.logo, width) if args.logo else None if args.tracker == "bytetrack": tracker = ByteTrackTracker( frame_rate=fps, track_high_thresh=args.bt_high, track_low_thresh=args.bt_low, new_track_thresh=args.bt_new, track_buffer=args.max_missed, match_thresh=args.bt_match, trail_length=args.trail_length, ) elif args.tracker == "iou": tracker = IouTracker( iou_threshold=args.track_iou, max_missed=args.max_missed, min_hits=args.min_hits, trail_length=args.trail_length, ) else: tracker = MotionTracker( max_missed=args.max_missed, min_hits=args.min_hits, new_track_thresh=args.new_track_conf, smooth_alpha=args.smooth_alpha, reconnect_buffer=args.reconnect_buffer, max_match_gap=args.max_match_gap, suppress_near_new_tracks=not args.allow_near_new_tracks, trail_length=args.trail_length, ) start = time.time() frame_index = 0 single_target_state: dict = {} stitch_state: dict = {} display_confirm_state: dict = {} rescue_state: dict = {} with output_csv.open("w", newline="", encoding="utf-8") as csv_file: csv_writer = csv.writer(csv_file) csv_writer.writerow(["frame", "track_id", "class_id", "class_name", "score", "x1", "y1", "x2", "y2"]) while True: ok, frame = cap.read() if not ok: break frame_index += 1 detections = detector.detect(frame) try: tracks = tracker.update(detections, frame.shape) except TypeError: tracks = tracker.update(detections) visible_tracks = filter_visible_tracks( tracks, min_score=args.display_conf, nms_iou=args.track_nms_iou, center_ratio=args.track_nms_center, ) if args.single_target: visible_tracks = select_single_target_track( visible_tracks, single_target_state, target_id=args.single_target_id, trail_length=args.trail_length, ) elif args.stitch_ids: visible_tracks = stitch_track_ids( visible_tracks, stitch_state, frame_index=frame_index, frame_shape=frame.shape, buffer_frames=args.stitch_buffer, min_gap=args.stitch_min_gap, trail_length=args.trail_length, ) visible_tracks = confirm_visible_tracks( visible_tracks, display_confirm_state, min_hits=args.display_min_hits, ) if args.rescue_detections: visible_tracks = rescue_missing_detections( detections, visible_tracks, rescue_state, min_score=args.rescue_conf, max_missed=args.rescue_max_missed, ) for track in visible_tracks: x1, y1, x2, y2 = track.xyxy csv_writer.writerow( [ frame_index, track.track_id, track.class_id, track.class_name, f"{track.score:.4f}", f"{x1:.2f}", f"{y1:.2f}", f"{x2:.2f}", f"{y2:.2f}", ] ) visualized = draw_tracks(frame, visible_tracks, frame_index) visualized = overlay_logo(visualized, logo) writer.write(visualized) if args.show: cv2.imshow("UAV detection and tracking", visualized) if cv2.waitKey(1) & 0xFF == ord("q"): break if frame_index % 30 == 0: elapsed = max(time.time() - start, 1e-6) print( f"[{video_path.name}] {frame_index}/{total_frames or '?'} frames, " f"{frame_index / elapsed:.2f} FPS" ) if args.limit_frames and frame_index >= args.limit_frames: break cap.release() writer.release() if args.show: cv2.destroyAllWindows() print(f"Saved video: {output_video}") print(f"Saved tracks: {output_csv}") def main() -> None: args = parse_args() try: from src.detector import YoloDetector except ImportError as exc: raise SystemExit( "Missing Python vision dependencies. Run `bash scripts/install_deps.sh` first." ) from exc input_path = Path(args.input) output_dir = Path(args.output) output_dir.mkdir(parents=True, exist_ok=True) detector = YoloDetector( model_path=args.model, conf=args.conf, iou=args.det_iou, classes=parse_classes(args.classes), display_name=args.display_class_name, device=args.device, imgsz=args.imgsz, ) for video_path in collect_videos(input_path): process_video(video_path, output_dir, detector, args) if __name__ == "__main__": main()