Initial project import
commit
4340aa0790
@ -0,0 +1,32 @@
|
|||||||
|
__pycache__/
|
||||||
|
tests/__pycache__/
|
||||||
|
.pytest_cache/
|
||||||
|
.venv/
|
||||||
|
venv/
|
||||||
|
|
||||||
|
cvat/
|
||||||
|
dataset_prep/
|
||||||
|
analysis_frames/
|
||||||
|
__debug_frames/
|
||||||
|
|
||||||
|
*.mp4
|
||||||
|
*.zip
|
||||||
|
*.csv
|
||||||
|
*.json
|
||||||
|
|
||||||
|
out_infer*.mp4
|
||||||
|
track_log_*.csv
|
||||||
|
track_summary_*.json
|
||||||
|
guidance_override*.csv
|
||||||
|
guidance_override*.json
|
||||||
|
smoke_*.csv
|
||||||
|
smoke_*.json
|
||||||
|
smoke_*.mp4
|
||||||
|
segment_*.csv
|
||||||
|
segment_*.json
|
||||||
|
segment_*.mp4
|
||||||
|
rollback_*.csv
|
||||||
|
|
||||||
|
!best.pt
|
||||||
|
!docs/superpowers/specs/2026-06-29-docker-offline-gpu-proto-udp-design.md
|
||||||
|
!docs/superpowers/plans/2026-06-29-docker-offline-gpu-proto-udp-implementation.md
|
||||||
@ -0,0 +1,51 @@
|
|||||||
|
FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04
|
||||||
|
|
||||||
|
ARG DEBIAN_FRONTEND=noninteractive
|
||||||
|
ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
|
||||||
|
|
||||||
|
ENV PYTHONUNBUFFERED=1 \
|
||||||
|
PIP_NO_CACHE_DIR=1 \
|
||||||
|
TZ=UTC
|
||||||
|
|
||||||
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||||
|
python3 \
|
||||||
|
python3-pip \
|
||||||
|
python3-venv \
|
||||||
|
ffmpeg \
|
||||||
|
libgl1 \
|
||||||
|
libglib2.0-0 \
|
||||||
|
libgomp1 \
|
||||||
|
libsm6 \
|
||||||
|
libxext6 \
|
||||||
|
libxrender1 \
|
||||||
|
v4l-utils \
|
||||||
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
COPY requirements-docker.txt /tmp/requirements-docker.txt
|
||||||
|
|
||||||
|
RUN python3 -m pip install --upgrade pip setuptools wheel && \
|
||||||
|
python3 -m pip install --index-url ${PYTORCH_INDEX_URL} torch torchvision && \
|
||||||
|
python3 -m pip install -r /tmp/requirements-docker.txt
|
||||||
|
|
||||||
|
COPY . /app
|
||||||
|
|
||||||
|
RUN mkdir -p /data/out /data/guidance /data/autopilot /app/docker && \
|
||||||
|
chmod +x /app/docker/entrypoint.sh
|
||||||
|
|
||||||
|
ENV FPV_MODEL_PATH=/app/best.pt \
|
||||||
|
FPV_SHOW_OUTPUT=0 \
|
||||||
|
FPV_SAVE_INFER_VIDEO=1 \
|
||||||
|
FPV_OUT_VIDEO_PATH=/data/out/out_infer.mp4 \
|
||||||
|
FPV_GUIDANCE_EXPORT_ENABLE=1 \
|
||||||
|
FPV_GUIDANCE_EXPORT_PATH=/data/guidance/guidance_state.json \
|
||||||
|
FPV_AUTOPILOT_ENABLE=1 \
|
||||||
|
FPV_AUTOPILOT_BACKEND=json \
|
||||||
|
FPV_AUTOPILOT_JSON_PATH=/data/autopilot/autopilot_cmd.json \
|
||||||
|
FPV_PROTO_UDP_ENABLE=0 \
|
||||||
|
FPV_PROTO_UDP_HOST=127.0.0.1 \
|
||||||
|
FPV_PROTO_UDP_PORT=5005
|
||||||
|
|
||||||
|
ENTRYPOINT ["/app/docker/entrypoint.sh"]
|
||||||
|
CMD ["python3", "main.py"]
|
||||||
@ -0,0 +1,70 @@
|
|||||||
|
import cv2
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
# 0 — встроенная USB/веб-камера
|
||||||
|
# Для RTSP можно указать строку:
|
||||||
|
# camera_url = "rtsp://login:password@192.168.1.100:554/Streaming/Channels/101"
|
||||||
|
camera_url = 0
|
||||||
|
|
||||||
|
cap = cv2.VideoCapture(camera_url)
|
||||||
|
|
||||||
|
if not cap.isOpened():
|
||||||
|
print("Ошибка: не удалось открыть камеру")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Настройки камеры
|
||||||
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||||
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||||
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||||
|
|
||||||
|
if fps <= 0 or fps > 120:
|
||||||
|
fps = 25
|
||||||
|
|
||||||
|
print(f"Камера открыта: {width}x{height}, FPS={fps}")
|
||||||
|
|
||||||
|
filename = datetime.now().strftime("record_%Y-%m-%d_%H-%M-%S.mp4")
|
||||||
|
|
||||||
|
# Кодек для MP4
|
||||||
|
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||||
|
|
||||||
|
writer = cv2.VideoWriter(
|
||||||
|
filename,
|
||||||
|
fourcc,
|
||||||
|
fps,
|
||||||
|
(width, height)
|
||||||
|
)
|
||||||
|
|
||||||
|
if not writer.isOpened():
|
||||||
|
print("Ошибка: не удалось создать файл записи")
|
||||||
|
cap.release()
|
||||||
|
return
|
||||||
|
|
||||||
|
print(f"Запись началась: {filename}")
|
||||||
|
print("Нажми Q для остановки")
|
||||||
|
|
||||||
|
while True:
|
||||||
|
ret, frame = cap.read()
|
||||||
|
|
||||||
|
if not ret:
|
||||||
|
print("Ошибка: кадр не получен")
|
||||||
|
break
|
||||||
|
|
||||||
|
writer.write(frame)
|
||||||
|
|
||||||
|
cv2.imshow("Camera recording", frame)
|
||||||
|
|
||||||
|
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||||
|
print("Остановка записи")
|
||||||
|
break
|
||||||
|
|
||||||
|
cap.release()
|
||||||
|
writer.release()
|
||||||
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
print(f"Видео сохранено: {filename}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,342 @@
|
|||||||
|
import time
|
||||||
|
import threading
|
||||||
|
from collections import deque
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from config import *
|
||||||
|
from helpers import clip_box, box_center, box_area
|
||||||
|
|
||||||
|
|
||||||
|
class AutoGazeROIWorker:
|
||||||
|
def __init__(self):
|
||||||
|
self.enabled = bool(AUTOGAZE_ENABLE)
|
||||||
|
self.ready = False
|
||||||
|
self.last_error = ""
|
||||||
|
|
||||||
|
self._autogaze_cls = None
|
||||||
|
self._processor_cls = None
|
||||||
|
if self.enabled:
|
||||||
|
try:
|
||||||
|
from autogaze.models.autogaze import AutoGaze, AutoGazeImageProcessor
|
||||||
|
self._autogaze_cls = AutoGaze
|
||||||
|
self._processor_cls = AutoGazeImageProcessor
|
||||||
|
except Exception as e:
|
||||||
|
self.last_error = f"import failed: {e}"
|
||||||
|
|
||||||
|
self.device = torch.device(f"cuda:{int(DEVICE)}") if torch.cuda.is_available() else torch.device("cpu")
|
||||||
|
|
||||||
|
self.clip_len = int(max(1, AUTOGAZE_CLIP_LEN))
|
||||||
|
self.update_every = int(max(1, AUTOGAZE_UPDATE_EVERY))
|
||||||
|
self.result_ttl = int(max(1, AUTOGAZE_RESULT_TTL))
|
||||||
|
self.topk = int(max(1, AUTOGAZE_TOPK))
|
||||||
|
self.min_active_cells = int(max(1, AUTOGAZE_MIN_ACTIVE_CELLS))
|
||||||
|
self.margin_cells = int(max(0, AUTOGAZE_ROI_MARGIN_CELLS))
|
||||||
|
self.min_side = float(max(8.0, AUTOGAZE_ROI_MIN_SIDE))
|
||||||
|
self.gazing_ratio = float(min(1.0, max(0.01, AUTOGAZE_GAZING_RATIO)))
|
||||||
|
self.task_loss_requirement = float(min(1.0, max(0.0, AUTOGAZE_TASK_LOSS_REQUIREMENT)))
|
||||||
|
|
||||||
|
self.model = None
|
||||||
|
self.processor = None
|
||||||
|
|
||||||
|
self.frame_buf = deque(maxlen=self.clip_len)
|
||||||
|
self.req = deque(maxlen=1)
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
self.running = False
|
||||||
|
self.thread = threading.Thread(target=self._loop, daemon=True)
|
||||||
|
self.latest = None
|
||||||
|
self.last_submit_frame = -10**9
|
||||||
|
|
||||||
|
def status_line(self):
|
||||||
|
if not self.enabled:
|
||||||
|
return "AutoGaze disabled by config"
|
||||||
|
if self.ready:
|
||||||
|
return f"AutoGaze ready: model={AUTOGAZE_MODEL_ID} device={self.device}"
|
||||||
|
if self.last_error:
|
||||||
|
return f"AutoGaze unavailable: {self.last_error}"
|
||||||
|
return "AutoGaze unavailable"
|
||||||
|
|
||||||
|
def start(self):
|
||||||
|
if not self.enabled:
|
||||||
|
return
|
||||||
|
if self._autogaze_cls is None or self._processor_cls is None:
|
||||||
|
return
|
||||||
|
if not self._load_model():
|
||||||
|
return
|
||||||
|
if self.running:
|
||||||
|
return
|
||||||
|
self.running = True
|
||||||
|
self.thread.start()
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
self.running = False
|
||||||
|
if self.thread.is_alive():
|
||||||
|
self.thread.join(timeout=1.0)
|
||||||
|
|
||||||
|
def submit(self, frame_eff_bgr, frame_id):
|
||||||
|
if (not self.running) or (not self.ready):
|
||||||
|
return
|
||||||
|
|
||||||
|
frame_copy = np.ascontiguousarray(frame_eff_bgr.copy())
|
||||||
|
fid = int(frame_id)
|
||||||
|
with self.lock:
|
||||||
|
self.frame_buf.append((fid, frame_copy))
|
||||||
|
if len(self.frame_buf) < self.clip_len:
|
||||||
|
return
|
||||||
|
if (fid - self.last_submit_frame) < self.update_every:
|
||||||
|
return
|
||||||
|
items = list(self.frame_buf)
|
||||||
|
clip = [f for _, f in items]
|
||||||
|
end_fid = int(items[-1][0])
|
||||||
|
self.req.append((clip, end_fid))
|
||||||
|
self.last_submit_frame = fid
|
||||||
|
|
||||||
|
def get_latest(self, current_frame_id, frame_w, frame_h, pred_ref_box=None):
|
||||||
|
info = {
|
||||||
|
"enabled": int(self.enabled),
|
||||||
|
"ready": int(self.ready),
|
||||||
|
"status": self.status_line(),
|
||||||
|
"stale": 1,
|
||||||
|
"age": -1,
|
||||||
|
"active_cells": 0,
|
||||||
|
"infer_ms": 0.0,
|
||||||
|
}
|
||||||
|
if not self.ready:
|
||||||
|
return None, [], info
|
||||||
|
|
||||||
|
with self.lock:
|
||||||
|
latest = None if self.latest is None else dict(self.latest)
|
||||||
|
|
||||||
|
if latest is None:
|
||||||
|
return None, [], info
|
||||||
|
|
||||||
|
age = int(current_frame_id) - int(latest["frame_id"])
|
||||||
|
info["age"] = age
|
||||||
|
info["active_cells"] = int(latest.get("active_cells", 0))
|
||||||
|
info["infer_ms"] = float(latest.get("infer_ms", 0.0))
|
||||||
|
if age > self.result_ttl:
|
||||||
|
return None, [], info
|
||||||
|
|
||||||
|
rois = []
|
||||||
|
for r in latest.get("rois", []):
|
||||||
|
rois.append(clip_box(np.array(r, dtype=np.float32), frame_w, frame_h))
|
||||||
|
|
||||||
|
roi = self._pick_roi(rois, pred_ref_box)
|
||||||
|
info["stale"] = 0
|
||||||
|
return roi, rois, info
|
||||||
|
|
||||||
|
def _load_model(self):
|
||||||
|
if self.ready and self.model is not None and self.processor is not None:
|
||||||
|
return True
|
||||||
|
|
||||||
|
try:
|
||||||
|
self.processor = self._processor_cls.from_pretrained(
|
||||||
|
AUTOGAZE_MODEL_ID,
|
||||||
|
local_files_only=bool(AUTOGAZE_LOCAL_FILES_ONLY),
|
||||||
|
)
|
||||||
|
self.model = self._autogaze_cls.from_pretrained(
|
||||||
|
AUTOGAZE_MODEL_ID,
|
||||||
|
use_flash_attn=bool(AUTOGAZE_USE_FLASH_ATTN),
|
||||||
|
local_files_only=bool(AUTOGAZE_LOCAL_FILES_ONLY),
|
||||||
|
)
|
||||||
|
self.model.to(self.device)
|
||||||
|
self.model.eval()
|
||||||
|
self.ready = True
|
||||||
|
self.last_error = ""
|
||||||
|
return True
|
||||||
|
except Exception as e:
|
||||||
|
self.model = None
|
||||||
|
self.processor = None
|
||||||
|
self.ready = False
|
||||||
|
self.last_error = f"load failed: {e}"
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _loop(self):
|
||||||
|
while self.running:
|
||||||
|
item = None
|
||||||
|
with self.lock:
|
||||||
|
if self.req:
|
||||||
|
item = self.req.pop()
|
||||||
|
self.req.clear()
|
||||||
|
|
||||||
|
if item is None:
|
||||||
|
time.sleep(0.001)
|
||||||
|
continue
|
||||||
|
|
||||||
|
clip_frames, frame_id = item
|
||||||
|
infer_ms = 0.0
|
||||||
|
rois = []
|
||||||
|
active_cells = 0
|
||||||
|
try:
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
rois, active_cells = self._infer_rois(clip_frames)
|
||||||
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
||||||
|
except Exception as e:
|
||||||
|
self.last_error = f"infer failed: {e}"
|
||||||
|
rois = []
|
||||||
|
active_cells = 0
|
||||||
|
infer_ms = 0.0
|
||||||
|
|
||||||
|
with self.lock:
|
||||||
|
self.latest = {
|
||||||
|
"frame_id": int(frame_id),
|
||||||
|
"rois": rois,
|
||||||
|
"active_cells": int(active_cells),
|
||||||
|
"infer_ms": float(infer_ms),
|
||||||
|
}
|
||||||
|
|
||||||
|
def _infer_rois(self, clip_frames_bgr):
|
||||||
|
if (self.model is None) or (self.processor is None):
|
||||||
|
return [], 0
|
||||||
|
if len(clip_frames_bgr) < self.clip_len:
|
||||||
|
return [], 0
|
||||||
|
|
||||||
|
h, w = clip_frames_bgr[-1].shape[:2]
|
||||||
|
clip_rgb = [cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in clip_frames_bgr[-self.clip_len:]]
|
||||||
|
processed = self.processor([clip_rgb], return_tensors="pt")
|
||||||
|
pixel_values = processed.get("pixel_values", None)
|
||||||
|
if pixel_values is None:
|
||||||
|
return [], 0
|
||||||
|
if isinstance(pixel_values, list):
|
||||||
|
pixel_values = torch.as_tensor(np.asarray(pixel_values))
|
||||||
|
if pixel_values.ndim == 4:
|
||||||
|
pixel_values = pixel_values.unsqueeze(0)
|
||||||
|
|
||||||
|
video = pixel_values.to(self.device)
|
||||||
|
if video.dtype != torch.float32 and self.device.type == "cpu":
|
||||||
|
video = video.float()
|
||||||
|
|
||||||
|
with torch.inference_mode():
|
||||||
|
outputs = self.model(
|
||||||
|
{"video": video},
|
||||||
|
gazing_ratio=self.gazing_ratio,
|
||||||
|
task_loss_requirement=self.task_loss_requirement,
|
||||||
|
generate_only=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
grid = self._extract_fine_grid(outputs)
|
||||||
|
if grid is None:
|
||||||
|
return [], 0
|
||||||
|
|
||||||
|
rois, active_cells = self._grid_to_rois(grid, w, h)
|
||||||
|
return rois, active_cells
|
||||||
|
|
||||||
|
def _extract_fine_grid(self, outputs):
|
||||||
|
gazing_mask = outputs.get("gazing_mask", None)
|
||||||
|
if gazing_mask is None or len(gazing_mask) == 0:
|
||||||
|
return None
|
||||||
|
|
||||||
|
fine = gazing_mask[-1]
|
||||||
|
if isinstance(fine, torch.Tensor):
|
||||||
|
arr = fine.detach().float().cpu().numpy()
|
||||||
|
else:
|
||||||
|
arr = np.asarray(fine)
|
||||||
|
|
||||||
|
if arr.ndim == 2:
|
||||||
|
arr = arr[None, ...]
|
||||||
|
if arr.ndim != 3 or arr.shape[0] <= 0 or arr.shape[1] <= 0:
|
||||||
|
return None
|
||||||
|
|
||||||
|
vec = arr[0, -1]
|
||||||
|
if vec.size <= 0:
|
||||||
|
return None
|
||||||
|
|
||||||
|
side = int(round(np.sqrt(float(vec.size))))
|
||||||
|
if side * side != int(vec.size):
|
||||||
|
return None
|
||||||
|
|
||||||
|
grid = vec.reshape(side, side)
|
||||||
|
return (grid > 0.5).astype(np.uint8)
|
||||||
|
|
||||||
|
def _grid_to_rois(self, grid, frame_w, frame_h):
|
||||||
|
if grid is None or grid.size == 0:
|
||||||
|
return [], 0
|
||||||
|
|
||||||
|
mask = (grid > 0).astype(np.uint8)
|
||||||
|
active_cells = int(mask.sum())
|
||||||
|
if active_cells < self.min_active_cells:
|
||||||
|
return [], active_cells
|
||||||
|
|
||||||
|
gh, gw = mask.shape
|
||||||
|
cell_w = float(frame_w) / float(max(1, gw))
|
||||||
|
cell_h = float(frame_h) / float(max(1, gh))
|
||||||
|
|
||||||
|
candidates = []
|
||||||
|
n_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=4)
|
||||||
|
for lbl in range(1, int(n_labels)):
|
||||||
|
x = int(stats[lbl, cv2.CC_STAT_LEFT])
|
||||||
|
y = int(stats[lbl, cv2.CC_STAT_TOP])
|
||||||
|
ww = int(stats[lbl, cv2.CC_STAT_WIDTH])
|
||||||
|
hh = int(stats[lbl, cv2.CC_STAT_HEIGHT])
|
||||||
|
area = int(stats[lbl, cv2.CC_STAT_AREA])
|
||||||
|
if area < self.min_active_cells:
|
||||||
|
continue
|
||||||
|
|
||||||
|
gx1 = max(0, x - self.margin_cells)
|
||||||
|
gy1 = max(0, y - self.margin_cells)
|
||||||
|
gx2 = min(gw, x + ww + self.margin_cells)
|
||||||
|
gy2 = min(gh, y + hh + self.margin_cells)
|
||||||
|
roi = np.array(
|
||||||
|
[gx1 * cell_w, gy1 * cell_h, gx2 * cell_w, gy2 * cell_h],
|
||||||
|
dtype=np.float32,
|
||||||
|
)
|
||||||
|
roi = self._ensure_min_side(roi, frame_w, frame_h)
|
||||||
|
candidates.append((area, roi))
|
||||||
|
|
||||||
|
if not candidates:
|
||||||
|
ys, xs = np.where(mask > 0)
|
||||||
|
if xs.size <= 0:
|
||||||
|
return [], active_cells
|
||||||
|
gx1 = max(0, int(xs.min()) - self.margin_cells)
|
||||||
|
gy1 = max(0, int(ys.min()) - self.margin_cells)
|
||||||
|
gx2 = min(gw, int(xs.max()) + 1 + self.margin_cells)
|
||||||
|
gy2 = min(gh, int(ys.max()) + 1 + self.margin_cells)
|
||||||
|
roi = np.array(
|
||||||
|
[gx1 * cell_w, gy1 * cell_h, gx2 * cell_w, gy2 * cell_h],
|
||||||
|
dtype=np.float32,
|
||||||
|
)
|
||||||
|
roi = self._ensure_min_side(roi, frame_w, frame_h)
|
||||||
|
candidates.append((active_cells, roi))
|
||||||
|
|
||||||
|
candidates.sort(key=lambda x: x[0], reverse=True)
|
||||||
|
rois = [np.array(r, dtype=np.float32) for _, r in candidates[: self.topk]]
|
||||||
|
return rois, active_cells
|
||||||
|
|
||||||
|
def _ensure_min_side(self, roi, frame_w, frame_h):
|
||||||
|
roi = clip_box(roi, frame_w, frame_h)
|
||||||
|
rw = float(roi[2] - roi[0])
|
||||||
|
rh = float(roi[3] - roi[1])
|
||||||
|
if rw >= self.min_side and rh >= self.min_side:
|
||||||
|
return roi
|
||||||
|
|
||||||
|
cx = 0.5 * (float(roi[0]) + float(roi[2]))
|
||||||
|
cy = 0.5 * (float(roi[1]) + float(roi[3]))
|
||||||
|
side = max(self.min_side, rw, rh)
|
||||||
|
return clip_box(
|
||||||
|
np.array([cx - 0.5 * side, cy - 0.5 * side, cx + 0.5 * side, cy + 0.5 * side], dtype=np.float32),
|
||||||
|
frame_w,
|
||||||
|
frame_h,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _pick_roi(self, rois, pred_ref_box):
|
||||||
|
if not rois:
|
||||||
|
return None
|
||||||
|
if pred_ref_box is None:
|
||||||
|
return rois[0]
|
||||||
|
|
||||||
|
pref = np.array(pred_ref_box, dtype=np.float32)
|
||||||
|
pc = box_center(pref)
|
||||||
|
|
||||||
|
best = None
|
||||||
|
best_key = (1e9, 1e9)
|
||||||
|
for r in rois:
|
||||||
|
c = box_center(r)
|
||||||
|
dist = float(np.linalg.norm(c - pc))
|
||||||
|
a = float(box_area(r))
|
||||||
|
key = (dist, -a)
|
||||||
|
if key < best_key:
|
||||||
|
best_key = key
|
||||||
|
best = r
|
||||||
|
return best
|
||||||
@ -0,0 +1,339 @@
|
|||||||
|
# bytetrack_min.py
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
def tlbr_to_xyah(tlbr):
|
||||||
|
x1, y1, x2, y2 = tlbr
|
||||||
|
w = max(1.0, x2 - x1)
|
||||||
|
h = max(1.0, y2 - y1)
|
||||||
|
cx = x1 + w * 0.5
|
||||||
|
cy = y1 + h * 0.5
|
||||||
|
a = w / h
|
||||||
|
return np.array([cx, cy, a, h], dtype=np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def xyah_to_tlbr(xyah):
|
||||||
|
cx, cy, a, h = [float(v) for v in xyah]
|
||||||
|
h = max(2.0, h)
|
||||||
|
w = max(2.0, a * h)
|
||||||
|
x1 = cx - w * 0.5
|
||||||
|
y1 = cy - h * 0.5
|
||||||
|
x2 = cx + w * 0.5
|
||||||
|
y2 = cy + h * 0.5
|
||||||
|
return np.array([x1, y1, x2, y2], dtype=np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def iou_tlbr(a, b):
|
||||||
|
ax1, ay1, ax2, ay2 = a
|
||||||
|
bx1, by1, bx2, by2 = b
|
||||||
|
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
||||||
|
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
||||||
|
iw = max(0.0, ix2 - ix1)
|
||||||
|
ih = max(0.0, iy2 - iy1)
|
||||||
|
inter = iw * ih
|
||||||
|
if inter <= 0:
|
||||||
|
return 0.0
|
||||||
|
area_a = max(1.0, (ax2 - ax1)) * max(1.0, (ay2 - ay1))
|
||||||
|
area_b = max(1.0, (bx2 - bx1)) * max(1.0, (by2 - by1))
|
||||||
|
union = area_a + area_b - inter
|
||||||
|
return float(inter / union) if union > 0 else 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def iou_matrix(tracks_tlbr, dets_tlbr):
|
||||||
|
if len(tracks_tlbr) == 0 or len(dets_tlbr) == 0:
|
||||||
|
return np.zeros((len(tracks_tlbr), len(dets_tlbr)), dtype=np.float32)
|
||||||
|
matrix = np.zeros((len(tracks_tlbr), len(dets_tlbr)), dtype=np.float32)
|
||||||
|
for i, track in enumerate(tracks_tlbr):
|
||||||
|
for j, det in enumerate(dets_tlbr):
|
||||||
|
matrix[i, j] = iou_tlbr(track, det)
|
||||||
|
return matrix
|
||||||
|
|
||||||
|
|
||||||
|
def hungarian(cost):
|
||||||
|
cost = cost.copy()
|
||||||
|
n, m = cost.shape
|
||||||
|
size = max(n, m)
|
||||||
|
|
||||||
|
pad = np.zeros((size, size), dtype=np.float32)
|
||||||
|
pad[:n, :m] = cost
|
||||||
|
big = float(cost.max() + 1.0) if cost.size else 1.0
|
||||||
|
if n < size:
|
||||||
|
pad[n:, :] = big
|
||||||
|
if m < size:
|
||||||
|
pad[:, m:] = big
|
||||||
|
cost = pad
|
||||||
|
|
||||||
|
size = cost.shape[0]
|
||||||
|
u = np.zeros(size, dtype=np.float32)
|
||||||
|
v = np.zeros(size, dtype=np.float32)
|
||||||
|
p = np.zeros(size, dtype=np.int32)
|
||||||
|
way = np.zeros(size, dtype=np.int32)
|
||||||
|
|
||||||
|
for i in range(1, size):
|
||||||
|
p[0] = i
|
||||||
|
j0 = 0
|
||||||
|
minv = np.full(size, np.inf, dtype=np.float32)
|
||||||
|
used = np.zeros(size, dtype=bool)
|
||||||
|
way.fill(0)
|
||||||
|
|
||||||
|
while True:
|
||||||
|
used[j0] = True
|
||||||
|
i0 = p[j0]
|
||||||
|
delta = np.inf
|
||||||
|
j1 = 0
|
||||||
|
for j in range(1, size):
|
||||||
|
if not used[j]:
|
||||||
|
cur = cost[i0, j] - u[i0] - v[j]
|
||||||
|
if cur < minv[j]:
|
||||||
|
minv[j] = cur
|
||||||
|
way[j] = j0
|
||||||
|
if minv[j] < delta:
|
||||||
|
delta = minv[j]
|
||||||
|
j1 = j
|
||||||
|
for j in range(size):
|
||||||
|
if used[j]:
|
||||||
|
u[p[j]] += delta
|
||||||
|
v[j] -= delta
|
||||||
|
else:
|
||||||
|
minv[j] -= delta
|
||||||
|
j0 = j1
|
||||||
|
if p[j0] == 0:
|
||||||
|
break
|
||||||
|
|
||||||
|
while True:
|
||||||
|
j1 = way[j0]
|
||||||
|
p[j0] = p[j1]
|
||||||
|
j0 = j1
|
||||||
|
if j0 == 0:
|
||||||
|
break
|
||||||
|
|
||||||
|
assignment = -np.ones(size, dtype=np.int32)
|
||||||
|
for j in range(1, size):
|
||||||
|
if p[j] != 0:
|
||||||
|
assignment[p[j]] = j
|
||||||
|
|
||||||
|
row_to_col = assignment[:n]
|
||||||
|
matches = []
|
||||||
|
for row, col in enumerate(row_to_col):
|
||||||
|
if 0 <= col < m:
|
||||||
|
matches.append((row, int(col)))
|
||||||
|
return matches
|
||||||
|
|
||||||
|
|
||||||
|
class KalmanXYAH:
|
||||||
|
def __init__(self):
|
||||||
|
self.ndim = 4
|
||||||
|
self.dim_x = 8
|
||||||
|
self._motion_mat = np.eye(self.dim_x, dtype=np.float32)
|
||||||
|
for i in range(self.ndim):
|
||||||
|
self._motion_mat[i, self.ndim + i] = 1.0
|
||||||
|
self._update_mat = np.eye(self.ndim, self.dim_x, dtype=np.float32)
|
||||||
|
self.std_pos = np.array([1.5, 1.5, 1e-2, 2.0], dtype=np.float32)
|
||||||
|
self.std_vel = np.array([5.0, 5.0, 1e-2, 5.0], dtype=np.float32)
|
||||||
|
|
||||||
|
def initiate(self, measurement_xyah):
|
||||||
|
mean = np.zeros((self.dim_x,), dtype=np.float32)
|
||||||
|
mean[:4] = measurement_xyah
|
||||||
|
cov = np.eye(self.dim_x, dtype=np.float32)
|
||||||
|
cov[:4, :4] *= 50.0
|
||||||
|
cov[4:, 4:] *= 200.0
|
||||||
|
return mean, cov
|
||||||
|
|
||||||
|
def predict(self, mean, cov, dt=1.0):
|
||||||
|
motion_mat = self._motion_mat.copy()
|
||||||
|
for i in range(self.ndim):
|
||||||
|
motion_mat[i, self.ndim + i] = float(dt)
|
||||||
|
q = np.diag(np.concatenate([self.std_pos ** 2, self.std_vel ** 2]).astype(np.float32))
|
||||||
|
mean = motion_mat @ mean
|
||||||
|
cov = motion_mat @ cov @ motion_mat.T + q
|
||||||
|
return mean, cov
|
||||||
|
|
||||||
|
def update(self, mean, cov, measurement_xyah):
|
||||||
|
r = np.diag((self.std_pos ** 2).astype(np.float32))
|
||||||
|
s = self._update_mat @ cov @ self._update_mat.T + r
|
||||||
|
k = cov @ self._update_mat.T @ np.linalg.inv(s)
|
||||||
|
y = measurement_xyah - (self._update_mat @ mean)
|
||||||
|
mean = mean + k @ y
|
||||||
|
cov = (np.eye(self.dim_x, dtype=np.float32) - k @ self._update_mat) @ cov
|
||||||
|
return mean, cov
|
||||||
|
|
||||||
|
|
||||||
|
class TrackState:
|
||||||
|
Tracked = 1
|
||||||
|
Lost = 2
|
||||||
|
Removed = 3
|
||||||
|
|
||||||
|
|
||||||
|
class STrack:
|
||||||
|
_next_id = 1
|
||||||
|
|
||||||
|
def __init__(self, tlbr, score, kf: KalmanXYAH):
|
||||||
|
self.kf = kf
|
||||||
|
self.mean = None
|
||||||
|
self.cov = None
|
||||||
|
self.tlbr = np.array(tlbr, dtype=np.float32)
|
||||||
|
self.score = float(score)
|
||||||
|
self.track_id = STrack._next_id
|
||||||
|
STrack._next_id += 1
|
||||||
|
self.state = TrackState.Tracked
|
||||||
|
self.is_activated = False
|
||||||
|
self.frame_id = 0
|
||||||
|
self.start_frame = 0
|
||||||
|
self.time_since_update = 0
|
||||||
|
self.hits = 0
|
||||||
|
|
||||||
|
def activate(self, frame_id):
|
||||||
|
self.frame_id = frame_id
|
||||||
|
self.start_frame = frame_id
|
||||||
|
self.time_since_update = 0
|
||||||
|
xyah = tlbr_to_xyah(self.tlbr)
|
||||||
|
self.mean, self.cov = self.kf.initiate(xyah)
|
||||||
|
self.is_activated = True
|
||||||
|
self.hits = 1
|
||||||
|
self.state = TrackState.Tracked
|
||||||
|
|
||||||
|
def predict(self, dt=1.0):
|
||||||
|
if self.mean is None:
|
||||||
|
return
|
||||||
|
self.mean, self.cov = self.kf.predict(self.mean, self.cov, dt=dt)
|
||||||
|
self.tlbr = xyah_to_tlbr(self.mean[:4])
|
||||||
|
self.time_since_update += 1
|
||||||
|
|
||||||
|
def update(self, tlbr, score, frame_id):
|
||||||
|
self.frame_id = frame_id
|
||||||
|
self.time_since_update = 0
|
||||||
|
self.score = float(score)
|
||||||
|
self.hits += 1
|
||||||
|
xyah = tlbr_to_xyah(tlbr)
|
||||||
|
self.mean, self.cov = self.kf.update(self.mean, self.cov, xyah)
|
||||||
|
self.tlbr = xyah_to_tlbr(self.mean[:4])
|
||||||
|
self.state = TrackState.Tracked
|
||||||
|
self.is_activated = True
|
||||||
|
|
||||||
|
def mark_lost(self):
|
||||||
|
self.state = TrackState.Lost
|
||||||
|
|
||||||
|
def mark_removed(self):
|
||||||
|
self.state = TrackState.Removed
|
||||||
|
|
||||||
|
|
||||||
|
class BYTETracker:
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
track_high_thresh=0.35,
|
||||||
|
track_low_thresh=0.12,
|
||||||
|
new_track_thresh=0.35,
|
||||||
|
match_thresh=0.70,
|
||||||
|
track_buffer=50,
|
||||||
|
min_hits=2,
|
||||||
|
):
|
||||||
|
self.high = float(track_high_thresh)
|
||||||
|
self.low = float(track_low_thresh)
|
||||||
|
self.new = float(new_track_thresh)
|
||||||
|
self.match_thresh = float(match_thresh)
|
||||||
|
self.track_buffer = int(track_buffer)
|
||||||
|
self.min_hits = int(min_hits)
|
||||||
|
self.kf = KalmanXYAH()
|
||||||
|
self.frame_id = 0
|
||||||
|
self.tracked = []
|
||||||
|
self.lost = []
|
||||||
|
self.removed = []
|
||||||
|
|
||||||
|
def update(self, dets_tlbr_score, dt=1.0):
|
||||||
|
self.frame_id += 1
|
||||||
|
frame_id = self.frame_id
|
||||||
|
|
||||||
|
if dets_tlbr_score is None:
|
||||||
|
for track in self.tracked:
|
||||||
|
track.predict(dt=dt)
|
||||||
|
kept_lost = []
|
||||||
|
for track in self.lost:
|
||||||
|
if (frame_id - track.frame_id) <= self.track_buffer:
|
||||||
|
kept_lost.append(track)
|
||||||
|
else:
|
||||||
|
track.mark_removed()
|
||||||
|
self.removed.append(track)
|
||||||
|
self.lost = kept_lost
|
||||||
|
out = []
|
||||||
|
for track in self.tracked:
|
||||||
|
if track.hits >= self.min_hits:
|
||||||
|
out.append(track)
|
||||||
|
return out
|
||||||
|
|
||||||
|
dets = dets_tlbr_score.astype(np.float32, copy=False)
|
||||||
|
scores = np.zeros((0,), dtype=np.float32) if len(dets) == 0 else dets[:, 4]
|
||||||
|
high_mask = scores >= self.high
|
||||||
|
low_mask = (scores >= self.low) & (scores < self.high)
|
||||||
|
dets_high = dets[high_mask]
|
||||||
|
dets_low = dets[low_mask]
|
||||||
|
|
||||||
|
for track in self.tracked:
|
||||||
|
track.predict(dt=dt)
|
||||||
|
|
||||||
|
matches_a, u_trk_a, u_det_a = self._associate(self.tracked, dets_high, iou_thresh=self.match_thresh)
|
||||||
|
for ti, di in matches_a:
|
||||||
|
track = self.tracked[ti]
|
||||||
|
det = dets_high[di]
|
||||||
|
track.update(det[:4], det[4], frame_id)
|
||||||
|
|
||||||
|
remaining_tracks = [self.tracked[i] for i in u_trk_a]
|
||||||
|
matches_b, u_trk_b, _ = self._associate(
|
||||||
|
remaining_tracks, dets_low, iou_thresh=max(0.10, self.match_thresh - 0.15)
|
||||||
|
)
|
||||||
|
for ti, di in matches_b:
|
||||||
|
track = remaining_tracks[ti]
|
||||||
|
det = dets_low[di]
|
||||||
|
track.update(det[:4], det[4], frame_id)
|
||||||
|
|
||||||
|
unmatched_after_b = [remaining_tracks[i] for i in u_trk_b]
|
||||||
|
for track in unmatched_after_b:
|
||||||
|
track.mark_lost()
|
||||||
|
|
||||||
|
new_tracked = [track for track in self.tracked if track.state == TrackState.Tracked]
|
||||||
|
newly_lost = [track for track in self.tracked if track.state == TrackState.Lost]
|
||||||
|
self.tracked = new_tracked
|
||||||
|
self.lost.extend(newly_lost)
|
||||||
|
|
||||||
|
for di in u_det_a:
|
||||||
|
det = dets_high[di]
|
||||||
|
if float(det[4]) >= self.new:
|
||||||
|
new_track = STrack(det[:4], det[4], self.kf)
|
||||||
|
new_track.activate(frame_id)
|
||||||
|
self.tracked.append(new_track)
|
||||||
|
|
||||||
|
kept_lost = []
|
||||||
|
for track in self.lost:
|
||||||
|
if (frame_id - track.frame_id) <= self.track_buffer:
|
||||||
|
kept_lost.append(track)
|
||||||
|
else:
|
||||||
|
track.mark_removed()
|
||||||
|
self.removed.append(track)
|
||||||
|
self.lost = kept_lost
|
||||||
|
|
||||||
|
out = []
|
||||||
|
for track in self.tracked:
|
||||||
|
if track.hits >= self.min_hits:
|
||||||
|
out.append(track)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def _associate(self, tracks, dets, iou_thresh):
|
||||||
|
if len(tracks) == 0 or len(dets) == 0:
|
||||||
|
return [], list(range(len(tracks))), list(range(len(dets)))
|
||||||
|
|
||||||
|
trk_boxes = [track.tlbr for track in tracks]
|
||||||
|
det_boxes = [det[:4] for det in dets]
|
||||||
|
ious = iou_matrix(trk_boxes, det_boxes)
|
||||||
|
cost = 1.0 - ious
|
||||||
|
matches = hungarian(cost)
|
||||||
|
|
||||||
|
matched = []
|
||||||
|
u_trk = set(range(len(tracks)))
|
||||||
|
u_det = set(range(len(dets)))
|
||||||
|
threshold = float(iou_thresh)
|
||||||
|
for ti, di in matches:
|
||||||
|
if ious[ti, di] >= threshold:
|
||||||
|
matched.append((ti, di))
|
||||||
|
u_trk.discard(ti)
|
||||||
|
u_det.discard(di)
|
||||||
|
|
||||||
|
return matched, sorted(list(u_trk)), sorted(list(u_det))
|
||||||
@ -0,0 +1,81 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import ultralytics.nn.modules as um
|
||||||
|
import ultralytics.nn.tasks as ut
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
class ChannelAttentionDyn(nn.Module):
|
||||||
|
def __init__(self, reduction=16):
|
||||||
|
super().__init__()
|
||||||
|
self.reduction = reduction
|
||||||
|
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.max_pool = nn.AdaptiveMaxPool2d(1)
|
||||||
|
self.mlp = None
|
||||||
|
self.sigmoid = nn.Sigmoid()
|
||||||
|
|
||||||
|
def _build(self, c, device, dtype):
|
||||||
|
hidden = max(c // self.reduction, 1)
|
||||||
|
self.mlp = nn.Sequential(
|
||||||
|
nn.Conv2d(c, hidden, 1, bias=False),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
nn.Conv2d(hidden, c, 1, bias=False),
|
||||||
|
).to(device=device, dtype=dtype)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if self.mlp is None:
|
||||||
|
self._build(x.shape[1], x.device, x.dtype)
|
||||||
|
else:
|
||||||
|
p = next(self.mlp.parameters())
|
||||||
|
if p.device != x.device or p.dtype != x.dtype:
|
||||||
|
self.mlp = self.mlp.to(device=x.device, dtype=x.dtype)
|
||||||
|
|
||||||
|
a = self.mlp(self.avg_pool(x))
|
||||||
|
m = self.mlp(self.max_pool(x))
|
||||||
|
return x * self.sigmoid(a + m)
|
||||||
|
|
||||||
|
|
||||||
|
class SpatialAttention(nn.Module):
|
||||||
|
def __init__(self, kernel_size=7):
|
||||||
|
super().__init__()
|
||||||
|
padding = kernel_size // 2
|
||||||
|
self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
|
||||||
|
self.sigmoid = nn.Sigmoid()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if self.conv.weight.device != x.device or self.conv.weight.dtype != x.dtype:
|
||||||
|
self.conv = self.conv.to(device=x.device, dtype=x.dtype)
|
||||||
|
|
||||||
|
avg = torch.mean(x, dim=1, keepdim=True)
|
||||||
|
mx, _ = torch.max(x, dim=1, keepdim=True)
|
||||||
|
w = self.sigmoid(self.conv(torch.cat([avg, mx], dim=1)))
|
||||||
|
return x * w
|
||||||
|
|
||||||
|
|
||||||
|
class CBAM(nn.Module):
|
||||||
|
def __init__(self, reduction=16, sa_kernel=7):
|
||||||
|
super().__init__()
|
||||||
|
self.ca = ChannelAttentionDyn(reduction=reduction)
|
||||||
|
self.sa = SpatialAttention(kernel_size=sa_kernel)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.sa(self.ca(x))
|
||||||
|
|
||||||
|
|
||||||
|
# Register so torch/ultralytics can deserialize CBAM checkpoints.
|
||||||
|
um.CBAM = CBAM
|
||||||
|
ut.CBAM = CBAM
|
||||||
|
um.ChannelAttentionDyn = ChannelAttentionDyn
|
||||||
|
ut.ChannelAttentionDyn = ChannelAttentionDyn
|
||||||
|
um.SpatialAttention = SpatialAttention
|
||||||
|
ut.SpatialAttention = SpatialAttention
|
||||||
|
globals()["ChannelAttentionDyn"] = ChannelAttentionDyn
|
||||||
|
globals()["SpatialAttention"] = SpatialAttention
|
||||||
|
globals()["CBAM"] = CBAM
|
||||||
|
|
||||||
|
# Compatibility for checkpoints serialized with __main__.CBAM
|
||||||
|
main_mod = sys.modules.get("__main__")
|
||||||
|
if main_mod is not None:
|
||||||
|
setattr(main_mod, "ChannelAttentionDyn", ChannelAttentionDyn)
|
||||||
|
setattr(main_mod, "SpatialAttention", SpatialAttention)
|
||||||
|
setattr(main_mod, "CBAM", CBAM)
|
||||||
@ -0,0 +1,135 @@
|
|||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import shutil
|
||||||
|
import zipfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Package extracted frame folders into per-video ZIP archives for CVAT."
|
||||||
|
)
|
||||||
|
parser.add_argument("--frames-root", required=True, help="Root with extracted frames, e.g. D:\\MAI_CVAT_frames_s15\\frames_all")
|
||||||
|
parser.add_argument("--output-root", required=True, help="Where ZIP archives and manifests will be written")
|
||||||
|
parser.add_argument("--group-filter", action="append", default=[], help="Optional substring filter for group names")
|
||||||
|
parser.add_argument("--max-files-per-zip", type=int, default=0, help="Optional limit; skip items above this count when > 0")
|
||||||
|
parser.add_argument("--overwrite", action="store_true")
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
args = parse_args()
|
||||||
|
frames_root = Path(args.frames_root).expanduser().resolve()
|
||||||
|
output_root = Path(args.output_root).expanduser().resolve()
|
||||||
|
|
||||||
|
if not frames_root.exists():
|
||||||
|
raise SystemExit(f"Frames root not found: {frames_root}")
|
||||||
|
|
||||||
|
filters = [s.strip().lower() for s in args.group_filter if s.strip()]
|
||||||
|
zip_root = output_root / "zips"
|
||||||
|
zip_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
rows = []
|
||||||
|
video_dirs = list(find_video_dirs(frames_root))
|
||||||
|
for video_dir in video_dirs:
|
||||||
|
rel_parts = video_dir.relative_to(frames_root).parts
|
||||||
|
if len(rel_parts) < 5:
|
||||||
|
continue
|
||||||
|
|
||||||
|
split, role, group_name, video_name, interval_name = rel_parts[:5]
|
||||||
|
if filters and not any(token in group_name.lower() for token in filters):
|
||||||
|
continue
|
||||||
|
|
||||||
|
files = sorted(
|
||||||
|
[
|
||||||
|
p for p in video_dir.iterdir()
|
||||||
|
if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png"}
|
||||||
|
]
|
||||||
|
)
|
||||||
|
file_count = len(files)
|
||||||
|
if file_count == 0:
|
||||||
|
continue
|
||||||
|
if args.max_files_per_zip > 0 and file_count > args.max_files_per_zip:
|
||||||
|
continue
|
||||||
|
|
||||||
|
safe_group = sanitize(group_name)
|
||||||
|
safe_video = sanitize(video_name)
|
||||||
|
safe_interval = sanitize(interval_name)
|
||||||
|
zip_name = f"{safe_group}__{safe_video}__{safe_interval}.zip"
|
||||||
|
zip_path = zip_root / zip_name
|
||||||
|
|
||||||
|
if zip_path.exists():
|
||||||
|
if args.overwrite:
|
||||||
|
zip_path.unlink()
|
||||||
|
else:
|
||||||
|
rows.append(build_row(split, role, group_name, video_name, interval_name, file_count, zip_path, video_dir))
|
||||||
|
continue
|
||||||
|
|
||||||
|
with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_STORED) as zf:
|
||||||
|
for file_path in files:
|
||||||
|
zf.write(file_path, arcname=file_path.name)
|
||||||
|
|
||||||
|
rows.append(build_row(split, role, group_name, video_name, interval_name, file_count, zip_path, video_dir))
|
||||||
|
|
||||||
|
manifest_path = output_root / "cvat_task_packages.csv"
|
||||||
|
write_csv(
|
||||||
|
manifest_path,
|
||||||
|
rows,
|
||||||
|
[
|
||||||
|
"split",
|
||||||
|
"role",
|
||||||
|
"group_name",
|
||||||
|
"video_name",
|
||||||
|
"interval_name",
|
||||||
|
"frame_count",
|
||||||
|
"zip_path",
|
||||||
|
"source_dir",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
print(f"Wrote {len(rows)} package rows to {manifest_path}")
|
||||||
|
print(f"ZIP archives are in {zip_root}")
|
||||||
|
|
||||||
|
|
||||||
|
def find_video_dirs(frames_root: Path):
|
||||||
|
for split_dir in sorted(p for p in frames_root.iterdir() if p.is_dir()):
|
||||||
|
for role_dir in sorted(p for p in split_dir.iterdir() if p.is_dir()):
|
||||||
|
for group_dir in sorted(p for p in role_dir.iterdir() if p.is_dir()):
|
||||||
|
for video_dir in sorted(p for p in group_dir.iterdir() if p.is_dir()):
|
||||||
|
for interval_dir in sorted(p for p in video_dir.iterdir() if p.is_dir()):
|
||||||
|
yield interval_dir
|
||||||
|
|
||||||
|
|
||||||
|
def sanitize(value: str):
|
||||||
|
out = []
|
||||||
|
for ch in value:
|
||||||
|
if ch in '<>:"/\\|?*':
|
||||||
|
out.append("_")
|
||||||
|
else:
|
||||||
|
out.append(ch)
|
||||||
|
return "".join(out).strip().rstrip(".")
|
||||||
|
|
||||||
|
|
||||||
|
def build_row(split, role, group_name, video_name, interval_name, file_count, zip_path: Path, source_dir: Path):
|
||||||
|
return {
|
||||||
|
"split": split,
|
||||||
|
"role": role,
|
||||||
|
"group_name": group_name,
|
||||||
|
"video_name": video_name,
|
||||||
|
"interval_name": interval_name,
|
||||||
|
"frame_count": file_count,
|
||||||
|
"zip_path": str(zip_path),
|
||||||
|
"source_dir": str(source_dir),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def write_csv(path: Path, rows, fieldnames):
|
||||||
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with path.open("w", newline="", encoding="utf-8-sig") as f:
|
||||||
|
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||||
|
writer.writeheader()
|
||||||
|
for row in rows:
|
||||||
|
writer.writerow(row)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,735 @@
|
|||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import shutil
|
||||||
|
import subprocess
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".mts", ".m2ts"}
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class VideoMeta:
|
||||||
|
abs_path: Path
|
||||||
|
rel_path: Path
|
||||||
|
group_name: str
|
||||||
|
file_name: str
|
||||||
|
size_bytes: int
|
||||||
|
width: int
|
||||||
|
height: int
|
||||||
|
fps: float
|
||||||
|
duration_sec: float
|
||||||
|
nb_frames: int
|
||||||
|
osd_present: bool
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Prepare grouped frame datasets from many videos for CVAT and YOLO."
|
||||||
|
)
|
||||||
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||||
|
|
||||||
|
catalog = subparsers.add_parser("catalog", help="Scan videos and write catalog CSV files.")
|
||||||
|
catalog.add_argument("--input-root", required=True, help="Root folder with grouped videos.")
|
||||||
|
catalog.add_argument("--output-root", required=True, help="Workspace folder for CSV outputs.")
|
||||||
|
|
||||||
|
preview = subparsers.add_parser(
|
||||||
|
"preview",
|
||||||
|
help="Extract low-rate preview frames for interval selection.",
|
||||||
|
)
|
||||||
|
preview.add_argument("--input-root", required=True)
|
||||||
|
preview.add_argument("--output-root", required=True)
|
||||||
|
preview.add_argument("--every-sec", type=float, default=2.0)
|
||||||
|
preview.add_argument("--image-ext", choices=("jpg", "png"), default="jpg")
|
||||||
|
preview.add_argument("--jpg-quality", type=int, default=95)
|
||||||
|
preview.add_argument("--overwrite", action="store_true")
|
||||||
|
|
||||||
|
extract = subparsers.add_parser(
|
||||||
|
"extract",
|
||||||
|
help="Extract frames for selected intervals from intervals CSV.",
|
||||||
|
)
|
||||||
|
extract.add_argument("--input-root", required=True)
|
||||||
|
extract.add_argument("--intervals-csv", required=True)
|
||||||
|
extract.add_argument("--output-root", required=True)
|
||||||
|
extract.add_argument("--overwrite", action="store_true")
|
||||||
|
|
||||||
|
extract_all = subparsers.add_parser(
|
||||||
|
"extract-all",
|
||||||
|
help="Extract frames from every video while preserving the input group structure.",
|
||||||
|
)
|
||||||
|
extract_all.add_argument("--input-root", required=True)
|
||||||
|
extract_all.add_argument("--output-root", required=True)
|
||||||
|
extract_all.add_argument("--frame-step", type=int, default=2)
|
||||||
|
extract_all.add_argument("--image-ext", choices=("jpg", "png"), default="png")
|
||||||
|
extract_all.add_argument("--jpg-quality", type=int, default=95)
|
||||||
|
extract_all.add_argument("--split", default="train")
|
||||||
|
extract_all.add_argument("--role", default="bulk_raw")
|
||||||
|
extract_all.add_argument(
|
||||||
|
"--group-filter",
|
||||||
|
action="append",
|
||||||
|
default=[],
|
||||||
|
help="Optional substring filter for group names. Can be repeated.",
|
||||||
|
)
|
||||||
|
extract_all.add_argument("--overwrite", action="store_true")
|
||||||
|
|
||||||
|
extract_flat_all = subparsers.add_parser(
|
||||||
|
"extract-flat-all",
|
||||||
|
help="Extract frames from every video into one flat output folder.",
|
||||||
|
)
|
||||||
|
extract_flat_all.add_argument("--input-root", required=True)
|
||||||
|
extract_flat_all.add_argument(
|
||||||
|
"--output-root",
|
||||||
|
required=True,
|
||||||
|
help="Flat folder where all images will be written.",
|
||||||
|
)
|
||||||
|
extract_flat_all.add_argument("--frame-step", type=int, default=1)
|
||||||
|
extract_flat_all.add_argument("--image-ext", choices=("jpg", "png"), default="jpg")
|
||||||
|
extract_flat_all.add_argument("--jpg-quality", type=int, default=95)
|
||||||
|
extract_flat_all.add_argument(
|
||||||
|
"--group-filter",
|
||||||
|
action="append",
|
||||||
|
default=[],
|
||||||
|
help="Optional substring filter for group names. Can be repeated.",
|
||||||
|
)
|
||||||
|
extract_flat_all.add_argument("--overwrite", action="store_true")
|
||||||
|
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
args = parse_args()
|
||||||
|
input_root = Path(args.input_root).expanduser().resolve()
|
||||||
|
output_root = Path(args.output_root).expanduser().resolve()
|
||||||
|
|
||||||
|
if not input_root.exists():
|
||||||
|
raise SystemExit(f"Input root not found: {input_root}")
|
||||||
|
|
||||||
|
ensure_tool("ffprobe")
|
||||||
|
ensure_tool("ffmpeg")
|
||||||
|
|
||||||
|
if args.command == "catalog":
|
||||||
|
videos = scan_videos(input_root)
|
||||||
|
write_catalog_outputs(videos, output_root)
|
||||||
|
print(f"Wrote catalog for {len(videos)} videos into {output_root}")
|
||||||
|
return
|
||||||
|
|
||||||
|
if args.command == "preview":
|
||||||
|
videos = scan_videos(input_root)
|
||||||
|
preview_root = output_root / "previews"
|
||||||
|
preview_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
manifest_path = output_root / "preview_manifest.csv"
|
||||||
|
rows = []
|
||||||
|
for meta in videos:
|
||||||
|
out_dir = preview_root / sanitize_path_component(meta.group_name) / sanitize_path_component(meta.rel_path.stem)
|
||||||
|
if args.overwrite and out_dir.exists():
|
||||||
|
shutil.rmtree(out_dir)
|
||||||
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
pattern = out_dir / f"preview_%06d.{args.image_ext}"
|
||||||
|
cmd = [
|
||||||
|
"ffmpeg",
|
||||||
|
"-hide_banner",
|
||||||
|
"-loglevel",
|
||||||
|
"warning",
|
||||||
|
"-y",
|
||||||
|
"-i",
|
||||||
|
str(meta.abs_path),
|
||||||
|
"-vf",
|
||||||
|
f"fps=1/{float(args.every_sec)}",
|
||||||
|
]
|
||||||
|
if args.image_ext == "jpg":
|
||||||
|
cmd.extend(["-q:v", str(jpeg_qscale(args.jpg_quality))])
|
||||||
|
cmd.append(str(pattern))
|
||||||
|
run_cmd(cmd)
|
||||||
|
|
||||||
|
files = sorted(out_dir.glob(f"*.{args.image_ext}"))
|
||||||
|
for idx, file_path in enumerate(files):
|
||||||
|
rows.append(
|
||||||
|
{
|
||||||
|
"group_name": meta.group_name,
|
||||||
|
"video_relpath": str(meta.rel_path).replace("\\", "/"),
|
||||||
|
"preview_file": str(file_path.relative_to(output_root)).replace("\\", "/"),
|
||||||
|
"preview_index": idx + 1,
|
||||||
|
"approx_time_sec": round(idx * float(args.every_sec), 3),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
write_csv(
|
||||||
|
manifest_path,
|
||||||
|
rows,
|
||||||
|
["group_name", "video_relpath", "preview_file", "preview_index", "approx_time_sec"],
|
||||||
|
)
|
||||||
|
print(f"Wrote previews into {preview_root}")
|
||||||
|
print(f"Wrote preview manifest: {manifest_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
if args.command == "extract":
|
||||||
|
video_map = {str(v.rel_path).replace("\\", "/"): v for v in scan_videos(input_root)}
|
||||||
|
intervals_csv = Path(args.intervals_csv).expanduser().resolve()
|
||||||
|
if not intervals_csv.exists():
|
||||||
|
raise SystemExit(f"Intervals CSV not found: {intervals_csv}")
|
||||||
|
|
||||||
|
rows = load_intervals(intervals_csv)
|
||||||
|
manifest_rows = []
|
||||||
|
extracted_root = output_root / "frames"
|
||||||
|
extracted_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
for row in rows:
|
||||||
|
if not is_enabled(row.get("enabled", "")):
|
||||||
|
continue
|
||||||
|
|
||||||
|
rel_key = normalize_relpath(row["video_relpath"])
|
||||||
|
if rel_key not in video_map:
|
||||||
|
raise SystemExit(f"Video not found in catalog: {rel_key}")
|
||||||
|
|
||||||
|
meta = video_map[rel_key]
|
||||||
|
split = row.get("split", "train").strip() or "train"
|
||||||
|
role = row.get("role", "positive").strip() or "positive"
|
||||||
|
frame_step = max(1, int(float(row.get("frame_step", "1") or "1")))
|
||||||
|
start_sec = clamp_float(row.get("start_sec", "0"), 0.0, meta.duration_sec)
|
||||||
|
end_sec_raw = row.get("end_sec", "")
|
||||||
|
end_sec = meta.duration_sec if end_sec_raw == "" else clamp_float(end_sec_raw, 0.0, meta.duration_sec)
|
||||||
|
if end_sec <= start_sec:
|
||||||
|
raise SystemExit(f"Invalid interval for {rel_key}: end_sec <= start_sec")
|
||||||
|
|
||||||
|
image_ext = (row.get("image_ext", "png").strip().lower() or "png")
|
||||||
|
if image_ext not in {"png", "jpg"}:
|
||||||
|
raise SystemExit(f"Unsupported image_ext '{image_ext}' in intervals CSV")
|
||||||
|
jpg_quality = max(70, min(100, int(float(row.get("jpg_quality", "95") or "95"))))
|
||||||
|
|
||||||
|
extract_rows = extract_range_to_dir(
|
||||||
|
meta=meta,
|
||||||
|
output_root=output_root,
|
||||||
|
extracted_root=extracted_root,
|
||||||
|
split=split,
|
||||||
|
role=role,
|
||||||
|
start_sec=start_sec,
|
||||||
|
end_sec=end_sec,
|
||||||
|
frame_step=frame_step,
|
||||||
|
image_ext=image_ext,
|
||||||
|
jpg_quality=jpg_quality,
|
||||||
|
overwrite=args.overwrite,
|
||||||
|
interval_name=format_interval_name(start_sec, end_sec, frame_step),
|
||||||
|
)
|
||||||
|
manifest_rows.extend(extract_rows)
|
||||||
|
|
||||||
|
manifest_path = output_root / "extracted_frames_manifest.csv"
|
||||||
|
write_csv(
|
||||||
|
manifest_path,
|
||||||
|
manifest_rows,
|
||||||
|
["split", "role", "group_name", "video_relpath", "output_file", "frame_idx", "time_sec", "frame_step", "image_ext"],
|
||||||
|
)
|
||||||
|
print(f"Wrote extracted frames into {extracted_root}")
|
||||||
|
print(f"Wrote extraction manifest: {manifest_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
if args.command == "extract-all":
|
||||||
|
videos = scan_videos(input_root)
|
||||||
|
if args.group_filter:
|
||||||
|
filters = [v.strip().lower() for v in args.group_filter if v.strip()]
|
||||||
|
videos = [meta for meta in videos if any(token in meta.group_name.lower() for token in filters)]
|
||||||
|
|
||||||
|
frame_step = max(1, int(args.frame_step))
|
||||||
|
image_ext = args.image_ext
|
||||||
|
jpg_quality = max(70, min(100, int(args.jpg_quality)))
|
||||||
|
split = args.split.strip() or "train"
|
||||||
|
role = args.role.strip() or "bulk_raw"
|
||||||
|
|
||||||
|
extracted_root = output_root / "frames_all"
|
||||||
|
extracted_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
manifest_rows = []
|
||||||
|
|
||||||
|
for meta in videos:
|
||||||
|
extract_rows = extract_range_to_dir(
|
||||||
|
meta=meta,
|
||||||
|
output_root=output_root,
|
||||||
|
extracted_root=extracted_root,
|
||||||
|
split=split,
|
||||||
|
role=role,
|
||||||
|
start_sec=0.0,
|
||||||
|
end_sec=meta.duration_sec,
|
||||||
|
frame_step=frame_step,
|
||||||
|
image_ext=image_ext,
|
||||||
|
jpg_quality=jpg_quality,
|
||||||
|
overwrite=args.overwrite,
|
||||||
|
interval_name=f"all_step{frame_step}",
|
||||||
|
)
|
||||||
|
manifest_rows.extend(extract_rows)
|
||||||
|
|
||||||
|
manifest_path = output_root / "extract_all_manifest.csv"
|
||||||
|
write_csv(
|
||||||
|
manifest_path,
|
||||||
|
manifest_rows,
|
||||||
|
["split", "role", "group_name", "video_relpath", "output_file", "frame_idx", "time_sec", "frame_step", "image_ext"],
|
||||||
|
)
|
||||||
|
print(f"Wrote bulk extracted frames into {extracted_root}")
|
||||||
|
print(f"Wrote bulk extraction manifest: {manifest_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
if args.command == "extract-flat-all":
|
||||||
|
videos = scan_videos(input_root)
|
||||||
|
if args.group_filter:
|
||||||
|
filters = [v.strip().lower() for v in args.group_filter if v.strip()]
|
||||||
|
videos = [meta for meta in videos if any(token in meta.group_name.lower() for token in filters)]
|
||||||
|
|
||||||
|
frame_step = max(1, int(args.frame_step))
|
||||||
|
image_ext = args.image_ext
|
||||||
|
jpg_quality = max(70, min(100, int(args.jpg_quality)))
|
||||||
|
flat_root = output_root
|
||||||
|
|
||||||
|
if args.overwrite and flat_root.exists():
|
||||||
|
shutil.rmtree(flat_root)
|
||||||
|
flat_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
manifest_rows = []
|
||||||
|
for meta in videos:
|
||||||
|
manifest_rows.extend(
|
||||||
|
extract_range_flat(
|
||||||
|
meta=meta,
|
||||||
|
flat_root=flat_root,
|
||||||
|
frame_step=frame_step,
|
||||||
|
image_ext=image_ext,
|
||||||
|
jpg_quality=jpg_quality,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
manifest_path = flat_root / "_flat_manifest.csv"
|
||||||
|
write_csv(
|
||||||
|
manifest_path,
|
||||||
|
manifest_rows,
|
||||||
|
["group_name", "video_relpath", "output_file", "frame_idx", "time_sec", "frame_step", "image_ext"],
|
||||||
|
)
|
||||||
|
print(f"Wrote flat extracted frames into {flat_root}")
|
||||||
|
print(f"Wrote flat extraction manifest: {manifest_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_tool(name: str):
|
||||||
|
if shutil.which(name) is None:
|
||||||
|
raise SystemExit(f"Required tool not found in PATH: {name}")
|
||||||
|
|
||||||
|
|
||||||
|
def scan_videos(input_root: Path):
|
||||||
|
videos = []
|
||||||
|
for path in sorted(input_root.rglob("*")):
|
||||||
|
if not path.is_file() or path.suffix.lower() not in VIDEO_EXTS:
|
||||||
|
continue
|
||||||
|
rel_path = path.relative_to(input_root)
|
||||||
|
group_name = rel_path.parts[0] if len(rel_path.parts) > 1 else "UNGROUPED"
|
||||||
|
info = probe_video(path)
|
||||||
|
videos.append(
|
||||||
|
VideoMeta(
|
||||||
|
abs_path=path,
|
||||||
|
rel_path=rel_path,
|
||||||
|
group_name=group_name,
|
||||||
|
file_name=path.name,
|
||||||
|
size_bytes=path.stat().st_size,
|
||||||
|
width=info["width"],
|
||||||
|
height=info["height"],
|
||||||
|
fps=info["fps"],
|
||||||
|
duration_sec=info["duration_sec"],
|
||||||
|
nb_frames=info["nb_frames"],
|
||||||
|
osd_present=("без osd" not in group_name.lower()),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return videos
|
||||||
|
|
||||||
|
|
||||||
|
def probe_video(path: Path):
|
||||||
|
cmd = [
|
||||||
|
"ffprobe",
|
||||||
|
"-v",
|
||||||
|
"error",
|
||||||
|
"-select_streams",
|
||||||
|
"v:0",
|
||||||
|
"-show_entries",
|
||||||
|
"stream=width,height,r_frame_rate,avg_frame_rate,nb_frames,duration",
|
||||||
|
"-of",
|
||||||
|
"json",
|
||||||
|
str(path),
|
||||||
|
]
|
||||||
|
data = json.loads(run_cmd(cmd))
|
||||||
|
streams = data.get("streams", [])
|
||||||
|
if not streams:
|
||||||
|
raise SystemExit(f"No video stream found: {path}")
|
||||||
|
stream = streams[0]
|
||||||
|
fps_str = stream.get("avg_frame_rate") or stream.get("r_frame_rate") or "0/1"
|
||||||
|
fps = parse_fraction(fps_str)
|
||||||
|
duration = float(stream.get("duration") or 0.0)
|
||||||
|
nb_frames_raw = stream.get("nb_frames")
|
||||||
|
nb_frames = int(nb_frames_raw) if str(nb_frames_raw).isdigit() else int(round(duration * fps))
|
||||||
|
return {
|
||||||
|
"width": int(stream.get("width") or 0),
|
||||||
|
"height": int(stream.get("height") or 0),
|
||||||
|
"fps": fps,
|
||||||
|
"duration_sec": duration,
|
||||||
|
"nb_frames": nb_frames,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def extract_range_to_dir(
|
||||||
|
meta: VideoMeta,
|
||||||
|
output_root: Path,
|
||||||
|
extracted_root: Path,
|
||||||
|
split: str,
|
||||||
|
role: str,
|
||||||
|
start_sec: float,
|
||||||
|
end_sec: float,
|
||||||
|
frame_step: int,
|
||||||
|
image_ext: str,
|
||||||
|
jpg_quality: int,
|
||||||
|
overwrite: bool,
|
||||||
|
interval_name: str,
|
||||||
|
):
|
||||||
|
target_dir = (
|
||||||
|
extracted_root
|
||||||
|
/ sanitize_path_component(split)
|
||||||
|
/ sanitize_path_component(role)
|
||||||
|
/ sanitize_path_component(meta.group_name)
|
||||||
|
/ sanitize_path_component(meta.rel_path.stem)
|
||||||
|
/ sanitize_path_component(interval_name)
|
||||||
|
)
|
||||||
|
|
||||||
|
if overwrite and target_dir.exists():
|
||||||
|
shutil.rmtree(target_dir)
|
||||||
|
|
||||||
|
existing_files = sorted(target_dir.glob(f"*.{image_ext}")) if target_dir.exists() else []
|
||||||
|
if existing_files and (not overwrite):
|
||||||
|
return build_manifest_rows(
|
||||||
|
output_root=output_root,
|
||||||
|
meta=meta,
|
||||||
|
split=split,
|
||||||
|
role=role,
|
||||||
|
files=existing_files,
|
||||||
|
frame_step=frame_step,
|
||||||
|
)
|
||||||
|
|
||||||
|
target_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
tmp_dir = target_dir / "_tmp"
|
||||||
|
if tmp_dir.exists():
|
||||||
|
shutil.rmtree(tmp_dir)
|
||||||
|
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
tmp_pattern = tmp_dir / f"tmp_%06d.{image_ext}"
|
||||||
|
cmd = [
|
||||||
|
"ffmpeg",
|
||||||
|
"-hide_banner",
|
||||||
|
"-loglevel",
|
||||||
|
"warning",
|
||||||
|
"-y",
|
||||||
|
"-ss",
|
||||||
|
f"{start_sec:.3f}",
|
||||||
|
"-to",
|
||||||
|
f"{end_sec:.3f}",
|
||||||
|
"-i",
|
||||||
|
str(meta.abs_path),
|
||||||
|
"-vf",
|
||||||
|
f"select=not(mod(n\\,{frame_step}))",
|
||||||
|
"-vsync",
|
||||||
|
"vfr",
|
||||||
|
]
|
||||||
|
if image_ext == "jpg":
|
||||||
|
cmd.extend(["-q:v", str(jpeg_qscale(jpg_quality))])
|
||||||
|
cmd.append(str(tmp_pattern))
|
||||||
|
run_cmd(cmd)
|
||||||
|
|
||||||
|
files = sorted(tmp_dir.glob(f"*.{image_ext}"))
|
||||||
|
start_frame_idx = max(0, int(round(start_sec * meta.fps)))
|
||||||
|
final_files = []
|
||||||
|
for idx, tmp_file in enumerate(files):
|
||||||
|
frame_idx = start_frame_idx + idx * frame_step
|
||||||
|
time_sec = frame_idx / max(1e-6, meta.fps)
|
||||||
|
final_name = (
|
||||||
|
f"{sanitize_path_component(meta.rel_path.stem)}__"
|
||||||
|
f"f{frame_idx:06d}__t{time_sec:010.3f}.{image_ext}"
|
||||||
|
)
|
||||||
|
final_path = target_dir / final_name
|
||||||
|
tmp_file.replace(final_path)
|
||||||
|
final_files.append(final_path)
|
||||||
|
|
||||||
|
shutil.rmtree(tmp_dir)
|
||||||
|
return build_manifest_rows(
|
||||||
|
output_root=output_root,
|
||||||
|
meta=meta,
|
||||||
|
split=split,
|
||||||
|
role=role,
|
||||||
|
files=final_files,
|
||||||
|
frame_step=frame_step,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def build_manifest_rows(output_root: Path, meta: VideoMeta, split: str, role: str, files, frame_step: int):
|
||||||
|
rows = []
|
||||||
|
for file_path in files:
|
||||||
|
frame_idx, time_sec = parse_frame_info_from_name(file_path.stem, meta.fps)
|
||||||
|
rows.append(
|
||||||
|
{
|
||||||
|
"split": split,
|
||||||
|
"role": role,
|
||||||
|
"group_name": meta.group_name,
|
||||||
|
"video_relpath": normalize_relpath(meta.rel_path),
|
||||||
|
"output_file": normalize_relpath(file_path.relative_to(output_root)),
|
||||||
|
"frame_idx": frame_idx,
|
||||||
|
"time_sec": round(time_sec, 3),
|
||||||
|
"frame_step": frame_step,
|
||||||
|
"image_ext": file_path.suffix.lstrip(".").lower(),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return rows
|
||||||
|
|
||||||
|
|
||||||
|
def parse_frame_info_from_name(stem: str, fps: float):
|
||||||
|
frame_idx = 0
|
||||||
|
time_sec = 0.0
|
||||||
|
for part in stem.split("__"):
|
||||||
|
if part.startswith("f") and part[1:].isdigit():
|
||||||
|
frame_idx = int(part[1:])
|
||||||
|
elif part.startswith("t"):
|
||||||
|
try:
|
||||||
|
time_sec = float(part[1:])
|
||||||
|
except ValueError:
|
||||||
|
time_sec = frame_idx / max(1e-6, fps)
|
||||||
|
if time_sec <= 0.0 and frame_idx > 0:
|
||||||
|
time_sec = frame_idx / max(1e-6, fps)
|
||||||
|
return frame_idx, time_sec
|
||||||
|
|
||||||
|
|
||||||
|
def extract_range_flat(
|
||||||
|
meta: VideoMeta,
|
||||||
|
flat_root: Path,
|
||||||
|
frame_step: int,
|
||||||
|
image_ext: str,
|
||||||
|
jpg_quality: int,
|
||||||
|
):
|
||||||
|
tmp_dir = flat_root / f"_tmp_{sanitize_path_component(meta.rel_path.stem)}"
|
||||||
|
if tmp_dir.exists():
|
||||||
|
shutil.rmtree(tmp_dir)
|
||||||
|
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
tmp_pattern = tmp_dir / f"tmp_%06d.{image_ext}"
|
||||||
|
cmd = [
|
||||||
|
"ffmpeg",
|
||||||
|
"-hide_banner",
|
||||||
|
"-loglevel",
|
||||||
|
"warning",
|
||||||
|
"-y",
|
||||||
|
"-i",
|
||||||
|
str(meta.abs_path),
|
||||||
|
"-vf",
|
||||||
|
f"select=not(mod(n\\,{frame_step}))",
|
||||||
|
"-vsync",
|
||||||
|
"vfr",
|
||||||
|
]
|
||||||
|
if image_ext == "jpg":
|
||||||
|
cmd.extend(["-q:v", str(jpeg_qscale(jpg_quality))])
|
||||||
|
cmd.append(str(tmp_pattern))
|
||||||
|
run_cmd(cmd)
|
||||||
|
|
||||||
|
rows = []
|
||||||
|
files = sorted(tmp_dir.glob(f"*.{image_ext}"))
|
||||||
|
safe_group = sanitize_path_component(meta.group_name)
|
||||||
|
safe_video = sanitize_path_component(meta.rel_path.stem)
|
||||||
|
for idx, tmp_file in enumerate(files):
|
||||||
|
frame_idx = idx * frame_step
|
||||||
|
time_sec = frame_idx / max(1e-6, meta.fps)
|
||||||
|
final_name = f"{safe_group}__{safe_video}__f{frame_idx:06d}__t{time_sec:010.3f}.{image_ext}"
|
||||||
|
final_path = flat_root / final_name
|
||||||
|
tmp_file.replace(final_path)
|
||||||
|
rows.append(
|
||||||
|
{
|
||||||
|
"group_name": meta.group_name,
|
||||||
|
"video_relpath": normalize_relpath(meta.rel_path),
|
||||||
|
"output_file": final_path.name,
|
||||||
|
"frame_idx": frame_idx,
|
||||||
|
"time_sec": round(time_sec, 3),
|
||||||
|
"frame_step": frame_step,
|
||||||
|
"image_ext": image_ext,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
shutil.rmtree(tmp_dir)
|
||||||
|
return rows
|
||||||
|
|
||||||
|
|
||||||
|
def parse_fraction(value: str):
|
||||||
|
num, den = value.split("/")
|
||||||
|
den_v = float(den)
|
||||||
|
if abs(den_v) < 1e-9:
|
||||||
|
return 0.0
|
||||||
|
return float(num) / den_v
|
||||||
|
|
||||||
|
|
||||||
|
def write_catalog_outputs(videos, output_root: Path):
|
||||||
|
output_root.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
catalog_rows = []
|
||||||
|
summary_map = {}
|
||||||
|
template_rows = []
|
||||||
|
|
||||||
|
for meta in videos:
|
||||||
|
rel_path_str = normalize_relpath(meta.rel_path)
|
||||||
|
catalog_rows.append(
|
||||||
|
{
|
||||||
|
"group_name": meta.group_name,
|
||||||
|
"osd_present": int(meta.osd_present),
|
||||||
|
"video_relpath": rel_path_str,
|
||||||
|
"file_name": meta.file_name,
|
||||||
|
"size_bytes": meta.size_bytes,
|
||||||
|
"size_gb": round(meta.size_bytes / (1024 ** 3), 4),
|
||||||
|
"duration_sec": round(meta.duration_sec, 3),
|
||||||
|
"duration_min": round(meta.duration_sec / 60.0, 2),
|
||||||
|
"fps": round(meta.fps, 6),
|
||||||
|
"width": meta.width,
|
||||||
|
"height": meta.height,
|
||||||
|
"nb_frames": meta.nb_frames,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
key = meta.group_name
|
||||||
|
info = summary_map.setdefault(key, {"videos": 0, "bytes": 0, "duration": 0.0})
|
||||||
|
info["videos"] += 1
|
||||||
|
info["bytes"] += meta.size_bytes
|
||||||
|
info["duration"] += meta.duration_sec
|
||||||
|
|
||||||
|
template_rows.append(
|
||||||
|
{
|
||||||
|
"enabled": "no",
|
||||||
|
"split": "train",
|
||||||
|
"role": "positive",
|
||||||
|
"video_relpath": rel_path_str,
|
||||||
|
"group_name": meta.group_name,
|
||||||
|
"start_sec": "",
|
||||||
|
"end_sec": "",
|
||||||
|
"frame_step": "2",
|
||||||
|
"image_ext": "png",
|
||||||
|
"jpg_quality": "95",
|
||||||
|
"notes": "Duplicate row for more intervals or use role=hard_negative for sparse negatives.",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
summary_rows = []
|
||||||
|
for group_name in sorted(summary_map):
|
||||||
|
info = summary_map[group_name]
|
||||||
|
summary_rows.append(
|
||||||
|
{
|
||||||
|
"group_name": group_name,
|
||||||
|
"video_count": info["videos"],
|
||||||
|
"total_size_gb": round(info["bytes"] / (1024 ** 3), 3),
|
||||||
|
"total_duration_min": round(info["duration"] / 60.0, 2),
|
||||||
|
"osd_present": int("без osd" not in group_name.lower()),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
write_csv(
|
||||||
|
output_root / "video_catalog.csv",
|
||||||
|
catalog_rows,
|
||||||
|
[
|
||||||
|
"group_name",
|
||||||
|
"osd_present",
|
||||||
|
"video_relpath",
|
||||||
|
"file_name",
|
||||||
|
"size_bytes",
|
||||||
|
"size_gb",
|
||||||
|
"duration_sec",
|
||||||
|
"duration_min",
|
||||||
|
"fps",
|
||||||
|
"width",
|
||||||
|
"height",
|
||||||
|
"nb_frames",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
write_csv(
|
||||||
|
output_root / "group_summary.csv",
|
||||||
|
summary_rows,
|
||||||
|
["group_name", "video_count", "total_size_gb", "total_duration_min", "osd_present"],
|
||||||
|
)
|
||||||
|
write_csv(
|
||||||
|
output_root / "intervals_template.csv",
|
||||||
|
template_rows,
|
||||||
|
[
|
||||||
|
"enabled",
|
||||||
|
"split",
|
||||||
|
"role",
|
||||||
|
"video_relpath",
|
||||||
|
"group_name",
|
||||||
|
"start_sec",
|
||||||
|
"end_sec",
|
||||||
|
"frame_step",
|
||||||
|
"image_ext",
|
||||||
|
"jpg_quality",
|
||||||
|
"notes",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def write_csv(path: Path, rows, fieldnames):
|
||||||
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
with path.open("w", newline="", encoding="utf-8-sig") as f:
|
||||||
|
writer = csv.DictWriter(f, fieldnames=fieldnames)
|
||||||
|
writer.writeheader()
|
||||||
|
for row in rows:
|
||||||
|
writer.writerow(row)
|
||||||
|
|
||||||
|
|
||||||
|
def load_intervals(path: Path):
|
||||||
|
with path.open("r", newline="", encoding="utf-8-sig") as f:
|
||||||
|
return list(csv.DictReader(f))
|
||||||
|
|
||||||
|
|
||||||
|
def is_enabled(value: str):
|
||||||
|
return str(value).strip().lower() in {"1", "true", "yes", "y", "on"}
|
||||||
|
|
||||||
|
|
||||||
|
def clamp_float(value, min_value: float, max_value: float):
|
||||||
|
value_f = float(value)
|
||||||
|
return max(min_value, min(max_value, value_f))
|
||||||
|
|
||||||
|
|
||||||
|
def format_interval_name(start_sec: float, end_sec: float, frame_step: int):
|
||||||
|
return f"s{int(math.floor(start_sec)):06d}_e{int(math.ceil(end_sec)):06d}_step{frame_step}"
|
||||||
|
|
||||||
|
|
||||||
|
def sanitize_path_component(value: str):
|
||||||
|
keep = []
|
||||||
|
for ch in value:
|
||||||
|
if ch in '<>:"/\\|?*':
|
||||||
|
keep.append("_")
|
||||||
|
else:
|
||||||
|
keep.append(ch)
|
||||||
|
return "".join(keep).strip().rstrip(".")
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_relpath(path):
|
||||||
|
return str(path).replace("\\", "/")
|
||||||
|
|
||||||
|
|
||||||
|
def jpeg_qscale(quality_pct: int):
|
||||||
|
quality_pct = max(70, min(100, int(quality_pct)))
|
||||||
|
if quality_pct >= 98:
|
||||||
|
return 1
|
||||||
|
if quality_pct >= 92:
|
||||||
|
return 2
|
||||||
|
if quality_pct >= 88:
|
||||||
|
return 3
|
||||||
|
if quality_pct >= 82:
|
||||||
|
return 4
|
||||||
|
return 5
|
||||||
|
|
||||||
|
|
||||||
|
def run_cmd(cmd):
|
||||||
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||||
|
if result.returncode != 0:
|
||||||
|
raise SystemExit(
|
||||||
|
f"Command failed ({result.returncode}): {' '.join(cmd)}\n"
|
||||||
|
f"stdout:\n{result.stdout}\n"
|
||||||
|
f"stderr:\n{result.stderr}"
|
||||||
|
)
|
||||||
|
return result.stdout
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,159 @@
|
|||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace:
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Split a flat image folder by source-video prefixes from CSV."
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--flat-root",
|
||||||
|
type=Path,
|
||||||
|
required=True,
|
||||||
|
help="Flat folder with extracted images.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--split-csv",
|
||||||
|
type=Path,
|
||||||
|
required=True,
|
||||||
|
help="CSV with recommended_split and flat_prefix columns.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--output-root",
|
||||||
|
type=Path,
|
||||||
|
required=True,
|
||||||
|
help="Output folder that will contain train/valid/test.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--image-ext",
|
||||||
|
default="jpg",
|
||||||
|
help="Image extension in the flat folder, default: jpg.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--mode",
|
||||||
|
choices=("hardlink", "copy"),
|
||||||
|
default="hardlink",
|
||||||
|
help="Write mode for split files. hardlink avoids doubling disk usage.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--overwrite",
|
||||||
|
action="store_true",
|
||||||
|
help="Remove existing output_root before writing.",
|
||||||
|
)
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def load_split_rows(csv_path: Path) -> list[dict]:
|
||||||
|
with csv_path.open("r", encoding="utf-8-sig", newline="") as f:
|
||||||
|
rows = list(csv.DictReader(f))
|
||||||
|
required = {"recommended_split", "flat_prefix"}
|
||||||
|
if not rows:
|
||||||
|
raise ValueError(f"No rows found in {csv_path}")
|
||||||
|
missing = required - set(rows[0].keys())
|
||||||
|
if missing:
|
||||||
|
raise ValueError(f"Missing required CSV columns: {sorted(missing)}")
|
||||||
|
return rows
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_clean_dir(path: Path, overwrite: bool) -> None:
|
||||||
|
if path.exists():
|
||||||
|
if not overwrite:
|
||||||
|
raise FileExistsError(
|
||||||
|
f"Output path already exists: {path}. Use --overwrite to replace it."
|
||||||
|
)
|
||||||
|
shutil.rmtree(path)
|
||||||
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
|
||||||
|
def link_or_copy(src: Path, dst: Path, mode: str) -> str:
|
||||||
|
if mode == "hardlink":
|
||||||
|
try:
|
||||||
|
os.link(src, dst)
|
||||||
|
return "hardlink"
|
||||||
|
except OSError:
|
||||||
|
shutil.copy2(src, dst)
|
||||||
|
return "copy_fallback"
|
||||||
|
shutil.copy2(src, dst)
|
||||||
|
return "copy"
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
args = parse_args()
|
||||||
|
rows = load_split_rows(args.split_csv)
|
||||||
|
ext = args.image_ext.lower().lstrip(".")
|
||||||
|
|
||||||
|
ensure_clean_dir(args.output_root, args.overwrite)
|
||||||
|
manifest_path = args.output_root / "_split_manifest.csv"
|
||||||
|
|
||||||
|
split_name_map = {
|
||||||
|
"train": "train",
|
||||||
|
"val": "valid",
|
||||||
|
"valid": "valid",
|
||||||
|
"test": "test",
|
||||||
|
}
|
||||||
|
|
||||||
|
total_written = 0
|
||||||
|
total_missing = 0
|
||||||
|
|
||||||
|
with manifest_path.open("w", encoding="utf-8-sig", newline="") as mf:
|
||||||
|
writer = csv.DictWriter(
|
||||||
|
mf,
|
||||||
|
fieldnames=[
|
||||||
|
"recommended_split",
|
||||||
|
"output_split",
|
||||||
|
"group_name",
|
||||||
|
"file_name",
|
||||||
|
"flat_prefix",
|
||||||
|
"written_count",
|
||||||
|
"missing",
|
||||||
|
"write_mode",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
writer.writeheader()
|
||||||
|
|
||||||
|
for row in rows:
|
||||||
|
split_raw = row["recommended_split"].strip().lower()
|
||||||
|
output_split = split_name_map.get(split_raw)
|
||||||
|
if output_split is None:
|
||||||
|
raise ValueError(f"Unsupported split value: {row['recommended_split']}")
|
||||||
|
|
||||||
|
prefix = row["flat_prefix"].strip()
|
||||||
|
pattern = f"{prefix}__f*.{ext}"
|
||||||
|
matches = sorted(args.flat_root.glob(pattern))
|
||||||
|
|
||||||
|
out_dir = args.output_root / output_split
|
||||||
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
write_mode = ""
|
||||||
|
for src in matches:
|
||||||
|
dst = out_dir / src.name
|
||||||
|
write_mode = link_or_copy(src, dst, args.mode)
|
||||||
|
|
||||||
|
missing = 0 if matches else 1
|
||||||
|
total_missing += missing
|
||||||
|
total_written += len(matches)
|
||||||
|
|
||||||
|
writer.writerow(
|
||||||
|
{
|
||||||
|
"recommended_split": row["recommended_split"],
|
||||||
|
"output_split": output_split,
|
||||||
|
"group_name": row.get("group_name", ""),
|
||||||
|
"file_name": row.get("file_name", ""),
|
||||||
|
"flat_prefix": prefix,
|
||||||
|
"written_count": len(matches),
|
||||||
|
"missing": missing,
|
||||||
|
"write_mode": write_mode,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f"output_root={args.output_root}")
|
||||||
|
print(f"manifest={manifest_path}")
|
||||||
|
print(f"total_written={total_written}")
|
||||||
|
print(f"missing_prefixes={total_missing}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,347 @@
|
|||||||
|
import csv
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
from collections import Counter
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def _build_unique_path(base_path, default_suffix):
|
||||||
|
base = Path(base_path)
|
||||||
|
if base.suffix == "":
|
||||||
|
base = base.with_suffix(default_suffix)
|
||||||
|
|
||||||
|
parent = base.parent if str(base.parent) not in ("", ".") else Path.cwd()
|
||||||
|
parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
stamp = time.strftime("%Y%m%d_%H%M%S")
|
||||||
|
candidate = parent / f"{base.stem}_{stamp}{base.suffix}"
|
||||||
|
attempt = 1
|
||||||
|
while candidate.exists():
|
||||||
|
candidate = parent / f"{base.stem}_{stamp}_{attempt:02d}{base.suffix}"
|
||||||
|
attempt += 1
|
||||||
|
return candidate
|
||||||
|
|
||||||
|
|
||||||
|
def _box_to_metrics(box):
|
||||||
|
if box is None:
|
||||||
|
return "", "", "", ""
|
||||||
|
|
||||||
|
x1, y1, x2, y2 = [float(v) for v in box]
|
||||||
|
cx = 0.5 * (x1 + x2)
|
||||||
|
cy = 0.5 * (y1 + y2)
|
||||||
|
w = max(0.0, x2 - x1)
|
||||||
|
h = max(0.0, y2 - y1)
|
||||||
|
return round(cx, 3), round(cy, 3), round(w, 3), round(h, 3)
|
||||||
|
|
||||||
|
|
||||||
|
class TrackingDecisionLogger:
|
||||||
|
FIELDNAMES = [
|
||||||
|
"frame_id",
|
||||||
|
"timestamp_sec",
|
||||||
|
"dt_sec",
|
||||||
|
"source_kind",
|
||||||
|
"status",
|
||||||
|
"confirmed",
|
||||||
|
"target_id",
|
||||||
|
"target_track_score",
|
||||||
|
"locked",
|
||||||
|
"locked_cx",
|
||||||
|
"locked_cy",
|
||||||
|
"locked_w",
|
||||||
|
"locked_h",
|
||||||
|
"pred_cx",
|
||||||
|
"pred_cy",
|
||||||
|
"pred_w",
|
||||||
|
"pred_h",
|
||||||
|
"hit_streak",
|
||||||
|
"miss_streak",
|
||||||
|
"acquire_score",
|
||||||
|
"preacq_hits",
|
||||||
|
"switch_candidate_hits",
|
||||||
|
"best_score",
|
||||||
|
"have_yolo",
|
||||||
|
"yolo_mode",
|
||||||
|
"yolo_raw_count",
|
||||||
|
"det_count",
|
||||||
|
"merged_part_count",
|
||||||
|
"infer_ms",
|
||||||
|
"stale_lock_active",
|
||||||
|
"stale_lock_reason",
|
||||||
|
"fast_handoff_active",
|
||||||
|
"guidance_reset_event",
|
||||||
|
"used_prediction_hold",
|
||||||
|
"klt_valid",
|
||||||
|
"klt_quality",
|
||||||
|
"klt_points",
|
||||||
|
"speed",
|
||||||
|
"adaptive_chase_stage",
|
||||||
|
"approach_active",
|
||||||
|
"close_force_fullscan",
|
||||||
|
"fast_maneuver_guard",
|
||||||
|
"motion_active_zones",
|
||||||
|
"wavelet_active",
|
||||||
|
"wavelet_energy",
|
||||||
|
"wavelet_bonus",
|
||||||
|
"wavelet_hits",
|
||||||
|
"autogaze_ready",
|
||||||
|
"autogaze_stale",
|
||||||
|
"autogaze_age",
|
||||||
|
"autogaze_cells",
|
||||||
|
"autogaze_ms",
|
||||||
|
"guidance_active",
|
||||||
|
"guidance_status",
|
||||||
|
"guidance_confidence",
|
||||||
|
"guidance_error_x",
|
||||||
|
"guidance_error_y",
|
||||||
|
"guidance_cmd_x",
|
||||||
|
"guidance_cmd_y",
|
||||||
|
"guidance_on_target",
|
||||||
|
"events",
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, enabled, csv_path, summary_path, flush_every=30):
|
||||||
|
self.enabled = bool(enabled)
|
||||||
|
self.flush_every = max(1, int(flush_every))
|
||||||
|
self.csv_path = None
|
||||||
|
self.summary_path = None
|
||||||
|
self._csv_file = None
|
||||||
|
self._writer = None
|
||||||
|
self._rows_since_flush = 0
|
||||||
|
|
||||||
|
self.prev_confirmed = None
|
||||||
|
self.prev_locked = None
|
||||||
|
self.prev_target_id = None
|
||||||
|
|
||||||
|
self.frames = 0
|
||||||
|
self.confirmed_frames = 0
|
||||||
|
self.locked_frames = 0
|
||||||
|
self.prediction_hold_frames = 0
|
||||||
|
self.stale_lock_frames = 0
|
||||||
|
self.fast_handoff_frames = 0
|
||||||
|
self.autogaze_used_frames = 0
|
||||||
|
self.wavelet_active_frames = 0
|
||||||
|
self.guidance_active_frames = 0
|
||||||
|
self.guidance_on_target_frames = 0
|
||||||
|
self.guidance_reset_events = 0
|
||||||
|
self.guidance_confidence_sum = 0.0
|
||||||
|
self.max_hit_streak = 0
|
||||||
|
self.max_miss_streak = 0
|
||||||
|
self.klt_quality_confirmed_sum = 0.0
|
||||||
|
self.klt_quality_confirmed_count = 0
|
||||||
|
self.infer_ms_sum = 0.0
|
||||||
|
self.yolo_mode_counts = Counter()
|
||||||
|
self.event_counts = Counter()
|
||||||
|
|
||||||
|
if not self.enabled:
|
||||||
|
return
|
||||||
|
|
||||||
|
self.csv_path = _build_unique_path(csv_path, ".csv")
|
||||||
|
self.summary_path = _build_unique_path(summary_path, ".json")
|
||||||
|
self._csv_file = self.csv_path.open("w", newline="", encoding="utf-8")
|
||||||
|
self._writer = csv.DictWriter(self._csv_file, fieldnames=self.FIELDNAMES)
|
||||||
|
self._writer.writeheader()
|
||||||
|
|
||||||
|
def status_line(self):
|
||||||
|
if not self.enabled:
|
||||||
|
return "Tracking decision log disabled"
|
||||||
|
return f"Tracking decision log: csv={self.csv_path} summary={self.summary_path}"
|
||||||
|
|
||||||
|
def log_frame(self, **kwargs):
|
||||||
|
if not self.enabled or self._writer is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
confirmed = bool(kwargs.get("confirmed", False))
|
||||||
|
locked = bool(kwargs.get("locked", False))
|
||||||
|
target_id = kwargs.get("target_id", None)
|
||||||
|
|
||||||
|
events = []
|
||||||
|
if self.prev_confirmed is not None:
|
||||||
|
if (not self.prev_confirmed) and confirmed:
|
||||||
|
events.append("confirm_enter")
|
||||||
|
elif self.prev_confirmed and (not confirmed):
|
||||||
|
events.append("confirm_exit")
|
||||||
|
|
||||||
|
if (not self.prev_locked) and locked:
|
||||||
|
events.append("lock_acquired")
|
||||||
|
elif self.prev_locked and (not locked):
|
||||||
|
events.append("lock_lost")
|
||||||
|
|
||||||
|
if (
|
||||||
|
confirmed
|
||||||
|
and (self.prev_target_id is not None)
|
||||||
|
and (target_id is not None)
|
||||||
|
and (self.prev_target_id != target_id)
|
||||||
|
):
|
||||||
|
events.append("target_switch")
|
||||||
|
guidance_reset_event = str(kwargs.get("guidance_reset_event", "")).strip()
|
||||||
|
if guidance_reset_event == "stale_lock_break":
|
||||||
|
events.append("stale_lock_break")
|
||||||
|
if bool(kwargs.get("fast_handoff_active", False)) and ("target_switch" in events):
|
||||||
|
events.append("fast_handoff_switch")
|
||||||
|
if guidance_reset_event:
|
||||||
|
events.append("guidance_reset")
|
||||||
|
|
||||||
|
locked_cx, locked_cy, locked_w, locked_h = _box_to_metrics(kwargs.get("locked_box"))
|
||||||
|
pred_cx, pred_cy, pred_w, pred_h = _box_to_metrics(kwargs.get("pred_box"))
|
||||||
|
|
||||||
|
row = {
|
||||||
|
"frame_id": int(kwargs.get("frame_id", -1)),
|
||||||
|
"timestamp_sec": round(float(kwargs.get("timestamp_sec", 0.0)), 6),
|
||||||
|
"dt_sec": round(float(kwargs.get("dt_sec", 0.0)), 6),
|
||||||
|
"source_kind": kwargs.get("source_kind", ""),
|
||||||
|
"status": kwargs.get("status", ""),
|
||||||
|
"confirmed": int(confirmed),
|
||||||
|
"target_id": "" if target_id is None else int(target_id),
|
||||||
|
"target_track_score": self._maybe_float(kwargs.get("target_track_score")),
|
||||||
|
"locked": int(locked),
|
||||||
|
"locked_cx": locked_cx,
|
||||||
|
"locked_cy": locked_cy,
|
||||||
|
"locked_w": locked_w,
|
||||||
|
"locked_h": locked_h,
|
||||||
|
"pred_cx": pred_cx,
|
||||||
|
"pred_cy": pred_cy,
|
||||||
|
"pred_w": pred_w,
|
||||||
|
"pred_h": pred_h,
|
||||||
|
"hit_streak": int(kwargs.get("hit_streak", 0)),
|
||||||
|
"miss_streak": int(kwargs.get("miss_streak", 0)),
|
||||||
|
"acquire_score": int(kwargs.get("acquire_score", 0)),
|
||||||
|
"preacq_hits": int(kwargs.get("preacq_hits", 0)),
|
||||||
|
"switch_candidate_hits": int(kwargs.get("switch_candidate_hits", 0)),
|
||||||
|
"best_score": self._maybe_float(kwargs.get("best_score")),
|
||||||
|
"have_yolo": int(bool(kwargs.get("have_yolo", False))),
|
||||||
|
"yolo_mode": kwargs.get("yolo_mode", ""),
|
||||||
|
"yolo_raw_count": int(kwargs.get("yolo_raw_count", 0)),
|
||||||
|
"det_count": int(kwargs.get("det_count", 0)),
|
||||||
|
"merged_part_count": int(kwargs.get("merged_part_count", 0)),
|
||||||
|
"infer_ms": round(float(kwargs.get("infer_ms", 0.0)), 3),
|
||||||
|
"stale_lock_active": int(bool(kwargs.get("stale_lock_active", False))),
|
||||||
|
"stale_lock_reason": kwargs.get("stale_lock_reason", ""),
|
||||||
|
"fast_handoff_active": int(bool(kwargs.get("fast_handoff_active", False))),
|
||||||
|
"guidance_reset_event": guidance_reset_event,
|
||||||
|
"used_prediction_hold": int(bool(kwargs.get("used_prediction_hold", False))),
|
||||||
|
"klt_valid": int(bool(kwargs.get("klt_valid", False))),
|
||||||
|
"klt_quality": round(float(kwargs.get("klt_quality", 0.0)), 4),
|
||||||
|
"klt_points": int(kwargs.get("klt_points", 0)),
|
||||||
|
"speed": round(float(kwargs.get("speed", 0.0)), 3),
|
||||||
|
"adaptive_chase_stage": kwargs.get("adaptive_chase_stage", ""),
|
||||||
|
"approach_active": int(bool(kwargs.get("approach_active", False))),
|
||||||
|
"close_force_fullscan": int(bool(kwargs.get("close_force_fullscan", False))),
|
||||||
|
"fast_maneuver_guard": int(bool(kwargs.get("fast_maneuver_guard", False))),
|
||||||
|
"motion_active_zones": int(kwargs.get("motion_active_zones", 0)),
|
||||||
|
"wavelet_active": int(bool(kwargs.get("wavelet_active", False))),
|
||||||
|
"wavelet_energy": round(float(kwargs.get("wavelet_energy", 0.0)), 4),
|
||||||
|
"wavelet_bonus": round(float(kwargs.get("wavelet_bonus", 0.0)), 4),
|
||||||
|
"wavelet_hits": int(kwargs.get("wavelet_hits", 0)),
|
||||||
|
"autogaze_ready": int(bool(kwargs.get("autogaze_ready", False))),
|
||||||
|
"autogaze_stale": int(bool(kwargs.get("autogaze_stale", True))),
|
||||||
|
"autogaze_age": int(kwargs.get("autogaze_age", -1)),
|
||||||
|
"autogaze_cells": int(kwargs.get("autogaze_cells", 0)),
|
||||||
|
"autogaze_ms": round(float(kwargs.get("autogaze_ms", 0.0)), 3),
|
||||||
|
"guidance_active": int(bool(kwargs.get("guidance_active", False))),
|
||||||
|
"guidance_status": kwargs.get("guidance_status", ""),
|
||||||
|
"guidance_confidence": round(float(kwargs.get("guidance_confidence", 0.0)), 4),
|
||||||
|
"guidance_error_x": round(float(kwargs.get("guidance_error_x", 0.0)), 4),
|
||||||
|
"guidance_error_y": round(float(kwargs.get("guidance_error_y", 0.0)), 4),
|
||||||
|
"guidance_cmd_x": round(float(kwargs.get("guidance_cmd_x", 0.0)), 4),
|
||||||
|
"guidance_cmd_y": round(float(kwargs.get("guidance_cmd_y", 0.0)), 4),
|
||||||
|
"guidance_on_target": int(bool(kwargs.get("guidance_on_target", False))),
|
||||||
|
"events": "|".join(events),
|
||||||
|
}
|
||||||
|
|
||||||
|
self._writer.writerow(row)
|
||||||
|
self._rows_since_flush += 1
|
||||||
|
if self._rows_since_flush >= self.flush_every:
|
||||||
|
self._csv_file.flush()
|
||||||
|
self._rows_since_flush = 0
|
||||||
|
|
||||||
|
self.frames += 1
|
||||||
|
self.confirmed_frames += int(confirmed)
|
||||||
|
self.locked_frames += int(locked)
|
||||||
|
self.prediction_hold_frames += int(bool(kwargs.get("used_prediction_hold", False)))
|
||||||
|
self.stale_lock_frames += int(bool(kwargs.get("stale_lock_active", False)))
|
||||||
|
self.fast_handoff_frames += int(bool(kwargs.get("fast_handoff_active", False)))
|
||||||
|
self.autogaze_used_frames += int(bool(kwargs.get("autogaze_ready", False)) and (not bool(kwargs.get("autogaze_stale", True))))
|
||||||
|
self.wavelet_active_frames += int(bool(kwargs.get("wavelet_active", False)))
|
||||||
|
self.guidance_active_frames += int(bool(kwargs.get("guidance_active", False)))
|
||||||
|
self.guidance_on_target_frames += int(bool(kwargs.get("guidance_on_target", False)))
|
||||||
|
self.guidance_reset_events += int(bool(guidance_reset_event))
|
||||||
|
self.guidance_confidence_sum += float(kwargs.get("guidance_confidence", 0.0))
|
||||||
|
self.max_hit_streak = max(self.max_hit_streak, int(kwargs.get("hit_streak", 0)))
|
||||||
|
self.max_miss_streak = max(self.max_miss_streak, int(kwargs.get("miss_streak", 0)))
|
||||||
|
self.infer_ms_sum += float(kwargs.get("infer_ms", 0.0))
|
||||||
|
|
||||||
|
if confirmed:
|
||||||
|
self.klt_quality_confirmed_sum += float(kwargs.get("klt_quality", 0.0))
|
||||||
|
self.klt_quality_confirmed_count += 1
|
||||||
|
|
||||||
|
yolo_mode = str(kwargs.get("yolo_mode", "")).strip()
|
||||||
|
if yolo_mode:
|
||||||
|
self.yolo_mode_counts[yolo_mode] += 1
|
||||||
|
for event in events:
|
||||||
|
self.event_counts[event] += 1
|
||||||
|
|
||||||
|
self.prev_confirmed = confirmed
|
||||||
|
self.prev_locked = locked
|
||||||
|
self.prev_target_id = target_id
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
if not self.enabled:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if self._csv_file is not None:
|
||||||
|
self._csv_file.flush()
|
||||||
|
self._csv_file.close()
|
||||||
|
self._csv_file = None
|
||||||
|
|
||||||
|
avg_klt_quality_confirmed = 0.0
|
||||||
|
if self.klt_quality_confirmed_count > 0:
|
||||||
|
avg_klt_quality_confirmed = self.klt_quality_confirmed_sum / float(self.klt_quality_confirmed_count)
|
||||||
|
|
||||||
|
avg_infer_ms = 0.0
|
||||||
|
if self.frames > 0:
|
||||||
|
avg_infer_ms = self.infer_ms_sum / float(self.frames)
|
||||||
|
|
||||||
|
avg_guidance_confidence = 0.0
|
||||||
|
if self.frames > 0:
|
||||||
|
avg_guidance_confidence = self.guidance_confidence_sum / float(self.frames)
|
||||||
|
|
||||||
|
summary = {
|
||||||
|
"frames": self.frames,
|
||||||
|
"confirmed_frames": self.confirmed_frames,
|
||||||
|
"recover_frames": max(0, self.frames - self.confirmed_frames),
|
||||||
|
"locked_frames": self.locked_frames,
|
||||||
|
"prediction_hold_frames": self.prediction_hold_frames,
|
||||||
|
"stale_lock_frames": self.stale_lock_frames,
|
||||||
|
"fast_handoff_frames": self.fast_handoff_frames,
|
||||||
|
"autogaze_used_frames": self.autogaze_used_frames,
|
||||||
|
"wavelet_active_frames": self.wavelet_active_frames,
|
||||||
|
"guidance_active_frames": self.guidance_active_frames,
|
||||||
|
"guidance_on_target_frames": self.guidance_on_target_frames,
|
||||||
|
"guidance_reset_events": self.guidance_reset_events,
|
||||||
|
"confirm_entries": int(self.event_counts.get("confirm_enter", 0)),
|
||||||
|
"confirm_exits": int(self.event_counts.get("confirm_exit", 0)),
|
||||||
|
"lock_acquired_events": int(self.event_counts.get("lock_acquired", 0)),
|
||||||
|
"lock_lost_events": int(self.event_counts.get("lock_lost", 0)),
|
||||||
|
"target_switches": int(self.event_counts.get("target_switch", 0)),
|
||||||
|
"max_hit_streak": self.max_hit_streak,
|
||||||
|
"max_miss_streak": self.max_miss_streak,
|
||||||
|
"avg_klt_quality_confirmed": round(avg_klt_quality_confirmed, 6),
|
||||||
|
"avg_guidance_confidence": round(avg_guidance_confidence, 6),
|
||||||
|
"avg_infer_ms": round(avg_infer_ms, 6),
|
||||||
|
"yolo_mode_counts": dict(self.yolo_mode_counts),
|
||||||
|
"event_counts": dict(self.event_counts),
|
||||||
|
"csv_path": str(self.csv_path),
|
||||||
|
"summary_path": str(self.summary_path),
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.summary_path is not None:
|
||||||
|
with self.summary_path.open("w", encoding="utf-8") as f:
|
||||||
|
json.dump(summary, f, ensure_ascii=True, indent=2)
|
||||||
|
|
||||||
|
return summary
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _maybe_float(value):
|
||||||
|
if value is None:
|
||||||
|
return ""
|
||||||
|
return round(float(value), 6)
|
||||||
@ -0,0 +1,34 @@
|
|||||||
|
services:
|
||||||
|
fpv-tracker:
|
||||||
|
build:
|
||||||
|
context: .
|
||||||
|
dockerfile: Dockerfile
|
||||||
|
args:
|
||||||
|
PYTORCH_INDEX_URL: https://download.pytorch.org/whl/cu128
|
||||||
|
image: fpv-tracker:cu128-offline
|
||||||
|
gpus: all
|
||||||
|
restart: unless-stopped
|
||||||
|
stdin_open: true
|
||||||
|
tty: true
|
||||||
|
environment:
|
||||||
|
FPV_MODEL_PATH: /app/best.pt
|
||||||
|
FPV_SOURCE: "0"
|
||||||
|
FPV_SHOW_OUTPUT: "0"
|
||||||
|
FPV_SAVE_INFER_VIDEO: "1"
|
||||||
|
FPV_OUT_VIDEO_PATH: /data/out/out_infer.mp4
|
||||||
|
FPV_GUIDANCE_EXPORT_ENABLE: "1"
|
||||||
|
FPV_GUIDANCE_EXPORT_PATH: /data/guidance/guidance_state.json
|
||||||
|
FPV_AUTOPILOT_ENABLE: "1"
|
||||||
|
FPV_AUTOPILOT_BACKEND: "json"
|
||||||
|
FPV_AUTOPILOT_JSON_PATH: /data/autopilot/autopilot_cmd.json
|
||||||
|
FPV_PROTO_UDP_ENABLE: "0"
|
||||||
|
FPV_PROTO_UDP_HOST: "192.168.1.10"
|
||||||
|
FPV_PROTO_UDP_PORT: "5005"
|
||||||
|
volumes:
|
||||||
|
- ./runtime-data:/data
|
||||||
|
ports:
|
||||||
|
- "5600:5600/udp"
|
||||||
|
|
||||||
|
# Для Linux-камеры можно раскомментировать:
|
||||||
|
# devices:
|
||||||
|
# - /dev/video0:/dev/video0
|
||||||
@ -0,0 +1,20 @@
|
|||||||
|
#!/usr/bin/env bash
|
||||||
|
set -euo pipefail
|
||||||
|
|
||||||
|
mkdir -p /data/out /data/guidance /data/autopilot
|
||||||
|
|
||||||
|
if [[ -n "${FPV_OUT_VIDEO_PATH:-}" ]]; then
|
||||||
|
mkdir -p "$(dirname "${FPV_OUT_VIDEO_PATH}")"
|
||||||
|
fi
|
||||||
|
|
||||||
|
if [[ -n "${FPV_GUIDANCE_EXPORT_PATH:-}" ]]; then
|
||||||
|
mkdir -p "$(dirname "${FPV_GUIDANCE_EXPORT_PATH}")"
|
||||||
|
fi
|
||||||
|
|
||||||
|
if [[ -n "${FPV_AUTOPILOT_JSON_PATH:-}" ]]; then
|
||||||
|
mkdir -p "$(dirname "${FPV_AUTOPILOT_JSON_PATH}")"
|
||||||
|
fi
|
||||||
|
|
||||||
|
echo "[entrypoint] source=${FPV_SOURCE:-<config default>} backend=${FPV_AUTOPILOT_BACKEND:-json} proto_udp=${FPV_PROTO_UDP_ENABLE:-0}"
|
||||||
|
|
||||||
|
exec "$@"
|
||||||
@ -0,0 +1,542 @@
|
|||||||
|
# Fast Handoff And Stale Lock Implementation Plan
|
||||||
|
|
||||||
|
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||||
|
|
||||||
|
**Goal:** Make the tracker abandon stale lock anchors quickly during sharp turns, switch to valid new YOLO/ByteTrack candidates in `1-2` frames when safe, and reset guidance smoothing so the blue point and yellow line immediately follow the new target region.
|
||||||
|
|
||||||
|
**Architecture:** Extract the stale-lock and fast-handoff decision logic into a small pure-Python helper module so it can be unit-tested without running the full video pipeline. Then integrate that logic into `main.py`, add event-driven guidance reset hooks in `guidance.py`, and extend `decision_logger.py` so before/after behavior can be measured from logs. The plan assumes the current workspace is not a git repository; commit commands are included for use once the code is executed inside a repo-backed workspace.
|
||||||
|
|
||||||
|
**Tech Stack:** Python 3.12, pytest, NumPy, OpenCV-based tracking pipeline, existing local CSV/JSON diagnostics.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## File Structure
|
||||||
|
|
||||||
|
- Create: `target_handoff.py`
|
||||||
|
- Responsibility: pure decision helpers for stale-lock detection, fast-handoff threshold selection, and lightweight switch event representation.
|
||||||
|
- Create: `tests/test_target_handoff.py`
|
||||||
|
- Responsibility: unit tests for stale-lock break and fast-handoff policy.
|
||||||
|
- Modify: `guidance.py`
|
||||||
|
- Responsibility: add explicit reset hook for smoothing state on stale-break and target switch.
|
||||||
|
- Create: `tests/test_guidance_reset.py`
|
||||||
|
- Responsibility: unit tests that prove aim/command smoothing is reset correctly.
|
||||||
|
- Modify: `main.py`
|
||||||
|
- Responsibility: compute stale-lock inputs, invoke handoff helpers, suppress stale prediction hold, propagate reset events to guidance, and emit diagnostics.
|
||||||
|
- Modify: `decision_logger.py`
|
||||||
|
- Responsibility: add CSV fields and summary accounting for stale-lock/handoff/guidance-reset events.
|
||||||
|
- Create: `tests/test_decision_logger_events.py`
|
||||||
|
- Responsibility: verify new CSV fields and event serialization.
|
||||||
|
- Modify: `config.py`
|
||||||
|
- Responsibility: add narrowly scoped tuning knobs for stale-lock break and fast handoff.
|
||||||
|
|
||||||
|
## Task 1: Add Pure Handoff Decision Module
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Create: `target_handoff.py`
|
||||||
|
- Test: `tests/test_target_handoff.py`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Write the failing tests**
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from target_handoff import (
|
||||||
|
HandoffDecision,
|
||||||
|
compute_fast_handoff_hits,
|
||||||
|
evaluate_stale_lock,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_evaluate_stale_lock_requires_fresh_candidate_and_weak_anchor():
|
||||||
|
decision = evaluate_stale_lock(
|
||||||
|
confirmed=True,
|
||||||
|
have_fresh_candidate=True,
|
||||||
|
target_has_fresh_update=False,
|
||||||
|
candidate_dist_px=420.0,
|
||||||
|
pred_diag_px=70.0,
|
||||||
|
miss_streak=8,
|
||||||
|
klt_valid=False,
|
||||||
|
klt_quality=0.18,
|
||||||
|
klt_iou=0.12,
|
||||||
|
candidate_motion_ok=True,
|
||||||
|
target_motion_ok=False,
|
||||||
|
stale_break_dist_diag=3.5,
|
||||||
|
stale_break_min_miss=2,
|
||||||
|
stale_break_klt_quality=0.35,
|
||||||
|
stale_break_klt_iou=0.25,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decision.stale_lock_active is True
|
||||||
|
assert "target_not_fresh" in decision.reasons
|
||||||
|
assert "weak_klt_quality" in decision.reasons
|
||||||
|
|
||||||
|
|
||||||
|
def test_evaluate_stale_lock_rejects_close_candidate_when_anchor_is_healthy():
|
||||||
|
decision = evaluate_stale_lock(
|
||||||
|
confirmed=True,
|
||||||
|
have_fresh_candidate=True,
|
||||||
|
target_has_fresh_update=True,
|
||||||
|
candidate_dist_px=90.0,
|
||||||
|
pred_diag_px=70.0,
|
||||||
|
miss_streak=0,
|
||||||
|
klt_valid=True,
|
||||||
|
klt_quality=0.92,
|
||||||
|
klt_iou=0.71,
|
||||||
|
candidate_motion_ok=True,
|
||||||
|
target_motion_ok=True,
|
||||||
|
stale_break_dist_diag=3.5,
|
||||||
|
stale_break_min_miss=2,
|
||||||
|
stale_break_klt_quality=0.35,
|
||||||
|
stale_break_klt_iou=0.25,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decision.stale_lock_active is False
|
||||||
|
assert decision.reasons == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_compute_fast_handoff_hits_reduces_confirmation_under_stale_lock():
|
||||||
|
hits = compute_fast_handoff_hits(
|
||||||
|
stale_lock_active=True,
|
||||||
|
motion_switch_mode=False,
|
||||||
|
approach_extra_hits=2,
|
||||||
|
fast_maneuver_extra_hits=2,
|
||||||
|
default_switch_hits=4,
|
||||||
|
fast_handoff_hits=2,
|
||||||
|
motion_switch_hits=6,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert hits == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_handoff_decision_repr_is_stable_for_logging():
|
||||||
|
decision = HandoffDecision(
|
||||||
|
stale_lock_active=True,
|
||||||
|
allow_fast_handoff=True,
|
||||||
|
required_switch_hits=2,
|
||||||
|
disable_prediction_hold=True,
|
||||||
|
reasons=["target_not_fresh", "weak_klt_quality"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decision.reason_text() == "target_not_fresh|weak_klt_quality"
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 2: Run test to verify it fails**
|
||||||
|
|
||||||
|
Run: `pytest tests/test_target_handoff.py -v`
|
||||||
|
|
||||||
|
Expected: FAIL with `ModuleNotFoundError: No module named 'target_handoff'`
|
||||||
|
|
||||||
|
- [ ] **Step 3: Write minimal implementation**
|
||||||
|
|
||||||
|
```python
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class HandoffDecision:
|
||||||
|
stale_lock_active: bool
|
||||||
|
allow_fast_handoff: bool
|
||||||
|
required_switch_hits: int
|
||||||
|
disable_prediction_hold: bool
|
||||||
|
reasons: list[str]
|
||||||
|
|
||||||
|
def reason_text(self) -> str:
|
||||||
|
return "|".join(self.reasons)
|
||||||
|
|
||||||
|
|
||||||
|
def evaluate_stale_lock(
|
||||||
|
*,
|
||||||
|
confirmed,
|
||||||
|
have_fresh_candidate,
|
||||||
|
target_has_fresh_update,
|
||||||
|
candidate_dist_px,
|
||||||
|
pred_diag_px,
|
||||||
|
miss_streak,
|
||||||
|
klt_valid,
|
||||||
|
klt_quality,
|
||||||
|
klt_iou,
|
||||||
|
candidate_motion_ok,
|
||||||
|
target_motion_ok,
|
||||||
|
stale_break_dist_diag,
|
||||||
|
stale_break_min_miss,
|
||||||
|
stale_break_klt_quality,
|
||||||
|
stale_break_klt_iou,
|
||||||
|
):
|
||||||
|
reasons = []
|
||||||
|
if not confirmed or not have_fresh_candidate:
|
||||||
|
return HandoffDecision(False, False, 0, False, [])
|
||||||
|
|
||||||
|
far_enough = candidate_dist_px >= max(40.0, stale_break_dist_diag * max(pred_diag_px, 1.0))
|
||||||
|
weak_anchor = (
|
||||||
|
(not klt_valid)
|
||||||
|
or (klt_quality < stale_break_klt_quality)
|
||||||
|
or (klt_iou < stale_break_klt_iou)
|
||||||
|
or (miss_streak >= stale_break_min_miss)
|
||||||
|
or (candidate_motion_ok and not target_motion_ok)
|
||||||
|
or (not target_has_fresh_update)
|
||||||
|
)
|
||||||
|
|
||||||
|
if not target_has_fresh_update:
|
||||||
|
reasons.append("target_not_fresh")
|
||||||
|
if not klt_valid:
|
||||||
|
reasons.append("klt_invalid")
|
||||||
|
if klt_quality < stale_break_klt_quality:
|
||||||
|
reasons.append("weak_klt_quality")
|
||||||
|
if klt_iou < stale_break_klt_iou:
|
||||||
|
reasons.append("weak_klt_iou")
|
||||||
|
if miss_streak >= stale_break_min_miss:
|
||||||
|
reasons.append("miss_streak")
|
||||||
|
if candidate_motion_ok and not target_motion_ok:
|
||||||
|
reasons.append("motion_conflict")
|
||||||
|
|
||||||
|
stale = bool(far_enough and weak_anchor)
|
||||||
|
return HandoffDecision(
|
||||||
|
stale_lock_active=stale,
|
||||||
|
allow_fast_handoff=stale,
|
||||||
|
required_switch_hits=0,
|
||||||
|
disable_prediction_hold=stale,
|
||||||
|
reasons=reasons if stale else [],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def compute_fast_handoff_hits(
|
||||||
|
*,
|
||||||
|
stale_lock_active,
|
||||||
|
motion_switch_mode,
|
||||||
|
approach_extra_hits,
|
||||||
|
fast_maneuver_extra_hits,
|
||||||
|
default_switch_hits,
|
||||||
|
fast_handoff_hits,
|
||||||
|
motion_switch_hits,
|
||||||
|
):
|
||||||
|
if stale_lock_active:
|
||||||
|
return int(fast_handoff_hits)
|
||||||
|
|
||||||
|
hits = int(default_switch_hits) + int(approach_extra_hits) + int(fast_maneuver_extra_hits)
|
||||||
|
if motion_switch_mode:
|
||||||
|
hits = max(hits, int(motion_switch_hits))
|
||||||
|
return int(hits)
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 4: Run test to verify it passes**
|
||||||
|
|
||||||
|
Run: `pytest tests/test_target_handoff.py -v`
|
||||||
|
|
||||||
|
Expected: `4 passed`
|
||||||
|
|
||||||
|
- [ ] **Step 5: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add target_handoff.py tests/test_target_handoff.py
|
||||||
|
git commit -m "feat: add stale-lock and fast-handoff decision helpers"
|
||||||
|
```
|
||||||
|
|
||||||
|
## Task 2: Add Guidance Reset Hook
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Modify: `guidance.py`
|
||||||
|
- Create: `tests/test_guidance_reset.py`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Write the failing tests**
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from guidance import ScreenGuidanceController
|
||||||
|
|
||||||
|
|
||||||
|
def test_guidance_reset_reinitializes_smoothed_point():
|
||||||
|
ctrl = ScreenGuidanceController()
|
||||||
|
ctrl.smoothed_point = np.array([100.0, 200.0], dtype=np.float32)
|
||||||
|
ctrl.smoothed_cmd = np.array([0.7, -0.4], dtype=np.float32)
|
||||||
|
|
||||||
|
ctrl.reset_for_target_switch(np.array([500.0, 300.0], dtype=np.float32))
|
||||||
|
|
||||||
|
assert np.allclose(ctrl.smoothed_point, np.array([500.0, 300.0], dtype=np.float32))
|
||||||
|
assert np.allclose(ctrl.smoothed_cmd, np.zeros(2, dtype=np.float32))
|
||||||
|
|
||||||
|
|
||||||
|
def test_guidance_reset_can_damp_instead_of_zero_when_requested():
|
||||||
|
ctrl = ScreenGuidanceController()
|
||||||
|
ctrl.smoothed_cmd = np.array([0.8, -0.6], dtype=np.float32)
|
||||||
|
|
||||||
|
ctrl.reset_for_target_switch(np.array([10.0, 20.0], dtype=np.float32), cmd_damp=0.25)
|
||||||
|
|
||||||
|
assert np.allclose(ctrl.smoothed_cmd, np.array([0.2, -0.15], dtype=np.float32))
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 2: Run test to verify it fails**
|
||||||
|
|
||||||
|
Run: `pytest tests/test_guidance_reset.py -v`
|
||||||
|
|
||||||
|
Expected: FAIL with `AttributeError: 'ScreenGuidanceController' object has no attribute 'reset_for_target_switch'`
|
||||||
|
|
||||||
|
- [ ] **Step 3: Write minimal implementation**
|
||||||
|
|
||||||
|
```python
|
||||||
|
def reset_for_target_switch(self, aim_point, cmd_damp=0.0):
|
||||||
|
aim = np.array(aim_point, dtype=np.float32)
|
||||||
|
self.smoothed_point = aim.copy()
|
||||||
|
damp = float(clamp(cmd_damp, 0.0, 1.0))
|
||||||
|
self.smoothed_cmd = (self.smoothed_cmd * damp).astype(np.float32)
|
||||||
|
```
|
||||||
|
|
||||||
|
Add it inside `ScreenGuidanceController` near `status_line()` / `idle_state()` in `guidance.py`.
|
||||||
|
|
||||||
|
- [ ] **Step 4: Run test to verify it passes**
|
||||||
|
|
||||||
|
Run: `pytest tests/test_guidance_reset.py -v`
|
||||||
|
|
||||||
|
Expected: `2 passed`
|
||||||
|
|
||||||
|
- [ ] **Step 5: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add guidance.py tests/test_guidance_reset.py
|
||||||
|
git commit -m "feat: add guidance reset hook for target handoff"
|
||||||
|
```
|
||||||
|
|
||||||
|
## Task 3: Integrate Stale Lock And Fast Handoff Into Main Loop
|
||||||
|
|
||||||
|
**Files:**
|
||||||
|
- Modify: `config.py`
|
||||||
|
- Modify: `main.py`
|
||||||
|
- Modify: `decision_logger.py`
|
||||||
|
- Create: `tests/test_decision_logger_events.py`
|
||||||
|
|
||||||
|
- [ ] **Step 1: Write the failing logger test**
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from decision_logger import FIELDNAMES
|
||||||
|
|
||||||
|
|
||||||
|
def test_decision_logger_contains_handoff_fields():
|
||||||
|
assert "stale_lock_active" in FIELDNAMES
|
||||||
|
assert "stale_lock_reason" in FIELDNAMES
|
||||||
|
assert "fast_handoff_active" in FIELDNAMES
|
||||||
|
assert "guidance_reset_event" in FIELDNAMES
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 2: Run test to verify it fails**
|
||||||
|
|
||||||
|
Run: `pytest tests/test_decision_logger_events.py -v`
|
||||||
|
|
||||||
|
Expected: FAIL because the new field names are missing from `FIELDNAMES`
|
||||||
|
|
||||||
|
- [ ] **Step 3: Add config defaults**
|
||||||
|
|
||||||
|
Add these exact settings near the existing single-target lock policy section in `config.py`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
STALE_LOCK_BREAK_ENABLE = True
|
||||||
|
STALE_LOCK_BREAK_DIST_DIAG = 3.5
|
||||||
|
STALE_LOCK_BREAK_MIN_MISS = 2
|
||||||
|
STALE_LOCK_BREAK_KLT_QUALITY = 0.35
|
||||||
|
STALE_LOCK_BREAK_KLT_IOU = 0.25
|
||||||
|
FAST_HANDOFF_ENABLE = True
|
||||||
|
FAST_HANDOFF_CONFIRM_HITS = 2
|
||||||
|
FAST_HANDOFF_GUIDANCE_CMD_DAMP = 0.0
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 4: Wire stale-lock and fast-handoff in `main.py`**
|
||||||
|
|
||||||
|
Add imports near the top of `main.py`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from target_handoff import compute_fast_handoff_hits, evaluate_stale_lock
|
||||||
|
```
|
||||||
|
|
||||||
|
Add per-frame state defaults near the other per-frame flags:
|
||||||
|
|
||||||
|
```python
|
||||||
|
stale_lock_active = False
|
||||||
|
stale_lock_reason = ""
|
||||||
|
fast_handoff_active = False
|
||||||
|
guidance_reset_event = ""
|
||||||
|
```
|
||||||
|
|
||||||
|
After `best` candidate selection is known and before switch confirmation is finalized, add stale-lock evaluation using current values:
|
||||||
|
|
||||||
|
```python
|
||||||
|
handoff_decision = evaluate_stale_lock(
|
||||||
|
confirmed=confirmed,
|
||||||
|
have_fresh_candidate=(best is not None and best.time_since_update == 0),
|
||||||
|
target_has_fresh_update=(target_track is not None and target_track.time_since_update == 0),
|
||||||
|
candidate_dist_px=float(dist) if pred_ref is not None else 0.0,
|
||||||
|
pred_diag_px=float(pred_diag),
|
||||||
|
miss_streak=int(miss_streak),
|
||||||
|
klt_valid=bool(klt_valid),
|
||||||
|
klt_quality=float(klt.quality),
|
||||||
|
klt_iou=float(iou(klt.box, locked_box_eff)) if (klt.box is not None and locked_box_eff is not None) else 0.0,
|
||||||
|
candidate_motion_ok=bool(residual_ok),
|
||||||
|
target_motion_ok=bool(target_track is not None and target_track.time_since_update == 0),
|
||||||
|
stale_break_dist_diag=float(STALE_LOCK_BREAK_DIST_DIAG),
|
||||||
|
stale_break_min_miss=int(STALE_LOCK_BREAK_MIN_MISS),
|
||||||
|
stale_break_klt_quality=float(STALE_LOCK_BREAK_KLT_QUALITY),
|
||||||
|
stale_break_klt_iou=float(STALE_LOCK_BREAK_KLT_IOU),
|
||||||
|
)
|
||||||
|
stale_lock_active = bool(STALE_LOCK_BREAK_ENABLE and handoff_decision.stale_lock_active)
|
||||||
|
stale_lock_reason = handoff_decision.reason_text()
|
||||||
|
```
|
||||||
|
|
||||||
|
When computing switch confirmation hits, replace the current manual calculation with:
|
||||||
|
|
||||||
|
```python
|
||||||
|
req_switch_hits = compute_fast_handoff_hits(
|
||||||
|
stale_lock_active=bool(FAST_HANDOFF_ENABLE and stale_lock_active),
|
||||||
|
motion_switch_mode=bool(motion_switch_mode),
|
||||||
|
approach_extra_hits=int(approach_switch_extra_hits),
|
||||||
|
fast_maneuver_extra_hits=int(max(0, SWITCH_FAST_MANEUVER_EXTRA_HITS)) if fast_maneuver_guard else 0,
|
||||||
|
default_switch_hits=int(TARGET_SWITCH_CONFIRM_HITS),
|
||||||
|
fast_handoff_hits=int(FAST_HANDOFF_CONFIRM_HITS),
|
||||||
|
motion_switch_hits=int(MOTION_CONF_SWITCH_CONFIRM_HITS),
|
||||||
|
)
|
||||||
|
fast_handoff_active = bool(FAST_HANDOFF_ENABLE and stale_lock_active and req_switch_hits == int(FAST_HANDOFF_CONFIRM_HITS))
|
||||||
|
```
|
||||||
|
|
||||||
|
Disable prediction hold while stale-lock is active by modifying the existing hold branch:
|
||||||
|
|
||||||
|
```python
|
||||||
|
if confirmed and pred_box_eff is not None and (not stale_lock_active) and miss_streak < KALMAN_HOLD_MAX:
|
||||||
|
```
|
||||||
|
|
||||||
|
When a fast handoff or target switch is accepted, reset guidance:
|
||||||
|
|
||||||
|
```python
|
||||||
|
old_target_id = target_id
|
||||||
|
```
|
||||||
|
|
||||||
|
before replacing `target_id`, then after `locked_box_eff = chosen`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
if confirmed and target_track is not None and old_target_id is not None and target_track.track_id != old_target_id:
|
||||||
|
guidance_ctrl.reset_for_target_switch(
|
||||||
|
box_center(locked_box_eff),
|
||||||
|
cmd_damp=float(FAST_HANDOFF_GUIDANCE_CMD_DAMP),
|
||||||
|
)
|
||||||
|
guidance_reset_event = "target_switch"
|
||||||
|
elif stale_lock_active and locked_box_eff is not None:
|
||||||
|
guidance_ctrl.reset_for_target_switch(
|
||||||
|
box_center(locked_box_eff),
|
||||||
|
cmd_damp=float(FAST_HANDOFF_GUIDANCE_CMD_DAMP),
|
||||||
|
)
|
||||||
|
guidance_reset_event = "stale_lock_break"
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 5: Add logger fields**
|
||||||
|
|
||||||
|
Extend `decision_logger.py` field names and row payload with:
|
||||||
|
|
||||||
|
```python
|
||||||
|
"stale_lock_active",
|
||||||
|
"stale_lock_reason",
|
||||||
|
"fast_handoff_active",
|
||||||
|
"guidance_reset_event",
|
||||||
|
```
|
||||||
|
|
||||||
|
and map them in `log_frame()`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
"stale_lock_active": int(bool(kwargs.get("stale_lock_active", False))),
|
||||||
|
"stale_lock_reason": kwargs.get("stale_lock_reason", ""),
|
||||||
|
"fast_handoff_active": int(bool(kwargs.get("fast_handoff_active", False))),
|
||||||
|
"guidance_reset_event": kwargs.get("guidance_reset_event", ""),
|
||||||
|
```
|
||||||
|
|
||||||
|
Also append event names if present:
|
||||||
|
|
||||||
|
```python
|
||||||
|
if kwargs.get("stale_lock_active", False):
|
||||||
|
events.append("stale_lock_break")
|
||||||
|
if kwargs.get("fast_handoff_active", False):
|
||||||
|
events.append("fast_handoff_switch")
|
||||||
|
if kwargs.get("guidance_reset_event", ""):
|
||||||
|
events.append("guidance_reset")
|
||||||
|
```
|
||||||
|
|
||||||
|
- [ ] **Step 6: Run focused tests**
|
||||||
|
|
||||||
|
Run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pytest tests/test_target_handoff.py tests/test_guidance_reset.py tests/test_decision_logger_events.py tests/test_track_score_gates.py -v
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected: all tests PASS
|
||||||
|
|
||||||
|
- [ ] **Step 7: Run a smoke video pass**
|
||||||
|
|
||||||
|
Run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
py -3 main.py
|
||||||
|
```
|
||||||
|
|
||||||
|
Expected:
|
||||||
|
- no import/runtime errors
|
||||||
|
- new CSV fields appear in `track_log_*.csv`
|
||||||
|
- problem segment near `231.5s - 232.3s` shows earlier stale-lock break and faster visual realignment
|
||||||
|
|
||||||
|
- [ ] **Step 8: Commit**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git add config.py main.py guidance.py decision_logger.py tests/test_decision_logger_events.py
|
||||||
|
git commit -m "feat: add stale-lock break and fast handoff integration"
|
||||||
|
```
|
||||||
|
|
||||||
|
## Self-Review
|
||||||
|
|
||||||
|
### Spec coverage
|
||||||
|
|
||||||
|
- Stale-lock detection: Task 1 + Task 3
|
||||||
|
- Fast handoff switch policy: Task 1 + Task 3
|
||||||
|
- Prediction hold restrictions: Task 3
|
||||||
|
- Guidance reset on switch: Task 2 + Task 3
|
||||||
|
- Logging and diagnostics: Task 3
|
||||||
|
- Verification on known failure segment: Task 3 smoke pass
|
||||||
|
|
||||||
|
No spec gaps remain for the first-pass implementation.
|
||||||
|
|
||||||
|
### Placeholder scan
|
||||||
|
|
||||||
|
- No `TBD`, `TODO`, or deferred placeholders included.
|
||||||
|
- Every code-changing step contains explicit code snippets.
|
||||||
|
- Every verification step contains an exact command and expected result.
|
||||||
|
|
||||||
|
### Type consistency
|
||||||
|
|
||||||
|
- `HandoffDecision` property names are consistent across tests and integration steps.
|
||||||
|
- `reset_for_target_switch()` is referenced consistently in tests and integration.
|
||||||
|
- New logger field names match the values referenced in `main.py`.
|
||||||
|
|
||||||
|
## Execution Handoff
|
||||||
|
|
||||||
|
Plan complete and saved to `docs/superpowers/plans/2026-06-24-fast-handoff-stale-lock-implementation.md`. Two execution options:
|
||||||
|
|
||||||
|
**1. Subagent-Driven (recommended)** - I dispatch a fresh subagent per task, review between tasks, fast iteration
|
||||||
|
|
||||||
|
**2. Inline Execution** - Execute tasks in this session using executing-plans, batch execution with checkpoints
|
||||||
|
|
||||||
|
**Which approach?**
|
||||||
@ -0,0 +1,218 @@
|
|||||||
|
# Fast Handoff And Stale Lock Design
|
||||||
|
|
||||||
|
## Goal
|
||||||
|
|
||||||
|
Reduce the delay between a valid new YOLO detection and the moment when:
|
||||||
|
- `target_id` switches to the new target track,
|
||||||
|
- `locked_box_eff` stops following stale KLT/Kalman state,
|
||||||
|
- the blue point and yellow guidance line stop pointing at the old scene region.
|
||||||
|
|
||||||
|
The fix must improve turning behavior without materially increasing false switches to terrain, OSD, or random high-contrast background features.
|
||||||
|
|
||||||
|
## Current Problem
|
||||||
|
|
||||||
|
The current pipeline has three separate lag sources that compound during sharp turns:
|
||||||
|
|
||||||
|
1. `main.py` keeps the old target anchor alive through prediction hold and KLT, even when YOLO has already found a better candidate elsewhere.
|
||||||
|
2. The target switch policy requires multiple confirming frames before accepting the new candidate, especially in chase/close regimes.
|
||||||
|
3. `guidance.py` smooths both aim point and command state, so even after a switch the yellow line continues to drift toward the new target instead of snapping quickly.
|
||||||
|
|
||||||
|
Observed failure mode:
|
||||||
|
- YOLO draws fresh boxes around the real target in a new image region.
|
||||||
|
- `target_id`, `locked_box_eff`, and `pred_box_eff` still describe the stale region.
|
||||||
|
- KLT may still look "valid" because it is tracking texture on the wrong patch.
|
||||||
|
- Guidance uses the stale lock, so the blue point and yellow line remain pointed away from the true target.
|
||||||
|
|
||||||
|
## Recommended Approach
|
||||||
|
|
||||||
|
Implement a three-part fix:
|
||||||
|
|
||||||
|
1. Add stale-lock detection in `main.py`.
|
||||||
|
2. Add a fast-handoff switch path for valid replacement tracks.
|
||||||
|
3. Reset guidance smoothing state on stale-break and on confirmed target switch.
|
||||||
|
|
||||||
|
This is intentionally an incremental fix that preserves the existing architecture:
|
||||||
|
- YOLO
|
||||||
|
- ByteTrack
|
||||||
|
- SafeKalman
|
||||||
|
- Hybrid KLT tracker
|
||||||
|
- guidance overlay/export
|
||||||
|
|
||||||
|
## Design
|
||||||
|
|
||||||
|
### 1. Stale-Lock Detection
|
||||||
|
|
||||||
|
Add a new state evaluation step in `main.py` after fresh YOLO tracks are available and before final target selection/hold behavior is applied.
|
||||||
|
|
||||||
|
The stale-lock detector should mark the current lock as stale when all of the following are true:
|
||||||
|
|
||||||
|
- The system is in `confirmed` state.
|
||||||
|
- There is at least one fresh YOLO candidate track.
|
||||||
|
- The current `target_id` either has no fresh update or is significantly less plausible than another candidate.
|
||||||
|
- The best candidate is far from the current `pred_ref` / `locked_box_eff`.
|
||||||
|
- The current anchor is weak by at least one anchor-quality signal.
|
||||||
|
|
||||||
|
Anchor-quality signals:
|
||||||
|
- `klt_valid` is false.
|
||||||
|
- `klt_quality` is below a configurable stale threshold.
|
||||||
|
- `klt.box` exists but has poor IoU with `locked_box_eff`.
|
||||||
|
- `miss_streak` is already elevated.
|
||||||
|
- A fresh candidate is motion-consistent while the current target is not.
|
||||||
|
|
||||||
|
Distance signal:
|
||||||
|
- Compute distance from candidate center to `pred_ref` center.
|
||||||
|
- Compare against a dedicated stale-break distance threshold based on target diagonal, not just the existing reacquisition threshold.
|
||||||
|
|
||||||
|
Important behavior:
|
||||||
|
- Marking the old lock as stale does not immediately force a blind jump to a new candidate.
|
||||||
|
- It only disables continued trust in the stale anchor for guidance and prediction-hold purposes.
|
||||||
|
|
||||||
|
### 2. Fast Handoff Switch Policy
|
||||||
|
|
||||||
|
Keep the current conservative switch logic as the default path.
|
||||||
|
|
||||||
|
Add a second switch path named fast handoff, enabled only when the stale-lock detector fires.
|
||||||
|
|
||||||
|
Fast handoff should accept a replacement candidate after only `1-2` confirming frames if the candidate passes strict gating:
|
||||||
|
|
||||||
|
- fresh ByteTrack update (`time_since_update == 0`)
|
||||||
|
- score above switch floor
|
||||||
|
- residual/motion gate passes when available
|
||||||
|
- not inside blocked OSD zones unless near the expected path
|
||||||
|
- distance is large enough to prove it is not the stale lock, but not absurdly far
|
||||||
|
- optional trajectory agreement if reliable target velocity exists
|
||||||
|
|
||||||
|
Fast handoff should bypass or reduce:
|
||||||
|
- `TARGET_SWITCH_CONFIRM_HITS`
|
||||||
|
- `APPROACH_*_SWITCH_EXTRA_HITS`
|
||||||
|
- `SWITCH_FAST_MANEUVER_EXTRA_HITS`
|
||||||
|
|
||||||
|
Fast handoff should also reduce or bypass old-lock waiting:
|
||||||
|
- `TARGET_SWITCH_MISS_FRAMES`
|
||||||
|
- `APPROACH_*_SWITCH_EXTRA_MISS`
|
||||||
|
- `SWITCH_FAST_MANEUVER_EXTRA_MISS`
|
||||||
|
|
||||||
|
If stale-lock is active and no candidate survives gates:
|
||||||
|
- do not keep driving guidance from old stale lock indefinitely,
|
||||||
|
- prefer degrading to neutral/low-confidence hold rather than continuing to steer at the wrong patch.
|
||||||
|
|
||||||
|
### 3. Prediction Hold Restrictions
|
||||||
|
|
||||||
|
The current `pred_box_eff`/Kalman hold is useful when YOLO drops out briefly, but harmful once a stale-lock condition is detected.
|
||||||
|
|
||||||
|
Modify hold behavior so that:
|
||||||
|
- normal prediction hold remains available during benign short misses,
|
||||||
|
- stale-lock condition disables or sharply shortens prediction hold,
|
||||||
|
- KLT-only continuation is not allowed to dominate when fresh YOLO is contradicting it.
|
||||||
|
|
||||||
|
This change should affect the branch that currently sets:
|
||||||
|
- `locked_box_eff = pred_box_eff`
|
||||||
|
- increments `miss_streak`
|
||||||
|
- keeps `confirmed = True`
|
||||||
|
|
||||||
|
Expected result:
|
||||||
|
- the system stops visually "dragging" the old lock across the frame for up to `KALMAN_HOLD_MAX` frames while the real target is already detected somewhere else.
|
||||||
|
|
||||||
|
### 4. Guidance Reset On Switch
|
||||||
|
|
||||||
|
Update `guidance.py` so that guidance smoothing state is reset when:
|
||||||
|
- a stale-lock break occurs, or
|
||||||
|
- `target_id` changes to a new confirmed target.
|
||||||
|
|
||||||
|
Required behavior:
|
||||||
|
- `smoothed_point` should be reinitialized to the new aim point.
|
||||||
|
- `smoothed_cmd` should be zeroed or strongly damped.
|
||||||
|
- command continuity should favor fast visual alignment over preserving momentum from the old target.
|
||||||
|
|
||||||
|
This reset must be explicit and event-driven.
|
||||||
|
|
||||||
|
Without it, even a correct target switch still leaves the yellow line lagging for multiple frames because:
|
||||||
|
- aim point is EMA-smoothed,
|
||||||
|
- command is EMA-smoothed.
|
||||||
|
|
||||||
|
### 5. Logging And Diagnostics
|
||||||
|
|
||||||
|
Extend diagnostics so the fix can be evaluated from logs rather than by eye only.
|
||||||
|
|
||||||
|
Add fields/events such as:
|
||||||
|
- `stale_lock_active`
|
||||||
|
- `stale_lock_reason`
|
||||||
|
- `fast_handoff_candidate_id`
|
||||||
|
- `fast_handoff_active`
|
||||||
|
- `guidance_reset_event`
|
||||||
|
|
||||||
|
Add event tags such as:
|
||||||
|
- `stale_lock_break`
|
||||||
|
- `fast_handoff_switch`
|
||||||
|
- `guidance_reset`
|
||||||
|
|
||||||
|
This allows measuring:
|
||||||
|
- frames between first valid new candidate and switch
|
||||||
|
- frames where YOLO sees the new target while lock still follows the stale region
|
||||||
|
- false switch counts before/after the fix
|
||||||
|
|
||||||
|
## Implementation Boundaries
|
||||||
|
|
||||||
|
Files expected to change:
|
||||||
|
- `main.py`
|
||||||
|
- `guidance.py`
|
||||||
|
- `decision_logger.py`
|
||||||
|
- possibly `config.py` for new thresholds
|
||||||
|
|
||||||
|
Files not in scope for the first pass:
|
||||||
|
- `range_estimation.py`
|
||||||
|
- `proportional_navigation.py`
|
||||||
|
- `autopilot_bridge.py`
|
||||||
|
- major ByteTrack internals rewrite
|
||||||
|
|
||||||
|
## Suggested New Config Controls
|
||||||
|
|
||||||
|
Add narrowly scoped tuning parameters, for example:
|
||||||
|
|
||||||
|
- stale-lock enable flag
|
||||||
|
- stale-lock distance threshold by target diagonal
|
||||||
|
- stale-lock KLT quality floor
|
||||||
|
- stale-lock max miss before break
|
||||||
|
- fast handoff enable flag
|
||||||
|
- fast handoff confirm hits
|
||||||
|
- fast handoff max switch distance
|
||||||
|
- guidance reset damping factor
|
||||||
|
|
||||||
|
The first implementation should keep defaults conservative but clearly more responsive than the current chase preset.
|
||||||
|
|
||||||
|
## Success Criteria
|
||||||
|
|
||||||
|
The fix is successful if:
|
||||||
|
|
||||||
|
- During sharp turns, once YOLO consistently detects the target in a new region, the stale lock is abandoned within `1-3` frames.
|
||||||
|
- `target_id` switches to the correct candidate substantially earlier than today.
|
||||||
|
- The blue point stops following the stale region almost immediately after stale-break.
|
||||||
|
- The yellow line reorients quickly after switch instead of easing over many frames.
|
||||||
|
- False switches to terrain/OSD do not materially increase.
|
||||||
|
|
||||||
|
## Verification Plan
|
||||||
|
|
||||||
|
Use the most recent problem video and compare before/after on the known failure segment around `231.5s - 232.3s`.
|
||||||
|
|
||||||
|
Measure:
|
||||||
|
- first frame where new YOLO candidate is valid
|
||||||
|
- frame where stale lock is broken
|
||||||
|
- frame where `target_switch` occurs
|
||||||
|
- frame where guidance visually realigns
|
||||||
|
|
||||||
|
Also review aggregate log metrics:
|
||||||
|
- target switch count
|
||||||
|
- lock lost count
|
||||||
|
- time spent in stale-lock-active state
|
||||||
|
- rate of false handoffs
|
||||||
|
|
||||||
|
## Risks
|
||||||
|
|
||||||
|
- If stale-lock conditions are too loose, the tracker may become jumpy.
|
||||||
|
- If fast handoff ignores too many gates, terrain switches may increase.
|
||||||
|
- If guidance reset is too aggressive, command output may become visually jerky.
|
||||||
|
|
||||||
|
The design therefore uses:
|
||||||
|
- explicit stale-lock conditions,
|
||||||
|
- gated fast handoff,
|
||||||
|
- targeted guidance reset only on real switch/break events.
|
||||||
@ -0,0 +1,316 @@
|
|||||||
|
import json
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from config import *
|
||||||
|
from helpers import box_center, clamp, clip_box
|
||||||
|
|
||||||
|
|
||||||
|
def _curve_command(value, deadzone, expo, gain):
|
||||||
|
v = float(np.clip(value, -1.0, 1.0))
|
||||||
|
av = abs(v)
|
||||||
|
if av <= deadzone:
|
||||||
|
return 0.0
|
||||||
|
t = (av - deadzone) / max(1e-6, 1.0 - deadzone)
|
||||||
|
out = min(1.0, float(gain) * (t ** float(expo)))
|
||||||
|
return out if v >= 0.0 else -out
|
||||||
|
|
||||||
|
|
||||||
|
class ScreenGuidanceController:
|
||||||
|
def __init__(self):
|
||||||
|
self.enabled = bool(GUIDANCE_ENABLE)
|
||||||
|
self.export_enable = bool(self.enabled and GUIDANCE_EXPORT_ENABLE)
|
||||||
|
self.export_every = int(max(1, GUIDANCE_EXPORT_EVERY))
|
||||||
|
self.export_path = Path(GUIDANCE_EXPORT_PATH)
|
||||||
|
if self.export_enable:
|
||||||
|
self.export_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
self.cmd_keep = float(clamp(GUIDANCE_CMD_SMOOTH, 0.0, 0.97))
|
||||||
|
self.point_keep = float(clamp(GUIDANCE_POINT_SMOOTH, 0.0, 0.97))
|
||||||
|
self.last_export_frame = -10**9
|
||||||
|
self.export_retry_at = 0.0
|
||||||
|
self.last_export_warn_ts = 0.0
|
||||||
|
self.smoothed_cmd = np.zeros(2, dtype=np.float32)
|
||||||
|
self.smoothed_point = None
|
||||||
|
|
||||||
|
def status_line(self):
|
||||||
|
if not self.enabled:
|
||||||
|
return "Guidance disabled"
|
||||||
|
if self.export_enable:
|
||||||
|
return f"Guidance ready: export={self.export_path}"
|
||||||
|
return "Guidance ready"
|
||||||
|
|
||||||
|
def reset_for_target_switch(self, aim_point, cmd_damp=0.0):
|
||||||
|
aim = np.array(aim_point, dtype=np.float32)
|
||||||
|
self.smoothed_point = aim.copy()
|
||||||
|
damp = float(clamp(cmd_damp, 0.0, 1.0))
|
||||||
|
self.smoothed_cmd = (self.smoothed_cmd * damp).astype(np.float32)
|
||||||
|
|
||||||
|
def idle_state(self):
|
||||||
|
return {
|
||||||
|
"active": False,
|
||||||
|
"status": "SEARCH",
|
||||||
|
"frame_id": -1,
|
||||||
|
"target_id": None,
|
||||||
|
"frame_w": None,
|
||||||
|
"frame_h": None,
|
||||||
|
"aim_x": None,
|
||||||
|
"aim_y": None,
|
||||||
|
"box_w": None,
|
||||||
|
"box_h": None,
|
||||||
|
"error_x": 0.0,
|
||||||
|
"error_y": 0.0,
|
||||||
|
"cmd_x": 0.0,
|
||||||
|
"cmd_y": 0.0,
|
||||||
|
"steer_x": 0.0,
|
||||||
|
"steer_y": 0.0,
|
||||||
|
"look_dx": 0.0,
|
||||||
|
"look_dy": 0.0,
|
||||||
|
"confidence": 0.0,
|
||||||
|
"on_target": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
def update(
|
||||||
|
self,
|
||||||
|
frame_id,
|
||||||
|
frame_w,
|
||||||
|
frame_h,
|
||||||
|
confirmed,
|
||||||
|
locked_box,
|
||||||
|
pred_box,
|
||||||
|
override_box,
|
||||||
|
target_id,
|
||||||
|
target_track_score,
|
||||||
|
klt_valid,
|
||||||
|
klt_quality,
|
||||||
|
miss_streak,
|
||||||
|
vx,
|
||||||
|
vy,
|
||||||
|
):
|
||||||
|
state = self.idle_state()
|
||||||
|
state["frame_id"] = int(frame_id)
|
||||||
|
state["target_id"] = None if target_id is None else int(target_id)
|
||||||
|
state["frame_w"] = int(frame_w)
|
||||||
|
state["frame_h"] = int(frame_h)
|
||||||
|
|
||||||
|
if not self.enabled:
|
||||||
|
return state
|
||||||
|
|
||||||
|
guide_box = None
|
||||||
|
if confirmed and override_box is not None:
|
||||||
|
guide_box = np.array(override_box, dtype=np.float32)
|
||||||
|
elif confirmed and locked_box is not None:
|
||||||
|
guide_box = np.array(locked_box, dtype=np.float32)
|
||||||
|
elif confirmed and pred_box is not None:
|
||||||
|
guide_box = np.array(pred_box, dtype=np.float32)
|
||||||
|
|
||||||
|
if guide_box is None:
|
||||||
|
self.smoothed_cmd *= 0.60
|
||||||
|
state["cmd_x"] = float(self.smoothed_cmd[0])
|
||||||
|
state["cmd_y"] = float(self.smoothed_cmd[1])
|
||||||
|
state["steer_x"] = float(self.smoothed_cmd[0])
|
||||||
|
state["steer_y"] = float(self.smoothed_cmd[1])
|
||||||
|
state["look_dx"] = float(self.smoothed_cmd[0] * float(GUIDANCE_LOOK_MAX_STEP_PX))
|
||||||
|
state["look_dy"] = float(self.smoothed_cmd[1] * float(GUIDANCE_LOOK_MAX_STEP_PX))
|
||||||
|
self._maybe_export(state)
|
||||||
|
return state
|
||||||
|
|
||||||
|
guide_box = clip_box(guide_box, frame_w, frame_h)
|
||||||
|
center = box_center(guide_box).astype(np.float32)
|
||||||
|
box_w = float(max(0.0, guide_box[2] - guide_box[0]))
|
||||||
|
box_h = float(max(0.0, guide_box[3] - guide_box[1]))
|
||||||
|
|
||||||
|
lead = np.array([float(vx), float(vy)], dtype=np.float32) * float(GUIDANCE_LEAD_SEC)
|
||||||
|
lead_norm = float(np.linalg.norm(lead))
|
||||||
|
if lead_norm > float(GUIDANCE_MAX_LEAD_PX):
|
||||||
|
lead *= float(GUIDANCE_MAX_LEAD_PX) / max(1e-6, lead_norm)
|
||||||
|
|
||||||
|
aim = center + lead
|
||||||
|
aim[1] += float(GUIDANCE_BOX_BIAS_Y) * max(2.0, float(guide_box[3] - guide_box[1]))
|
||||||
|
aim[0] = float(clamp(aim[0], 0.0, float(frame_w - 1)))
|
||||||
|
aim[1] = float(clamp(aim[1], 0.0, float(frame_h - 1)))
|
||||||
|
|
||||||
|
if self.smoothed_point is None:
|
||||||
|
self.smoothed_point = aim.copy()
|
||||||
|
else:
|
||||||
|
self.smoothed_point = (
|
||||||
|
(self.point_keep * self.smoothed_point)
|
||||||
|
+ ((1.0 - self.point_keep) * aim)
|
||||||
|
).astype(np.float32)
|
||||||
|
|
||||||
|
screen_center = np.array([0.5 * float(frame_w), 0.5 * float(frame_h)], dtype=np.float32)
|
||||||
|
error_px = self.smoothed_point - screen_center
|
||||||
|
error_norm = np.array(
|
||||||
|
[
|
||||||
|
float(error_px[0] / max(1.0, 0.5 * float(frame_w))),
|
||||||
|
float(error_px[1] / max(1.0, 0.5 * float(frame_h))),
|
||||||
|
],
|
||||||
|
dtype=np.float32,
|
||||||
|
)
|
||||||
|
error_norm = np.clip(error_norm, -1.0, 1.0)
|
||||||
|
|
||||||
|
raw_cmd = np.array(
|
||||||
|
[
|
||||||
|
_curve_command(error_norm[0], GUIDANCE_DEADZONE_NORM, GUIDANCE_EXPO, GUIDANCE_GAIN_X),
|
||||||
|
_curve_command(error_norm[1], GUIDANCE_DEADZONE_NORM, GUIDANCE_EXPO, GUIDANCE_GAIN_Y),
|
||||||
|
],
|
||||||
|
dtype=np.float32,
|
||||||
|
)
|
||||||
|
self.smoothed_cmd = (
|
||||||
|
(self.cmd_keep * self.smoothed_cmd)
|
||||||
|
+ ((1.0 - self.cmd_keep) * raw_cmd)
|
||||||
|
).astype(np.float32)
|
||||||
|
|
||||||
|
confidence = 0.0
|
||||||
|
confidence += 0.40 if confirmed else 0.0
|
||||||
|
confidence += 0.15 if locked_box is not None else 0.05
|
||||||
|
|
||||||
|
# KLT теперь главный источник доверия (klt_quality ≈ 0.93)
|
||||||
|
if klt_valid:
|
||||||
|
confidence += 0.30 * float(clamp(klt_quality, 0.0, 1.0))
|
||||||
|
else:
|
||||||
|
confidence += 0.05 * float(clamp(klt_quality, 0.0, 1.0))
|
||||||
|
|
||||||
|
# ByteTrack score — второстепенный
|
||||||
|
if target_track_score is not None:
|
||||||
|
confidence += 0.10 * float(clamp(target_track_score, 0.0, 1.0))
|
||||||
|
|
||||||
|
# Мягче штраф за miss (YOLO слабый, это временно)
|
||||||
|
confidence -= 0.025 * float(clamp(miss_streak, 0, 10))
|
||||||
|
|
||||||
|
confidence = float(clamp(confidence, 0.0, 1.0))
|
||||||
|
|
||||||
|
err_mag = float(np.linalg.norm(error_norm))
|
||||||
|
on_target = bool(
|
||||||
|
confirmed
|
||||||
|
and (confidence >= float(GUIDANCE_CONF_MIN))
|
||||||
|
and (err_mag <= float(GUIDANCE_ON_TARGET_NORM))
|
||||||
|
)
|
||||||
|
|
||||||
|
status = "SEARCH"
|
||||||
|
if confirmed and override_box is not None and miss_streak > 0:
|
||||||
|
status = "REACQ"
|
||||||
|
elif confirmed and miss_streak > 0:
|
||||||
|
status = "HOLD"
|
||||||
|
elif confirmed:
|
||||||
|
status = "LOCK"
|
||||||
|
|
||||||
|
state.update(
|
||||||
|
{
|
||||||
|
"active": bool(confirmed),
|
||||||
|
"status": status,
|
||||||
|
"aim_x": float(self.smoothed_point[0]),
|
||||||
|
"aim_y": float(self.smoothed_point[1]),
|
||||||
|
"box_w": box_w,
|
||||||
|
"box_h": box_h,
|
||||||
|
"error_x": float(error_norm[0]),
|
||||||
|
"error_y": float(error_norm[1]),
|
||||||
|
"cmd_x": float(self.smoothed_cmd[0]),
|
||||||
|
"cmd_y": float(self.smoothed_cmd[1]),
|
||||||
|
"steer_x": float(self.smoothed_cmd[0]),
|
||||||
|
"steer_y": float(self.smoothed_cmd[1]),
|
||||||
|
"look_dx": float(self.smoothed_cmd[0] * float(GUIDANCE_LOOK_MAX_STEP_PX)),
|
||||||
|
"look_dy": float(self.smoothed_cmd[1] * float(GUIDANCE_LOOK_MAX_STEP_PX)),
|
||||||
|
"confidence": confidence,
|
||||||
|
"on_target": on_target,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
self._maybe_export(state)
|
||||||
|
return state
|
||||||
|
|
||||||
|
def draw_overlay(self, frame_bgr, state, sx, sy):
|
||||||
|
if (not self.enabled) or (not DRAW_GUIDANCE):
|
||||||
|
return
|
||||||
|
|
||||||
|
frame_h, frame_w = frame_bgr.shape[:2]
|
||||||
|
cx = int(0.5 * frame_w)
|
||||||
|
cy = int(0.5 * frame_h)
|
||||||
|
|
||||||
|
color = (0, 220, 255) if state.get("active", False) else (120, 120, 120)
|
||||||
|
cv2.drawMarker(frame_bgr, (cx, cy), color, markerType=cv2.MARKER_CROSS, markerSize=18, thickness=1)
|
||||||
|
|
||||||
|
aim_x = state.get("aim_x", None)
|
||||||
|
aim_y = state.get("aim_y", None)
|
||||||
|
if aim_x is not None and aim_y is not None:
|
||||||
|
ax = int(float(aim_x) / max(1e-6, float(sx)))
|
||||||
|
ay = int(float(aim_y) / max(1e-6, float(sy)))
|
||||||
|
cv2.circle(frame_bgr, (ax, ay), 8, color, 2)
|
||||||
|
cv2.line(frame_bgr, (cx, cy), (ax, ay), color, 1, cv2.LINE_AA)
|
||||||
|
cv2.putText(
|
||||||
|
frame_bgr,
|
||||||
|
f"G {state['status']} steer=({state['steer_x']:+.2f},{state['steer_y']:+.2f}) conf={state['confidence']:.2f}",
|
||||||
|
(20, 165),
|
||||||
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||||||
|
0.60,
|
||||||
|
color,
|
||||||
|
2,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _maybe_export(self, state):
|
||||||
|
frame_id = int(state["frame_id"])
|
||||||
|
now = time.monotonic()
|
||||||
|
if (not self.export_enable) or ((frame_id - self.last_export_frame) < self.export_every) or (now < self.export_retry_at):
|
||||||
|
return
|
||||||
|
|
||||||
|
payload = {
|
||||||
|
"frame_id": frame_id,
|
||||||
|
"active": bool(state["active"]),
|
||||||
|
"status": str(state["status"]),
|
||||||
|
"target_id": state["target_id"],
|
||||||
|
"frame_w": state["frame_w"],
|
||||||
|
"frame_h": state["frame_h"],
|
||||||
|
"aim_x": state["aim_x"],
|
||||||
|
"aim_y": state["aim_y"],
|
||||||
|
"box_w": state["box_w"],
|
||||||
|
"box_h": state["box_h"],
|
||||||
|
"error_x": float(state["error_x"]),
|
||||||
|
"error_y": float(state["error_y"]),
|
||||||
|
"cmd_x": float(state["cmd_x"]),
|
||||||
|
"cmd_y": float(state["cmd_y"]),
|
||||||
|
"steer_x": float(state["steer_x"]),
|
||||||
|
"steer_y": float(state["steer_y"]),
|
||||||
|
"look_dx": float(state["look_dx"]),
|
||||||
|
"look_dy": float(state["look_dy"]),
|
||||||
|
"confidence": float(state["confidence"]),
|
||||||
|
"on_target": bool(state["on_target"]),
|
||||||
|
}
|
||||||
|
tmp_path = self.export_path.with_suffix(self.export_path.suffix + ".tmp")
|
||||||
|
try:
|
||||||
|
with tmp_path.open("w", encoding="utf-8") as f:
|
||||||
|
json.dump(payload, f, ensure_ascii=True, indent=2)
|
||||||
|
self._replace_export_file(tmp_path)
|
||||||
|
except PermissionError as exc:
|
||||||
|
self.export_retry_at = time.monotonic() + 0.25
|
||||||
|
self._warn_export_error(frame_id, exc)
|
||||||
|
return
|
||||||
|
except OSError as exc:
|
||||||
|
self.export_retry_at = time.monotonic() + 0.50
|
||||||
|
self._warn_export_error(frame_id, exc)
|
||||||
|
return
|
||||||
|
|
||||||
|
self.export_retry_at = 0.0
|
||||||
|
self.last_export_frame = frame_id
|
||||||
|
|
||||||
|
def _replace_export_file(self, tmp_path):
|
||||||
|
last_err = None
|
||||||
|
for delay_sec in (0.0, 0.01, 0.03):
|
||||||
|
if delay_sec > 0.0:
|
||||||
|
time.sleep(delay_sec)
|
||||||
|
try:
|
||||||
|
tmp_path.replace(self.export_path)
|
||||||
|
return
|
||||||
|
except PermissionError as exc:
|
||||||
|
last_err = exc
|
||||||
|
|
||||||
|
if last_err is not None:
|
||||||
|
raise last_err
|
||||||
|
|
||||||
|
def _warn_export_error(self, frame_id, exc):
|
||||||
|
now = time.monotonic()
|
||||||
|
if (now - self.last_export_warn_ts) < 2.0:
|
||||||
|
return
|
||||||
|
print(f"[guidance] export skipped at frame {frame_id} for {self.export_path}: {exc}")
|
||||||
|
self.last_export_warn_ts = now
|
||||||
@ -0,0 +1,644 @@
|
|||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from config import *
|
||||||
|
|
||||||
|
# Helpers and utility functions
|
||||||
|
# =========================
|
||||||
|
# HELPERS
|
||||||
|
# =========================
|
||||||
|
|
||||||
|
def is_box_outside(box, w, h, margin=20):
|
||||||
|
x1, y1, x2, y2 = [float(v) for v in box]
|
||||||
|
return (x2 < -margin) or (y2 < -margin) or (x1 > (w + margin)) or (y1 > (h + margin))
|
||||||
|
|
||||||
|
|
||||||
|
def pick_best_det(dets_eff, ew, eh):
|
||||||
|
if not dets_eff:
|
||||||
|
return None
|
||||||
|
best = None
|
||||||
|
best_s = -1.0
|
||||||
|
for d in dets_eff:
|
||||||
|
b = clip_box(d[:4], ew, eh)
|
||||||
|
s = float(d[4])
|
||||||
|
a = box_area(b)
|
||||||
|
score = s + 0.00015 * a
|
||||||
|
if score > best_s:
|
||||||
|
best_s = score
|
||||||
|
best = b
|
||||||
|
return best
|
||||||
|
|
||||||
|
|
||||||
|
def pick_best_det_with_score(dets_eff, ew, eh):
|
||||||
|
if not dets_eff:
|
||||||
|
return None, 0.0
|
||||||
|
best = None
|
||||||
|
best_s = -1.0
|
||||||
|
best_conf = 0.0
|
||||||
|
for d in dets_eff:
|
||||||
|
b = clip_box(d[:4], ew, eh)
|
||||||
|
conf = float(d[4])
|
||||||
|
a = box_area(b)
|
||||||
|
score = conf + 0.00015 * a
|
||||||
|
if score > best_s:
|
||||||
|
best_s = score
|
||||||
|
best = b
|
||||||
|
best_conf = conf
|
||||||
|
return best, best_conf
|
||||||
|
|
||||||
|
|
||||||
|
def clamp(x, a, b):
|
||||||
|
return max(a, min(b, x))
|
||||||
|
|
||||||
|
|
||||||
|
def box_center(box):
|
||||||
|
x1, y1, x2, y2 = box
|
||||||
|
return np.array([(x1 + x2) * 0.5, (y1 + y2) * 0.5], dtype=np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def box_wh(box):
|
||||||
|
x1, y1, x2, y2 = box
|
||||||
|
return np.array([max(1.0, x2 - x1), max(1.0, y2 - y1)], dtype=np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def box_area(box):
|
||||||
|
w, h = box_wh(box)
|
||||||
|
return float(w * h)
|
||||||
|
|
||||||
|
|
||||||
|
def iou(a, b):
|
||||||
|
ax1, ay1, ax2, ay2 = a
|
||||||
|
bx1, by1, bx2, by2 = b
|
||||||
|
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
||||||
|
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
||||||
|
inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
|
||||||
|
if inter <= 0:
|
||||||
|
return 0.0
|
||||||
|
union = box_area(a) + box_area(b) - inter
|
||||||
|
return 0.0 if union <= 0 else float(inter / union)
|
||||||
|
|
||||||
|
|
||||||
|
def clip_box(box, w, h):
|
||||||
|
x1, y1, x2, y2 = box
|
||||||
|
x1 = clamp(float(x1), 0.0, float(w - 2))
|
||||||
|
y1 = clamp(float(y1), 0.0, float(h - 2))
|
||||||
|
x2 = clamp(float(x2), 1.0, float(w - 1))
|
||||||
|
y2 = clamp(float(y2), 1.0, float(h - 1))
|
||||||
|
if x2 <= x1 + 1:
|
||||||
|
x2 = x1 + 2
|
||||||
|
if y2 <= y1 + 1:
|
||||||
|
y2 = y1 + 2
|
||||||
|
return np.array([x1, y1, x2, y2], dtype=np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def crop_roi(frame, roi_box):
|
||||||
|
x1, y1, x2, y2 = map(int, roi_box)
|
||||||
|
return frame[y1:y2, x1:x2], x1, y1
|
||||||
|
|
||||||
|
|
||||||
|
def compute_hsv_hist(frame, box):
|
||||||
|
x1, y1, x2, y2 = map(int, box)
|
||||||
|
patch = frame[y1:y2, x1:x2]
|
||||||
|
if patch.size == 0:
|
||||||
|
return None
|
||||||
|
hsv = cv2.cvtColor(patch, cv2.COLOR_BGR2HSV)
|
||||||
|
hist = cv2.calcHist([hsv], [0, 1], None, list(HSV_HIST_BINS), [0, 180, 0, 256])
|
||||||
|
cv2.normalize(hist, hist, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)
|
||||||
|
return hist
|
||||||
|
|
||||||
|
|
||||||
|
def hsv_sim(h1, h2):
|
||||||
|
if h1 is None or h2 is None:
|
||||||
|
return 0.0
|
||||||
|
d = cv2.compareHist(h1, h2, cv2.HISTCMP_BHATTACHARYYA)
|
||||||
|
return float(1.0 - d)
|
||||||
|
|
||||||
|
|
||||||
|
def safe_ratio(a, b):
|
||||||
|
a = float(max(1e-3, a))
|
||||||
|
b = float(max(1e-3, b))
|
||||||
|
return max(a / b, b / a)
|
||||||
|
|
||||||
|
|
||||||
|
def box_ar(box):
|
||||||
|
w, h = box_wh(box)
|
||||||
|
return float(w / max(1e-3, h))
|
||||||
|
|
||||||
|
|
||||||
|
def blend_hist(ref_hist, new_hist, alpha=0.85):
|
||||||
|
if ref_hist is None:
|
||||||
|
return new_hist
|
||||||
|
if new_hist is None:
|
||||||
|
return ref_hist
|
||||||
|
out = (alpha * ref_hist) + ((1.0 - alpha) * new_hist)
|
||||||
|
out = out.astype(np.float32, copy=False)
|
||||||
|
cv2.normalize(out, out, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def preprocess_for_yolo(frame_bgr):
|
||||||
|
if (not ANALOG_FPV_MODE) or (not APPLY_YOLO_PREPROC):
|
||||||
|
return frame_bgr
|
||||||
|
|
||||||
|
out = frame_bgr
|
||||||
|
if PRE_BLUR_K >= 3:
|
||||||
|
k = int(PRE_BLUR_K)
|
||||||
|
if (k % 2) == 0:
|
||||||
|
k += 1
|
||||||
|
out = cv2.GaussianBlur(out, (k, k), 0)
|
||||||
|
|
||||||
|
ycrcb = cv2.cvtColor(out, cv2.COLOR_BGR2YCrCb)
|
||||||
|
y, cr, cb = cv2.split(ycrcb)
|
||||||
|
clahe = cv2.createCLAHE(clipLimit=float(PRE_CLAHE_CLIP), tileGridSize=PRE_CLAHE_TILE)
|
||||||
|
y = clahe.apply(y)
|
||||||
|
out = cv2.cvtColor(cv2.merge((y, cr, cb)), cv2.COLOR_YCrCb2BGR)
|
||||||
|
|
||||||
|
if PRE_UNSHARP > 1e-6:
|
||||||
|
blur = cv2.GaussianBlur(out, (0, 0), 1.1)
|
||||||
|
out = cv2.addWeighted(out, 1.0 + float(PRE_UNSHARP), blur, -float(PRE_UNSHARP), 0)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def in_norm_rect(x, y, w, h, rect_norm):
|
||||||
|
rx1, ry1, rx2, ry2 = rect_norm
|
||||||
|
return (rx1 * w) <= x <= (rx2 * w) and (ry1 * h) <= y <= (ry2 * h)
|
||||||
|
|
||||||
|
|
||||||
|
def in_osd_zone(center_x, center_y, frame_w, frame_h):
|
||||||
|
for z in REJECT_OSD_ZONES_NORM:
|
||||||
|
if in_norm_rect(center_x, center_y, frame_w, frame_h, z):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def filter_yolo_boxes_with_scores(result, frame_w, frame_h, offset_x=0, offset_y=0, min_conf=0.12):
|
||||||
|
dets_out = []
|
||||||
|
if result.boxes is None or len(result.boxes) == 0:
|
||||||
|
return dets_out
|
||||||
|
|
||||||
|
xyxy = result.boxes.xyxy.detach().cpu().numpy()
|
||||||
|
confs = result.boxes.conf.detach().cpu().numpy()
|
||||||
|
clss = result.boxes.cls.detach().cpu().numpy().astype(int)
|
||||||
|
|
||||||
|
for b, c, cls_id in zip(xyxy, confs, clss):
|
||||||
|
score = float(c)
|
||||||
|
if score < float(min_conf):
|
||||||
|
continue
|
||||||
|
if TARGET_CLASS_ID is not None and cls_id != TARGET_CLASS_ID:
|
||||||
|
continue
|
||||||
|
|
||||||
|
x1, y1, x2, y2 = map(float, b)
|
||||||
|
ww = max(0.0, x2 - x1)
|
||||||
|
hh = max(0.0, y2 - y1)
|
||||||
|
if ww <= 1.0 or hh <= 1.0:
|
||||||
|
continue
|
||||||
|
|
||||||
|
a = ww * hh
|
||||||
|
if a < MIN_BOX_AREA:
|
||||||
|
continue
|
||||||
|
if MAX_BOX_AREA > 0 and a > MAX_BOX_AREA:
|
||||||
|
continue
|
||||||
|
|
||||||
|
ar = ww / hh
|
||||||
|
if not (ASPECT_RANGE[0] <= ar <= ASPECT_RANGE[1]):
|
||||||
|
continue
|
||||||
|
|
||||||
|
gx1, gy1, gx2, gy2 = x1 + offset_x, y1 + offset_y, x2 + offset_x, y2 + offset_y
|
||||||
|
if REJECT_OSD_ZONES:
|
||||||
|
cx = 0.5 * (gx1 + gx2)
|
||||||
|
cy = 0.5 * (gy1 + gy2)
|
||||||
|
if (a <= REJECT_OSD_SMALL_AREA_MAX) and in_osd_zone(cx, cy, frame_w, frame_h):
|
||||||
|
continue
|
||||||
|
|
||||||
|
dets_out.append(np.array([gx1, gy1, gx2, gy2, score], dtype=np.float32))
|
||||||
|
|
||||||
|
return dets_out
|
||||||
|
|
||||||
|
|
||||||
|
def merge_close_part_dets(dets_eff, ref_box, frame_w, frame_h):
|
||||||
|
if (not PART_MERGE_ENABLE) or ref_box is None or len(dets_eff) < int(max(2, PART_MERGE_MIN_PARTS)):
|
||||||
|
return dets_eff, 0
|
||||||
|
|
||||||
|
ref = clip_box(ref_box, frame_w, frame_h)
|
||||||
|
ref_center = box_center(ref)
|
||||||
|
ref_diag = max(float(np.linalg.norm(box_wh(ref))), float(PART_MERGE_MIN_REF_DIAG))
|
||||||
|
ref_area = box_area(ref)
|
||||||
|
|
||||||
|
boxes = [clip_box(d[:4], frame_w, frame_h) for d in dets_eff]
|
||||||
|
scores = [float(d[4]) for d in dets_eff]
|
||||||
|
|
||||||
|
near_ids = []
|
||||||
|
for idx, b in enumerate(boxes):
|
||||||
|
c = box_center(b)
|
||||||
|
dist = float(np.linalg.norm(c - ref_center))
|
||||||
|
if (dist <= float(PART_MERGE_REF_DIST_DIAG) * ref_diag) or (iou(b, ref) >= float(PART_MERGE_IOU_FLOOR)):
|
||||||
|
near_ids.append(idx)
|
||||||
|
|
||||||
|
if len(near_ids) < int(max(2, PART_MERGE_MIN_PARTS)):
|
||||||
|
return dets_eff, 0
|
||||||
|
|
||||||
|
seed = max(near_ids, key=lambda ii: scores[ii])
|
||||||
|
cluster = {int(seed)}
|
||||||
|
|
||||||
|
changed = True
|
||||||
|
while changed:
|
||||||
|
changed = False
|
||||||
|
cluster_boxes = [boxes[ii] for ii in cluster]
|
||||||
|
ux1 = min(float(b[0]) for b in cluster_boxes)
|
||||||
|
uy1 = min(float(b[1]) for b in cluster_boxes)
|
||||||
|
ux2 = max(float(b[2]) for b in cluster_boxes)
|
||||||
|
uy2 = max(float(b[3]) for b in cluster_boxes)
|
||||||
|
union_box = clip_box([ux1, uy1, ux2, uy2], frame_w, frame_h)
|
||||||
|
union_center = box_center(union_box)
|
||||||
|
union_diag = max(float(np.linalg.norm(box_wh(union_box))), ref_diag)
|
||||||
|
|
||||||
|
for idx in near_ids:
|
||||||
|
if idx in cluster:
|
||||||
|
continue
|
||||||
|
b = boxes[idx]
|
||||||
|
c = box_center(b)
|
||||||
|
center_close = float(np.linalg.norm(c - union_center)) <= (float(PART_MERGE_CLUSTER_DIST_DIAG) * union_diag)
|
||||||
|
overlap = iou(b, union_box) >= float(PART_MERGE_IOU_FLOOR)
|
||||||
|
if center_close or overlap:
|
||||||
|
cluster.add(int(idx))
|
||||||
|
changed = True
|
||||||
|
|
||||||
|
if len(cluster) < int(max(2, PART_MERGE_MIN_PARTS)):
|
||||||
|
return dets_eff, 0
|
||||||
|
|
||||||
|
cluster_boxes = [boxes[ii] for ii in cluster]
|
||||||
|
ux1 = min(float(b[0]) for b in cluster_boxes)
|
||||||
|
uy1 = min(float(b[1]) for b in cluster_boxes)
|
||||||
|
ux2 = max(float(b[2]) for b in cluster_boxes)
|
||||||
|
uy2 = max(float(b[3]) for b in cluster_boxes)
|
||||||
|
union_box = clip_box([ux1, uy1, ux2, uy2], frame_w, frame_h)
|
||||||
|
union_area = box_area(union_box)
|
||||||
|
max_part_area = max(box_area(b) for b in cluster_boxes)
|
||||||
|
union_ar = box_ar(union_box)
|
||||||
|
|
||||||
|
if union_area < (float(PART_MERGE_MIN_UNION_GAIN) * max_part_area):
|
||||||
|
return dets_eff, 0
|
||||||
|
if union_area > (float(PART_MERGE_MAX_UNION_RATIO) * max(1.0, ref_area)):
|
||||||
|
return dets_eff, 0
|
||||||
|
if union_ar > float(PART_MERGE_MAX_ASPECT):
|
||||||
|
return dets_eff, 0
|
||||||
|
|
||||||
|
merged_score = min(0.99, max(scores[ii] for ii in cluster) + float(PART_MERGE_SCORE_BONUS))
|
||||||
|
merged_det = np.array([union_box[0], union_box[1], union_box[2], union_box[3], merged_score], dtype=np.float32)
|
||||||
|
|
||||||
|
out = [dets_eff[idx] for idx in range(len(dets_eff)) if idx not in cluster]
|
||||||
|
out.append(merged_det)
|
||||||
|
return out, int(len(cluster))
|
||||||
|
|
||||||
|
|
||||||
|
def _odd_ksize(k):
|
||||||
|
k = int(max(1, k))
|
||||||
|
if (k % 2) == 0:
|
||||||
|
k += 1
|
||||||
|
return k
|
||||||
|
|
||||||
|
|
||||||
|
def build_motion_mask(prev_gray, gray_now, affine=None):
|
||||||
|
if prev_gray is None or gray_now is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
h, w = gray_now.shape[:2]
|
||||||
|
ref = prev_gray
|
||||||
|
if affine is not None:
|
||||||
|
ref = cv2.warpAffine(
|
||||||
|
prev_gray,
|
||||||
|
affine,
|
||||||
|
(w, h),
|
||||||
|
flags=cv2.INTER_LINEAR,
|
||||||
|
borderMode=cv2.BORDER_REPLICATE
|
||||||
|
)
|
||||||
|
|
||||||
|
cur = gray_now
|
||||||
|
if MOTION_BLUR_K >= 3:
|
||||||
|
bk = _odd_ksize(MOTION_BLUR_K)
|
||||||
|
ref = cv2.GaussianBlur(ref, (bk, bk), 0)
|
||||||
|
cur = cv2.GaussianBlur(cur, (bk, bk), 0)
|
||||||
|
|
||||||
|
diff = cv2.absdiff(cur, ref)
|
||||||
|
_, mask = cv2.threshold(diff, int(MOTION_DIFF_THR), 255, cv2.THRESH_BINARY)
|
||||||
|
|
||||||
|
mk = int(max(1, MOTION_MORPH_K))
|
||||||
|
if mk > 1:
|
||||||
|
kernel = np.ones((mk, mk), dtype=np.uint8)
|
||||||
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
|
||||||
|
mask = cv2.dilate(mask, kernel, iterations=1)
|
||||||
|
|
||||||
|
return mask
|
||||||
|
|
||||||
|
|
||||||
|
def motion_zones_to_roi(mask, frame_w, frame_h):
|
||||||
|
if mask is None or mask.size == 0:
|
||||||
|
return None, 0
|
||||||
|
|
||||||
|
gx = int(max(1, MOTION_ZONE_GRID[0]))
|
||||||
|
gy = int(max(1, MOTION_ZONE_GRID[1]))
|
||||||
|
zone_w = float(frame_w) / float(gx)
|
||||||
|
zone_h = float(frame_h) / float(gy)
|
||||||
|
|
||||||
|
active = []
|
||||||
|
for iy in range(gy):
|
||||||
|
y1 = int(round(iy * zone_h))
|
||||||
|
y2 = int(round((iy + 1) * zone_h))
|
||||||
|
y2 = max(y2, y1 + 1)
|
||||||
|
for ix in range(gx):
|
||||||
|
x1 = int(round(ix * zone_w))
|
||||||
|
x2 = int(round((ix + 1) * zone_w))
|
||||||
|
x2 = max(x2, x1 + 1)
|
||||||
|
patch = mask[y1:y2, x1:x2]
|
||||||
|
if patch.size == 0:
|
||||||
|
continue
|
||||||
|
px = int(cv2.countNonZero(patch))
|
||||||
|
area = max(1, (x2 - x1) * (y2 - y1))
|
||||||
|
min_need = max(int(MOTION_ZONE_MIN_PIXELS), int(area * float(MOTION_ZONE_MIN_RATIO)))
|
||||||
|
if px >= min_need:
|
||||||
|
active.append((x1, y1, x2, y2))
|
||||||
|
|
||||||
|
active_count = len(active)
|
||||||
|
if active_count == 0:
|
||||||
|
return None, 0
|
||||||
|
|
||||||
|
max_active = int(round(float(gx * gy) * float(MOTION_MAX_ACTIVE_ZONE_RATIO)))
|
||||||
|
max_active = max(1, max_active)
|
||||||
|
if active_count > max_active:
|
||||||
|
return None, active_count
|
||||||
|
|
||||||
|
ux1 = min(z[0] for z in active) - int(MOTION_ROI_MARGIN)
|
||||||
|
uy1 = min(z[1] for z in active) - int(MOTION_ROI_MARGIN)
|
||||||
|
ux2 = max(z[2] for z in active) + int(MOTION_ROI_MARGIN)
|
||||||
|
uy2 = max(z[3] for z in active) + int(MOTION_ROI_MARGIN)
|
||||||
|
|
||||||
|
roi = clip_box([ux1, uy1, ux2, uy2], frame_w, frame_h)
|
||||||
|
|
||||||
|
rw = float(roi[2] - roi[0])
|
||||||
|
rh = float(roi[3] - roi[1])
|
||||||
|
min_side = float(max(2, MOTION_ROI_MIN_SIZE))
|
||||||
|
if rw < min_side or rh < min_side:
|
||||||
|
cx, cy = box_center(roi)
|
||||||
|
rw = max(rw, min_side)
|
||||||
|
rh = max(rh, min_side)
|
||||||
|
roi = clip_box([cx - rw * 0.5, cy - rh * 0.5, cx + rw * 0.5, cy + rh * 0.5], frame_w, frame_h)
|
||||||
|
|
||||||
|
return roi, active_count
|
||||||
|
|
||||||
|
|
||||||
|
def box_motion_stats(mask, box):
|
||||||
|
if mask is None or mask.size == 0:
|
||||||
|
return 0, 0.0
|
||||||
|
h, w = mask.shape[:2]
|
||||||
|
x1, y1, x2, y2 = map(int, clip_box(box, w, h))
|
||||||
|
patch = mask[y1:y2, x1:x2]
|
||||||
|
if patch.size == 0:
|
||||||
|
return 0, 0.0
|
||||||
|
px = int(cv2.countNonZero(patch))
|
||||||
|
ratio = float(px) / float(max(1, patch.shape[0] * patch.shape[1]))
|
||||||
|
return px, ratio
|
||||||
|
|
||||||
|
|
||||||
|
def wavelet_energy_haar(gray, box, margin=8, min_side=20):
|
||||||
|
if gray is None:
|
||||||
|
return 0.0
|
||||||
|
h, w = gray.shape[:2]
|
||||||
|
b = clip_box(box, w, h)
|
||||||
|
cx, cy = box_center(b)
|
||||||
|
bw, bh = box_wh(b)
|
||||||
|
side = max(float(min_side), max(float(bw), float(bh)) + 2.0 * float(margin))
|
||||||
|
x1 = int(max(0, round(cx - 0.5 * side)))
|
||||||
|
y1 = int(max(0, round(cy - 0.5 * side)))
|
||||||
|
x2 = int(min(w, round(cx + 0.5 * side)))
|
||||||
|
y2 = int(min(h, round(cy + 0.5 * side)))
|
||||||
|
if x2 - x1 < 4 or y2 - y1 < 4:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
p = gray[y1:y2, x1:x2]
|
||||||
|
if p.size == 0:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
# 1-level Haar on 2x2 quads, then normalized high-frequency energy.
|
||||||
|
if (p.shape[0] % 2) == 1:
|
||||||
|
p = p[:-1, :]
|
||||||
|
if (p.shape[1] % 2) == 1:
|
||||||
|
p = p[:, :-1]
|
||||||
|
if p.shape[0] < 4 or p.shape[1] < 4:
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
f = p.astype(np.float32) / 255.0
|
||||||
|
a = f[0::2, 0::2]
|
||||||
|
b0 = f[0::2, 1::2]
|
||||||
|
c0 = f[1::2, 0::2]
|
||||||
|
d0 = f[1::2, 1::2]
|
||||||
|
|
||||||
|
lh = a - b0 + c0 - d0
|
||||||
|
hl = a + b0 - c0 - d0
|
||||||
|
hh = a - b0 - c0 + d0
|
||||||
|
ll = 0.25 * (a + b0 + c0 + d0)
|
||||||
|
|
||||||
|
hf = (np.mean(np.abs(lh)) + np.mean(np.abs(hl)) + np.mean(np.abs(hh))) / 3.0
|
||||||
|
lf = float(np.mean(np.abs(ll)) + 1e-3)
|
||||||
|
return float(hf / lf)
|
||||||
|
|
||||||
|
|
||||||
|
def wavelet_hot_roi(
|
||||||
|
gray,
|
||||||
|
frame_w,
|
||||||
|
frame_h,
|
||||||
|
motion_mask=None,
|
||||||
|
pred_ref=None,
|
||||||
|
margin=20,
|
||||||
|
min_side=84,
|
||||||
|
pred_scale=3.0,
|
||||||
|
min_peak_ratio=1.8,
|
||||||
|
):
|
||||||
|
if gray is None:
|
||||||
|
return None, 0.0
|
||||||
|
|
||||||
|
g = gray
|
||||||
|
if (g.shape[0] % 2) == 1:
|
||||||
|
g = g[:-1, :]
|
||||||
|
if (g.shape[1] % 2) == 1:
|
||||||
|
g = g[:, :-1]
|
||||||
|
if g.shape[0] < 8 or g.shape[1] < 8:
|
||||||
|
return None, 0.0
|
||||||
|
|
||||||
|
f = g.astype(np.float32) / 255.0
|
||||||
|
a = f[0::2, 0::2]
|
||||||
|
b0 = f[0::2, 1::2]
|
||||||
|
c0 = f[1::2, 0::2]
|
||||||
|
d0 = f[1::2, 1::2]
|
||||||
|
|
||||||
|
lh = a - b0 + c0 - d0
|
||||||
|
hl = a + b0 - c0 - d0
|
||||||
|
hh = a - b0 - c0 + d0
|
||||||
|
e = (np.abs(lh) + np.abs(hl) + np.abs(hh)) / 3.0
|
||||||
|
e = cv2.GaussianBlur(e, (0, 0), 1.0)
|
||||||
|
|
||||||
|
if motion_mask is not None and motion_mask.size > 0:
|
||||||
|
m = cv2.resize(
|
||||||
|
(motion_mask.astype(np.float32) / 255.0),
|
||||||
|
(e.shape[1], e.shape[0]),
|
||||||
|
interpolation=cv2.INTER_AREA,
|
||||||
|
)
|
||||||
|
e *= (0.25 + 0.75 * m)
|
||||||
|
|
||||||
|
if pred_ref is not None:
|
||||||
|
pc = box_center(pred_ref) * 0.5 # energy map is downsampled by 2
|
||||||
|
ys, xs = np.indices(e.shape, dtype=np.float32)
|
||||||
|
sig = max(8.0, 0.22 * float(max(e.shape[0], e.shape[1])))
|
||||||
|
prior = np.exp(-((xs - pc[0]) ** 2 + (ys - pc[1]) ** 2) / (2.0 * sig * sig))
|
||||||
|
e *= (0.70 + 0.30 * prior)
|
||||||
|
|
||||||
|
peak = float(np.max(e))
|
||||||
|
mean = float(np.mean(e) + 1e-6)
|
||||||
|
peak_ratio = peak / mean
|
||||||
|
if peak_ratio < float(min_peak_ratio):
|
||||||
|
return None, peak_ratio
|
||||||
|
|
||||||
|
iy, ix = np.unravel_index(np.argmax(e), e.shape)
|
||||||
|
cx = (float(ix) + 0.5) * 2.0
|
||||||
|
cy = (float(iy) + 0.5) * 2.0
|
||||||
|
|
||||||
|
base_side = max(float(min_side), float(2.0 * margin + 24.0))
|
||||||
|
if pred_ref is not None:
|
||||||
|
pdiag = float(np.linalg.norm(box_wh(pred_ref)))
|
||||||
|
side = max(base_side, float(pred_scale) * pdiag)
|
||||||
|
else:
|
||||||
|
side = base_side
|
||||||
|
|
||||||
|
roi = clip_box([cx - 0.5 * side, cy - 0.5 * side, cx + 0.5 * side, cy + 0.5 * side], frame_w, frame_h)
|
||||||
|
return roi, peak_ratio
|
||||||
|
|
||||||
|
|
||||||
|
def recover_det_is_valid(det, motion_mask):
|
||||||
|
score = float(det[4])
|
||||||
|
if score < float(RECOVER_MIN_SCORE):
|
||||||
|
return False
|
||||||
|
|
||||||
|
b = det[:4]
|
||||||
|
a = box_area(b)
|
||||||
|
if a <= float(RECOVER_TINY_AREA_MAX):
|
||||||
|
if score < float(RECOVER_TINY_MIN_SCORE):
|
||||||
|
return False
|
||||||
|
px, ratio = box_motion_stats(motion_mask, b)
|
||||||
|
if px < int(RECOVER_TINY_MIN_MOTION_PIXELS) and ratio < float(RECOVER_TINY_MIN_MOTION_RATIO):
|
||||||
|
return False
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def make_focus_roi_from_box(box, frame_w, frame_h, margin=72, min_side=90):
|
||||||
|
b = clip_box(box, frame_w, frame_h)
|
||||||
|
cx, cy = box_center(b)
|
||||||
|
bw, bh = box_wh(b)
|
||||||
|
rw = max(float(min_side), float(bw) + 2.0 * float(margin))
|
||||||
|
rh = max(float(min_side), float(bh) + 2.0 * float(margin))
|
||||||
|
return clip_box([cx - 0.5 * rw, cy - 0.5 * rh, cx + 0.5 * rw, cy + 0.5 * rh], frame_w, frame_h)
|
||||||
|
|
||||||
|
|
||||||
|
def pick_preacq_det(dets_eff, pred_ref, frame_w, frame_h):
|
||||||
|
if not dets_eff or pred_ref is None:
|
||||||
|
return None
|
||||||
|
pc = box_center(pred_ref)
|
||||||
|
pdiag = float(np.linalg.norm(box_wh(pred_ref)))
|
||||||
|
near_lim = max(float(PREACQ_NEAR_MIN), float(PREACQ_NEAR_FACTOR) * pdiag)
|
||||||
|
|
||||||
|
best = None
|
||||||
|
best_s = -1e9
|
||||||
|
for d in dets_eff:
|
||||||
|
score = float(d[4])
|
||||||
|
if score < float(PREACQ_MIN_SCORE):
|
||||||
|
continue
|
||||||
|
b = clip_box(d[:4], frame_w, frame_h)
|
||||||
|
c = box_center(b)
|
||||||
|
dist = float(np.linalg.norm(c - pc))
|
||||||
|
if dist > near_lim:
|
||||||
|
continue
|
||||||
|
s = score - (0.0025 * dist)
|
||||||
|
if s > best_s:
|
||||||
|
best_s = s
|
||||||
|
best = np.array([b[0], b[1], b[2], b[3], score], dtype=np.float32)
|
||||||
|
return best
|
||||||
|
|
||||||
|
|
||||||
|
def track_residual_motion_ok(track, motion_mask, frame_w, frame_h):
|
||||||
|
if motion_mask is None:
|
||||||
|
return False
|
||||||
|
b = clip_box(track.tlbr, frame_w, frame_h)
|
||||||
|
px, ratio = box_motion_stats(motion_mask, b)
|
||||||
|
return (px >= int(EGO_RESIDUAL_MIN_PIXELS)) or (ratio >= float(EGO_RESIDUAL_MIN_RATIO))
|
||||||
|
|
||||||
|
|
||||||
|
def get_effective_frame(orig_bgr):
|
||||||
|
oh, ow = orig_bgr.shape[:2]
|
||||||
|
if not FORCE_EFFECTIVE_PAL:
|
||||||
|
return orig_bgr, 1.0, 1.0
|
||||||
|
eff = cv2.resize(orig_bgr, (EFFECTIVE_W, EFFECTIVE_H), interpolation=cv2.INTER_AREA)
|
||||||
|
sx = EFFECTIVE_W / float(ow)
|
||||||
|
sy = EFFECTIVE_H / float(oh)
|
||||||
|
return eff, sx, sy
|
||||||
|
|
||||||
|
|
||||||
|
def unscale_box(box_eff, sx, sy):
|
||||||
|
x1, y1, x2, y2 = box_eff
|
||||||
|
return np.array([x1 / sx, y1 / sy, x2 / sx, y2 / sy], dtype=np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def is_int_source(src):
|
||||||
|
if isinstance(src, int):
|
||||||
|
return True
|
||||||
|
if isinstance(src, str) and src.isdigit():
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def is_stream_source(src):
|
||||||
|
if not isinstance(src, str):
|
||||||
|
return False
|
||||||
|
s = src.lower()
|
||||||
|
return (
|
||||||
|
s.startswith("rtsp://")
|
||||||
|
or s.startswith("udp://")
|
||||||
|
or s.startswith("rtmp://")
|
||||||
|
or s.startswith("http://")
|
||||||
|
or s.startswith("https://")
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def open_source(source, backend=cv2.CAP_DSHOW):
|
||||||
|
if is_int_source(source):
|
||||||
|
cap = cv2.VideoCapture(int(source), backend)
|
||||||
|
return cap, "camera"
|
||||||
|
|
||||||
|
if isinstance(source, str):
|
||||||
|
if is_stream_source(source):
|
||||||
|
cap = cv2.VideoCapture(source)
|
||||||
|
return cap, "stream"
|
||||||
|
cap = cv2.VideoCapture(source)
|
||||||
|
return cap, "file"
|
||||||
|
|
||||||
|
raise ValueError(f"Unsupported SOURCE type: {type(source)}")
|
||||||
|
|
||||||
|
|
||||||
|
def get_frame_timestamp_seconds(cap, source_kind, frame_id, input_fps, loop_ts):
|
||||||
|
if TIMESTAMP_USE_SOURCE_CLOCK:
|
||||||
|
pos_ms = float(cap.get(cv2.CAP_PROP_POS_MSEC))
|
||||||
|
if np.isfinite(pos_ms) and pos_ms > 0.0:
|
||||||
|
return float(pos_ms * 1e-3), "source"
|
||||||
|
if source_kind == "file" and input_fps > 1.0:
|
||||||
|
return float(frame_id) / float(input_fps), "fps"
|
||||||
|
return float(loop_ts), "wall"
|
||||||
|
|
||||||
|
|
||||||
|
def sanitize_dt(frame_ts, prev_frame_ts, nominal_dt):
|
||||||
|
dt = float(nominal_dt)
|
||||||
|
if prev_frame_ts is not None:
|
||||||
|
raw_dt = float(frame_ts - prev_frame_ts)
|
||||||
|
if np.isfinite(raw_dt) and (1e-6 < raw_dt <= float(DT_RESET_GAP_SEC)):
|
||||||
|
dt = raw_dt
|
||||||
|
return float(clamp(dt, float(DT_MIN_SEC), float(DT_MAX_SEC)))
|
||||||
|
|
||||||
|
|
||||||
@ -0,0 +1,5 @@
|
|||||||
|
numpy==1.26.4
|
||||||
|
opencv-python==4.10.0.84
|
||||||
|
ultralytics==8.4.75
|
||||||
|
pymavlink==2.4.49
|
||||||
|
pyserial==3.5
|
||||||
@ -0,0 +1,63 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_bool(value):
|
||||||
|
if isinstance(value, bool):
|
||||||
|
return value
|
||||||
|
text = str(value).strip().lower()
|
||||||
|
return text in {"1", "true", "yes", "on"}
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_int(value):
|
||||||
|
return int(str(value).strip())
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_str(value):
|
||||||
|
return str(value)
|
||||||
|
|
||||||
|
|
||||||
|
def _parse_source(value):
|
||||||
|
text = str(value).strip()
|
||||||
|
if text and text.lstrip("+-").isdigit():
|
||||||
|
return int(text)
|
||||||
|
return text
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_env_overrides(namespace, spec):
|
||||||
|
for env_name, (key, parser) in spec.items():
|
||||||
|
if env_name not in os.environ:
|
||||||
|
continue
|
||||||
|
namespace[key] = parser(os.environ[env_name])
|
||||||
|
|
||||||
|
|
||||||
|
def apply_config_env_overrides(namespace):
|
||||||
|
_apply_env_overrides(
|
||||||
|
namespace,
|
||||||
|
{
|
||||||
|
"FPV_MODEL_PATH": ("MODEL_PATH", _parse_str),
|
||||||
|
"FPV_SOURCE": ("SOURCE", _parse_source),
|
||||||
|
"FPV_VIDEO_REALTIME": ("VIDEO_REALTIME", _parse_bool),
|
||||||
|
"FPV_SHOW_OUTPUT": ("SHOW_OUTPUT", _parse_bool),
|
||||||
|
"FPV_SAVE_INFER_VIDEO": ("SAVE_INFER_VIDEO", _parse_bool),
|
||||||
|
"FPV_OUT_VIDEO_PATH": ("OUT_VIDEO_PATH", _parse_str),
|
||||||
|
"FPV_GUIDANCE_EXPORT_ENABLE": ("GUIDANCE_EXPORT_ENABLE", _parse_bool),
|
||||||
|
"FPV_GUIDANCE_EXPORT_PATH": ("GUIDANCE_EXPORT_PATH", _parse_str),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def apply_intercept_env_overrides(namespace):
|
||||||
|
_apply_env_overrides(
|
||||||
|
namespace,
|
||||||
|
{
|
||||||
|
"FPV_AUTOPILOT_ENABLE": ("AUTOPILOT_ENABLE", _parse_bool),
|
||||||
|
"FPV_AUTOPILOT_BACKEND": ("AUTOPILOT_BACKEND", _parse_str),
|
||||||
|
"FPV_AUTOPILOT_JSON_PATH": ("AUTOPILOT_JSON_PATH", _parse_str),
|
||||||
|
"FPV_MAVLINK_CONNECTION": ("MAVLINK_CONNECTION", _parse_str),
|
||||||
|
"FPV_MSP_PORT": ("MSP_PORT", _parse_str),
|
||||||
|
"FPV_PROTO_UDP_ENABLE": ("PROTO_UDP_ENABLE", _parse_bool),
|
||||||
|
"FPV_PROTO_UDP_HOST": ("PROTO_UDP_HOST", _parse_str),
|
||||||
|
"FPV_PROTO_UDP_PORT": ("PROTO_UDP_PORT", _parse_int),
|
||||||
|
"FPV_PROTO_UDP_DESCRIPTOR": ("PROTO_UDP_DESCRIPTOR", _parse_int),
|
||||||
|
},
|
||||||
|
)
|
||||||
@ -0,0 +1,119 @@
|
|||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class HandoffDecision:
|
||||||
|
stale_lock_active: bool
|
||||||
|
allow_fast_handoff: bool
|
||||||
|
required_switch_hits: int
|
||||||
|
disable_prediction_hold: bool
|
||||||
|
reasons: list[str]
|
||||||
|
|
||||||
|
def reason_text(self) -> str:
|
||||||
|
return "|".join(self.reasons)
|
||||||
|
|
||||||
|
|
||||||
|
def evaluate_stale_lock(
|
||||||
|
*,
|
||||||
|
confirmed,
|
||||||
|
have_fresh_candidate,
|
||||||
|
target_has_fresh_update,
|
||||||
|
candidate_dist_px,
|
||||||
|
pred_diag_px,
|
||||||
|
miss_streak,
|
||||||
|
klt_valid,
|
||||||
|
klt_quality,
|
||||||
|
klt_iou,
|
||||||
|
candidate_motion_ok,
|
||||||
|
target_motion_ok,
|
||||||
|
stale_break_dist_diag,
|
||||||
|
stale_break_min_miss,
|
||||||
|
stale_break_klt_quality,
|
||||||
|
stale_break_klt_iou,
|
||||||
|
):
|
||||||
|
reasons = []
|
||||||
|
if not confirmed or not have_fresh_candidate:
|
||||||
|
return HandoffDecision(False, False, 0, False, [])
|
||||||
|
|
||||||
|
far_enough = candidate_dist_px >= max(40.0, stale_break_dist_diag * max(pred_diag_px, 1.0))
|
||||||
|
weak_anchor = (
|
||||||
|
(not klt_valid)
|
||||||
|
or (klt_quality < stale_break_klt_quality)
|
||||||
|
or (klt_iou < stale_break_klt_iou)
|
||||||
|
or (miss_streak >= stale_break_min_miss)
|
||||||
|
or (candidate_motion_ok and not target_motion_ok)
|
||||||
|
or (not target_has_fresh_update)
|
||||||
|
)
|
||||||
|
|
||||||
|
if not target_has_fresh_update:
|
||||||
|
reasons.append("target_not_fresh")
|
||||||
|
if not klt_valid:
|
||||||
|
reasons.append("klt_invalid")
|
||||||
|
if klt_quality < stale_break_klt_quality:
|
||||||
|
reasons.append("weak_klt_quality")
|
||||||
|
if klt_iou < stale_break_klt_iou:
|
||||||
|
reasons.append("weak_klt_iou")
|
||||||
|
if miss_streak >= stale_break_min_miss:
|
||||||
|
reasons.append("miss_streak")
|
||||||
|
if candidate_motion_ok and not target_motion_ok:
|
||||||
|
reasons.append("motion_conflict")
|
||||||
|
|
||||||
|
stale = bool(far_enough and weak_anchor)
|
||||||
|
return HandoffDecision(
|
||||||
|
stale_lock_active=stale,
|
||||||
|
allow_fast_handoff=stale,
|
||||||
|
required_switch_hits=0,
|
||||||
|
disable_prediction_hold=stale,
|
||||||
|
reasons=reasons if stale else [],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def compute_fast_handoff_hits(
|
||||||
|
*,
|
||||||
|
stale_lock_active,
|
||||||
|
motion_switch_mode,
|
||||||
|
approach_extra_hits,
|
||||||
|
fast_maneuver_extra_hits,
|
||||||
|
default_switch_hits,
|
||||||
|
fast_handoff_hits,
|
||||||
|
motion_switch_hits,
|
||||||
|
):
|
||||||
|
if stale_lock_active:
|
||||||
|
return int(fast_handoff_hits)
|
||||||
|
|
||||||
|
hits = int(default_switch_hits) + int(approach_extra_hits) + int(fast_maneuver_extra_hits)
|
||||||
|
if motion_switch_mode:
|
||||||
|
hits = max(hits, int(motion_switch_hits))
|
||||||
|
return int(hits)
|
||||||
|
|
||||||
|
|
||||||
|
def should_override_guidance(
|
||||||
|
*,
|
||||||
|
confirmed,
|
||||||
|
target_track_missing,
|
||||||
|
have_fresh_candidate,
|
||||||
|
candidate_matches_target,
|
||||||
|
miss_streak,
|
||||||
|
override_miss_ge,
|
||||||
|
):
|
||||||
|
return bool(
|
||||||
|
confirmed
|
||||||
|
and target_track_missing
|
||||||
|
and have_fresh_candidate
|
||||||
|
and (not candidate_matches_target)
|
||||||
|
and (int(miss_streak) >= int(override_miss_ge))
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def pick_guidance_override_box(*, track_candidate_box, raw_yolo_box):
|
||||||
|
if track_candidate_box is not None:
|
||||||
|
return track_candidate_box
|
||||||
|
return raw_yolo_box
|
||||||
|
|
||||||
|
|
||||||
|
def update_guidance_override_latch(*, new_box, prev_box, prev_ttl, max_ttl):
|
||||||
|
if new_box is not None:
|
||||||
|
return new_box, int(max(0, max_ttl))
|
||||||
|
if prev_box is not None and int(prev_ttl) > 0:
|
||||||
|
return prev_box, int(prev_ttl) - 1
|
||||||
|
return None, 0
|
||||||
@ -0,0 +1,60 @@
|
|||||||
|
import cv2
|
||||||
|
|
||||||
|
from config import *
|
||||||
|
from helpers import clamp, box_center, box_wh
|
||||||
|
|
||||||
|
# Template matching fallback
|
||||||
|
# =========================
|
||||||
|
|
||||||
|
def _crop_safe(img, x1, y1, x2, y2):
|
||||||
|
h, w = img.shape[:2]
|
||||||
|
x1 = int(clamp(x1, 0, w - 1))
|
||||||
|
y1 = int(clamp(y1, 0, h - 1))
|
||||||
|
x2 = int(clamp(x2, 0, w))
|
||||||
|
y2 = int(clamp(y2, 0, h))
|
||||||
|
if x2 <= x1 + 1 or y2 <= y1 + 1:
|
||||||
|
return None
|
||||||
|
return img[y1:y2, x1:x2]
|
||||||
|
|
||||||
|
|
||||||
|
def tm_update_template(gray_eff, box_eff):
|
||||||
|
if gray_eff is None or box_eff is None:
|
||||||
|
return None
|
||||||
|
cx, cy = box_center(box_eff)
|
||||||
|
half = TM_TEMPLATE_SIZE // 2
|
||||||
|
patch = _crop_safe(gray_eff, cx - half, cy - half, cx + half, cy + half)
|
||||||
|
if patch is None:
|
||||||
|
return None
|
||||||
|
if patch.shape[0] < 12 or patch.shape[1] < 12:
|
||||||
|
return None
|
||||||
|
return patch.copy()
|
||||||
|
|
||||||
|
|
||||||
|
def tm_search(gray_eff, template, pred_center, miss_streak):
|
||||||
|
if (not TM_ENABLE) or gray_eff is None or template is None or pred_center is None:
|
||||||
|
return None
|
||||||
|
h, w = gray_eff.shape[:2]
|
||||||
|
cx, cy = float(pred_center[0]), float(pred_center[1])
|
||||||
|
scale = TM_SEARCH_SCALE_BASE + TM_SEARCH_SCALE_PER_MISS * float(miss_streak)
|
||||||
|
th, tw = template.shape[:2]
|
||||||
|
win_w = int(max(tw * scale, tw + 20))
|
||||||
|
win_h = int(max(th * scale, th + 20))
|
||||||
|
x1 = int(cx - win_w / 2)
|
||||||
|
y1 = int(cy - win_h / 2)
|
||||||
|
x2 = int(cx + win_w / 2)
|
||||||
|
y2 = int(cy + win_h / 2)
|
||||||
|
roi = _crop_safe(gray_eff, x1, y1, x2, y2)
|
||||||
|
if roi is None or roi.shape[0] < th + 2 or roi.shape[1] < tw + 2:
|
||||||
|
return None
|
||||||
|
res = cv2.matchTemplate(roi, template, cv2.TM_CCOEFF_NORMED)
|
||||||
|
_, maxv, _, maxl = cv2.minMaxLoc(res)
|
||||||
|
if maxv < TM_MIN_SCORE:
|
||||||
|
return None
|
||||||
|
bx = x1 + maxl[0]
|
||||||
|
by = y1 + maxl[1]
|
||||||
|
mcx = bx + tw / 2
|
||||||
|
mcy = by + th / 2
|
||||||
|
return (mcx, mcy, float(maxv), tw, th)
|
||||||
|
|
||||||
|
|
||||||
|
# =========================
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@ -0,0 +1,63 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import struct
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from autopilot_bridge import (
|
||||||
|
AutopilotCommand,
|
||||||
|
build_proto_udp_payload,
|
||||||
|
select_proto_udp_target_state,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_select_proto_udp_target_state_uses_tracker_state_when_detection_is_lost():
|
||||||
|
assert select_proto_udp_target_state(confirmed=False, miss_streak=0) == 0
|
||||||
|
assert select_proto_udp_target_state(confirmed=True, miss_streak=0) == 1
|
||||||
|
assert select_proto_udp_target_state(confirmed=True, miss_streak=4) == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_proto_udp_payload_encodes_offsets_and_bbox_area():
|
||||||
|
cmd = AutopilotCommand()
|
||||||
|
cmd.target_id = 7
|
||||||
|
cmd.target_state = 1
|
||||||
|
cmd.frame_w = 1000
|
||||||
|
cmd.frame_h = 500
|
||||||
|
cmd.aim_x = 650.0
|
||||||
|
cmd.aim_y = 150.0
|
||||||
|
cmd.bbox_w = 200.0
|
||||||
|
cmd.bbox_h = 50.0
|
||||||
|
|
||||||
|
payload = build_proto_udp_payload(cmd)
|
||||||
|
unpacked = struct.unpack("<BBBhhbbB", payload)
|
||||||
|
|
||||||
|
assert unpacked == (
|
||||||
|
1, # descriptor
|
||||||
|
7, # object_id
|
||||||
|
1, # target_state
|
||||||
|
100, # vertical px: up is positive
|
||||||
|
150, # horizontal px: right is positive
|
||||||
|
40, # vertical percent relative to half frame height
|
||||||
|
30, # horizontal percent relative to half frame width
|
||||||
|
2, # bbox area percent
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_proto_udp_payload_returns_zeroed_target_fields_when_target_is_missing():
|
||||||
|
cmd = AutopilotCommand()
|
||||||
|
cmd.target_id = None
|
||||||
|
cmd.target_state = 0
|
||||||
|
cmd.frame_w = 1280
|
||||||
|
cmd.frame_h = 720
|
||||||
|
cmd.aim_x = None
|
||||||
|
cmd.aim_y = None
|
||||||
|
cmd.bbox_w = None
|
||||||
|
cmd.bbox_h = None
|
||||||
|
|
||||||
|
payload = build_proto_udp_payload(cmd)
|
||||||
|
unpacked = struct.unpack("<BBBhhbbB", payload)
|
||||||
|
|
||||||
|
assert unpacked == (1, 0, 0, 0, 0, 0, 0, 0)
|
||||||
@ -0,0 +1,16 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from decision_logger import TrackingDecisionLogger
|
||||||
|
|
||||||
|
|
||||||
|
def test_decision_logger_contains_handoff_fields():
|
||||||
|
assert "stale_lock_active" in TrackingDecisionLogger.FIELDNAMES
|
||||||
|
assert "stale_lock_reason" in TrackingDecisionLogger.FIELDNAMES
|
||||||
|
assert "fast_handoff_active" in TrackingDecisionLogger.FIELDNAMES
|
||||||
|
assert "guidance_reset_event" in TrackingDecisionLogger.FIELDNAMES
|
||||||
@ -0,0 +1,58 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from guidance import ScreenGuidanceController
|
||||||
|
|
||||||
|
|
||||||
|
def test_guidance_reset_reinitializes_smoothed_point():
|
||||||
|
ctrl = ScreenGuidanceController()
|
||||||
|
ctrl.smoothed_point = np.array([100.0, 200.0], dtype=np.float32)
|
||||||
|
ctrl.smoothed_cmd = np.array([0.7, -0.4], dtype=np.float32)
|
||||||
|
|
||||||
|
ctrl.reset_for_target_switch(np.array([500.0, 300.0], dtype=np.float32))
|
||||||
|
|
||||||
|
assert np.allclose(ctrl.smoothed_point, np.array([500.0, 300.0], dtype=np.float32))
|
||||||
|
assert np.allclose(ctrl.smoothed_cmd, np.zeros(2, dtype=np.float32))
|
||||||
|
|
||||||
|
|
||||||
|
def test_guidance_reset_can_damp_instead_of_zero_when_requested():
|
||||||
|
ctrl = ScreenGuidanceController()
|
||||||
|
ctrl.smoothed_cmd = np.array([0.8, -0.6], dtype=np.float32)
|
||||||
|
|
||||||
|
ctrl.reset_for_target_switch(np.array([10.0, 20.0], dtype=np.float32), cmd_damp=0.25)
|
||||||
|
|
||||||
|
assert np.allclose(ctrl.smoothed_cmd, np.array([0.2, -0.15], dtype=np.float32))
|
||||||
|
|
||||||
|
|
||||||
|
def test_guidance_override_box_takes_priority_over_stale_locked_box():
|
||||||
|
ctrl = ScreenGuidanceController()
|
||||||
|
|
||||||
|
state = ctrl.update(
|
||||||
|
frame_id=1,
|
||||||
|
frame_w=1000,
|
||||||
|
frame_h=600,
|
||||||
|
confirmed=True,
|
||||||
|
locked_box=np.array([50.0, 50.0, 150.0, 150.0], dtype=np.float32),
|
||||||
|
pred_box=np.array([60.0, 60.0, 160.0, 160.0], dtype=np.float32),
|
||||||
|
override_box=np.array([800.0, 100.0, 900.0, 200.0], dtype=np.float32),
|
||||||
|
target_id=1,
|
||||||
|
target_track_score=0.55,
|
||||||
|
klt_valid=True,
|
||||||
|
klt_quality=0.9,
|
||||||
|
miss_streak=5,
|
||||||
|
vx=0.0,
|
||||||
|
vy=0.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert state["status"] == "REACQ"
|
||||||
|
assert 845.0 <= state["aim_x"] <= 855.0
|
||||||
|
assert 145.0 <= state["aim_y"] <= 155.0
|
||||||
|
assert state["box_w"] == 100.0
|
||||||
|
assert state["box_h"] == 100.0
|
||||||
@ -0,0 +1,18 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import importlib
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
|
||||||
|
def test_local_bytetrack_module_is_importable_from_project_root():
|
||||||
|
module = importlib.import_module("bytetrack_min_aggressive")
|
||||||
|
|
||||||
|
assert hasattr(module, "BYTETracker")
|
||||||
|
|
||||||
|
|
||||||
|
def test_local_bytetrack_module_file_exists_in_project_root():
|
||||||
|
assert (PROJECT_ROOT / "bytetrack_min_aggressive.py").exists()
|
||||||
@ -0,0 +1,57 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from runtime_env import apply_config_env_overrides, apply_intercept_env_overrides
|
||||||
|
|
||||||
|
|
||||||
|
def test_apply_config_env_overrides_parses_source_device_and_flags(monkeypatch):
|
||||||
|
namespace = {
|
||||||
|
"SOURCE": "video.mp4",
|
||||||
|
"MODEL_PATH": "best.pt",
|
||||||
|
"SHOW_OUTPUT": True,
|
||||||
|
"SAVE_INFER_VIDEO": True,
|
||||||
|
"OUT_VIDEO_PATH": "out.mp4",
|
||||||
|
}
|
||||||
|
|
||||||
|
monkeypatch.setenv("FPV_SOURCE", "0")
|
||||||
|
monkeypatch.setenv("FPV_MODEL_PATH", "/models/best.pt")
|
||||||
|
monkeypatch.setenv("FPV_SHOW_OUTPUT", "false")
|
||||||
|
monkeypatch.setenv("FPV_SAVE_INFER_VIDEO", "0")
|
||||||
|
monkeypatch.setenv("FPV_OUT_VIDEO_PATH", "/outputs/run.mp4")
|
||||||
|
|
||||||
|
apply_config_env_overrides(namespace)
|
||||||
|
|
||||||
|
assert namespace["SOURCE"] == 0
|
||||||
|
assert namespace["MODEL_PATH"] == "/models/best.pt"
|
||||||
|
assert namespace["SHOW_OUTPUT"] is False
|
||||||
|
assert namespace["SAVE_INFER_VIDEO"] is False
|
||||||
|
assert namespace["OUT_VIDEO_PATH"] == "/outputs/run.mp4"
|
||||||
|
|
||||||
|
|
||||||
|
def test_apply_intercept_env_overrides_switches_to_proto_udp(monkeypatch):
|
||||||
|
namespace = {
|
||||||
|
"AUTOPILOT_ENABLE": True,
|
||||||
|
"AUTOPILOT_BACKEND": "json",
|
||||||
|
"PROTO_UDP_ENABLE": False,
|
||||||
|
"PROTO_UDP_HOST": "127.0.0.1",
|
||||||
|
"PROTO_UDP_PORT": 5005,
|
||||||
|
}
|
||||||
|
|
||||||
|
monkeypatch.setenv("FPV_AUTOPILOT_ENABLE", "true")
|
||||||
|
monkeypatch.setenv("FPV_AUTOPILOT_BACKEND", "proto_udp")
|
||||||
|
monkeypatch.setenv("FPV_PROTO_UDP_ENABLE", "1")
|
||||||
|
monkeypatch.setenv("FPV_PROTO_UDP_HOST", "192.168.0.55")
|
||||||
|
monkeypatch.setenv("FPV_PROTO_UDP_PORT", "6001")
|
||||||
|
|
||||||
|
apply_intercept_env_overrides(namespace)
|
||||||
|
|
||||||
|
assert namespace["AUTOPILOT_ENABLE"] is True
|
||||||
|
assert namespace["AUTOPILOT_BACKEND"] == "proto_udp"
|
||||||
|
assert namespace["PROTO_UDP_ENABLE"] is True
|
||||||
|
assert namespace["PROTO_UDP_HOST"] == "192.168.0.55"
|
||||||
|
assert namespace["PROTO_UDP_PORT"] == 6001
|
||||||
@ -0,0 +1,150 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from target_handoff import (
|
||||||
|
HandoffDecision,
|
||||||
|
compute_fast_handoff_hits,
|
||||||
|
evaluate_stale_lock,
|
||||||
|
pick_guidance_override_box,
|
||||||
|
should_override_guidance,
|
||||||
|
update_guidance_override_latch,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_evaluate_stale_lock_requires_fresh_candidate_and_weak_anchor():
|
||||||
|
decision = evaluate_stale_lock(
|
||||||
|
confirmed=True,
|
||||||
|
have_fresh_candidate=True,
|
||||||
|
target_has_fresh_update=False,
|
||||||
|
candidate_dist_px=420.0,
|
||||||
|
pred_diag_px=70.0,
|
||||||
|
miss_streak=8,
|
||||||
|
klt_valid=False,
|
||||||
|
klt_quality=0.18,
|
||||||
|
klt_iou=0.12,
|
||||||
|
candidate_motion_ok=True,
|
||||||
|
target_motion_ok=False,
|
||||||
|
stale_break_dist_diag=3.5,
|
||||||
|
stale_break_min_miss=2,
|
||||||
|
stale_break_klt_quality=0.35,
|
||||||
|
stale_break_klt_iou=0.25,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decision.stale_lock_active is True
|
||||||
|
assert "target_not_fresh" in decision.reasons
|
||||||
|
assert "weak_klt_quality" in decision.reasons
|
||||||
|
|
||||||
|
|
||||||
|
def test_evaluate_stale_lock_rejects_close_candidate_when_anchor_is_healthy():
|
||||||
|
decision = evaluate_stale_lock(
|
||||||
|
confirmed=True,
|
||||||
|
have_fresh_candidate=True,
|
||||||
|
target_has_fresh_update=True,
|
||||||
|
candidate_dist_px=90.0,
|
||||||
|
pred_diag_px=70.0,
|
||||||
|
miss_streak=0,
|
||||||
|
klt_valid=True,
|
||||||
|
klt_quality=0.92,
|
||||||
|
klt_iou=0.71,
|
||||||
|
candidate_motion_ok=True,
|
||||||
|
target_motion_ok=True,
|
||||||
|
stale_break_dist_diag=3.5,
|
||||||
|
stale_break_min_miss=2,
|
||||||
|
stale_break_klt_quality=0.35,
|
||||||
|
stale_break_klt_iou=0.25,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decision.stale_lock_active is False
|
||||||
|
assert decision.reasons == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_compute_fast_handoff_hits_reduces_confirmation_under_stale_lock():
|
||||||
|
hits = compute_fast_handoff_hits(
|
||||||
|
stale_lock_active=True,
|
||||||
|
motion_switch_mode=False,
|
||||||
|
approach_extra_hits=2,
|
||||||
|
fast_maneuver_extra_hits=2,
|
||||||
|
default_switch_hits=4,
|
||||||
|
fast_handoff_hits=2,
|
||||||
|
motion_switch_hits=6,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert hits == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_handoff_decision_repr_is_stable_for_logging():
|
||||||
|
decision = HandoffDecision(
|
||||||
|
stale_lock_active=True,
|
||||||
|
allow_fast_handoff=True,
|
||||||
|
required_switch_hits=2,
|
||||||
|
disable_prediction_hold=True,
|
||||||
|
reasons=["target_not_fresh", "weak_klt_quality"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert decision.reason_text() == "target_not_fresh|weak_klt_quality"
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_override_guidance_when_target_is_missing_but_fresh_candidate_exists():
|
||||||
|
assert should_override_guidance(
|
||||||
|
confirmed=True,
|
||||||
|
target_track_missing=True,
|
||||||
|
have_fresh_candidate=True,
|
||||||
|
candidate_matches_target=False,
|
||||||
|
miss_streak=10,
|
||||||
|
override_miss_ge=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_not_override_guidance_before_miss_threshold():
|
||||||
|
assert not should_override_guidance(
|
||||||
|
confirmed=True,
|
||||||
|
target_track_missing=True,
|
||||||
|
have_fresh_candidate=True,
|
||||||
|
candidate_matches_target=False,
|
||||||
|
miss_streak=1,
|
||||||
|
override_miss_ge=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_pick_guidance_override_box_falls_back_to_raw_yolo_box():
|
||||||
|
yolo_box = [100.0, 120.0, 180.0, 200.0]
|
||||||
|
|
||||||
|
picked = pick_guidance_override_box(
|
||||||
|
track_candidate_box=None,
|
||||||
|
raw_yolo_box=yolo_box,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert picked == yolo_box
|
||||||
|
|
||||||
|
|
||||||
|
def test_update_guidance_override_latch_keeps_recent_box_when_new_candidate_is_missing():
|
||||||
|
prev_box = [100.0, 120.0, 180.0, 200.0]
|
||||||
|
|
||||||
|
latched_box, ttl = update_guidance_override_latch(
|
||||||
|
new_box=None,
|
||||||
|
prev_box=prev_box,
|
||||||
|
prev_ttl=2,
|
||||||
|
max_ttl=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert latched_box == prev_box
|
||||||
|
assert ttl == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_update_guidance_override_latch_replaces_memory_when_new_candidate_arrives():
|
||||||
|
new_box = [300.0, 320.0, 380.0, 400.0]
|
||||||
|
|
||||||
|
latched_box, ttl = update_guidance_override_latch(
|
||||||
|
new_box=new_box,
|
||||||
|
prev_box=None,
|
||||||
|
prev_ttl=0,
|
||||||
|
max_ttl=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert latched_box == new_box
|
||||||
|
assert ttl == 2
|
||||||
@ -0,0 +1,66 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(PROJECT_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(PROJECT_ROOT))
|
||||||
|
|
||||||
|
from track_score_policy import required_track_score, track_passes_score_gate
|
||||||
|
|
||||||
|
|
||||||
|
def test_required_track_score_is_stricter_for_switch_than_for_acquire():
|
||||||
|
acquire_floor = required_track_score(
|
||||||
|
confirmed=False,
|
||||||
|
is_switch_candidate=False,
|
||||||
|
weak_reacq_guard=False,
|
||||||
|
)
|
||||||
|
switch_floor = required_track_score(
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=True,
|
||||||
|
weak_reacq_guard=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert switch_floor > acquire_floor
|
||||||
|
|
||||||
|
|
||||||
|
def test_weak_reacquire_uses_reacquire_floor():
|
||||||
|
reacquire_floor = required_track_score(
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=False,
|
||||||
|
weak_reacq_guard=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert track_passes_score_gate(
|
||||||
|
track_score=reacquire_floor,
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=False,
|
||||||
|
weak_reacq_guard=True,
|
||||||
|
)
|
||||||
|
assert not track_passes_score_gate(
|
||||||
|
track_score=reacquire_floor - 0.01,
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=False,
|
||||||
|
weak_reacq_guard=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_switch_candidate_must_clear_switch_floor():
|
||||||
|
switch_floor = required_track_score(
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=True,
|
||||||
|
weak_reacq_guard=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert track_passes_score_gate(
|
||||||
|
track_score=switch_floor,
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=True,
|
||||||
|
weak_reacq_guard=False,
|
||||||
|
)
|
||||||
|
assert not track_passes_score_gate(
|
||||||
|
track_score=switch_floor - 0.01,
|
||||||
|
confirmed=True,
|
||||||
|
is_switch_candidate=True,
|
||||||
|
weak_reacq_guard=False,
|
||||||
|
)
|
||||||
@ -0,0 +1,38 @@
|
|||||||
|
def required_track_score(
|
||||||
|
*,
|
||||||
|
confirmed,
|
||||||
|
is_switch_candidate,
|
||||||
|
weak_reacq_guard,
|
||||||
|
acquire_floor=0.08,
|
||||||
|
reacquire_floor=0.07,
|
||||||
|
switch_floor=0.10,
|
||||||
|
):
|
||||||
|
if not confirmed:
|
||||||
|
return float(acquire_floor)
|
||||||
|
|
||||||
|
floor = float(reacquire_floor)
|
||||||
|
if weak_reacq_guard:
|
||||||
|
floor = max(floor, float(reacquire_floor))
|
||||||
|
if is_switch_candidate:
|
||||||
|
floor = max(floor, float(switch_floor))
|
||||||
|
return float(floor)
|
||||||
|
|
||||||
|
|
||||||
|
def track_passes_score_gate(
|
||||||
|
*,
|
||||||
|
track_score,
|
||||||
|
confirmed,
|
||||||
|
is_switch_candidate,
|
||||||
|
weak_reacq_guard,
|
||||||
|
acquire_floor=0.08,
|
||||||
|
reacquire_floor=0.07,
|
||||||
|
switch_floor=0.10,
|
||||||
|
):
|
||||||
|
return float(track_score) >= required_track_score(
|
||||||
|
confirmed=confirmed,
|
||||||
|
is_switch_candidate=is_switch_candidate,
|
||||||
|
weak_reacq_guard=weak_reacq_guard,
|
||||||
|
acquire_floor=acquire_floor,
|
||||||
|
reacquire_floor=reacquire_floor,
|
||||||
|
switch_floor=switch_floor,
|
||||||
|
)
|
||||||
@ -0,0 +1,140 @@
|
|||||||
|
import cv2
|
||||||
|
import numpy as np
|
||||||
|
import time
|
||||||
|
import threading
|
||||||
|
from collections import deque
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from config import *
|
||||||
|
from helpers import clip_box, crop_roi, preprocess_for_yolo, filter_yolo_boxes_with_scores
|
||||||
|
|
||||||
|
# Async YOLO worker
|
||||||
|
# =========================
|
||||||
|
|
||||||
|
class YOLOWorker:
|
||||||
|
def __init__(self, model):
|
||||||
|
self.model = model
|
||||||
|
self.req = deque(maxlen=YOLO_QUEUE_MAX)
|
||||||
|
self.res = deque(maxlen=1)
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
self.running = False
|
||||||
|
self.thread = threading.Thread(target=self._loop, daemon=True)
|
||||||
|
|
||||||
|
def start(self):
|
||||||
|
self.running = True
|
||||||
|
self.thread.start()
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
self.running = False
|
||||||
|
self.thread.join(timeout=1.0)
|
||||||
|
|
||||||
|
def submit(self, frame_eff_bgr, roi_box_eff, mode, ts):
|
||||||
|
with self.lock:
|
||||||
|
self.req.append((frame_eff_bgr, roi_box_eff, mode, ts))
|
||||||
|
|
||||||
|
def try_get(self):
|
||||||
|
with self.lock:
|
||||||
|
if not self.res:
|
||||||
|
return None
|
||||||
|
return self.res.pop()
|
||||||
|
|
||||||
|
def _loop(self):
|
||||||
|
while self.running:
|
||||||
|
item = None
|
||||||
|
with self.lock:
|
||||||
|
if self.req:
|
||||||
|
item = self.req.pop()
|
||||||
|
self.req.clear()
|
||||||
|
|
||||||
|
if item is None:
|
||||||
|
time.sleep(0.001)
|
||||||
|
continue
|
||||||
|
|
||||||
|
frame, roi_box, mode, ts = item
|
||||||
|
h, w = frame.shape[:2]
|
||||||
|
frame_infer = preprocess_for_yolo(frame)
|
||||||
|
dets = []
|
||||||
|
infer_ms = 0.0
|
||||||
|
used_roi = False
|
||||||
|
|
||||||
|
try:
|
||||||
|
if roi_box is not None:
|
||||||
|
roi_box = clip_box(roi_box, w, h)
|
||||||
|
crop, ox, oy = crop_roi(frame_infer, roi_box)
|
||||||
|
if crop.size > 0:
|
||||||
|
used_roi = True
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
with torch.inference_mode():
|
||||||
|
r = self.model(
|
||||||
|
crop,
|
||||||
|
conf=YOLO_CONF_EFFECTIVE,
|
||||||
|
imgsz=IMG_SIZE_ROI,
|
||||||
|
verbose=False,
|
||||||
|
max_det=MAX_DET,
|
||||||
|
device=DEVICE,
|
||||||
|
half=USE_HALF
|
||||||
|
)[0]
|
||||||
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
||||||
|
dets = filter_yolo_boxes_with_scores(
|
||||||
|
r,
|
||||||
|
frame_w=w,
|
||||||
|
frame_h=h,
|
||||||
|
offset_x=ox,
|
||||||
|
offset_y=oy,
|
||||||
|
min_conf=BT_LOW
|
||||||
|
)
|
||||||
|
|
||||||
|
else:
|
||||||
|
sh, sw = frame_infer.shape[:2]
|
||||||
|
short = min(sh, sw)
|
||||||
|
scale = 1.0
|
||||||
|
target = IMG_SIZE_FULL
|
||||||
|
if short > target:
|
||||||
|
scale = target / float(short)
|
||||||
|
small = cv2.resize(
|
||||||
|
frame_infer,
|
||||||
|
(int(sw * scale), int(sh * scale)),
|
||||||
|
interpolation=cv2.INTER_AREA
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
small = frame_infer
|
||||||
|
|
||||||
|
t0 = time.perf_counter()
|
||||||
|
with torch.inference_mode():
|
||||||
|
r = self.model(
|
||||||
|
small,
|
||||||
|
conf=YOLO_CONF_EFFECTIVE,
|
||||||
|
imgsz=IMG_SIZE_FULL,
|
||||||
|
verbose=False,
|
||||||
|
max_det=MAX_DET,
|
||||||
|
device=DEVICE,
|
||||||
|
half=USE_HALF
|
||||||
|
)[0]
|
||||||
|
infer_ms = (time.perf_counter() - t0) * 1000.0
|
||||||
|
|
||||||
|
dets_s = filter_yolo_boxes_with_scores(
|
||||||
|
r,
|
||||||
|
frame_w=small.shape[1],
|
||||||
|
frame_h=small.shape[0],
|
||||||
|
offset_x=0,
|
||||||
|
offset_y=0,
|
||||||
|
min_conf=BT_LOW
|
||||||
|
)
|
||||||
|
if scale != 1.0:
|
||||||
|
inv = 1.0 / scale
|
||||||
|
dets = [
|
||||||
|
np.array([d[0] * inv, d[1] * inv, d[2] * inv, d[3] * inv, d[4]], dtype=np.float32)
|
||||||
|
for d in dets_s
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
dets = dets_s
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
dets = []
|
||||||
|
infer_ms = 0.0
|
||||||
|
|
||||||
|
with self.lock:
|
||||||
|
self.res.append((dets, ts, mode, infer_ms, used_roi))
|
||||||
|
|
||||||
|
|
||||||
|
# =========================
|
||||||
Loading…
Reference in New Issue