523 lines
21 KiB
Python
Executable File
523 lines
21 KiB
Python
Executable File
#!/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()
|