feat: interactive zone calibration tool

tools/calibrate_zone.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""Interactive zone calibration tool.
+
+Click 4 points for truck area, 4 for counting area, 2 for count line.
+Keys: s=save, c=clear/restart, q/ESC=quit.
+"""
+
+import argparse
+import json
+import sys
+import cv2
+import numpy as np
+
+ZONE_KEYS = ["truck_area", "counting_area"]
+ZONE_SIZES = [4, 4]
+LINE_POINTS = 2
+
+
+def load_zones(path: str) -> dict:
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except FileNotFoundError:
+        return {"presets": {}, "active_preset": None}
+
+
+def save_zones(config: dict, path: str):
+    with open(path, "w") as f:
+        json.dump(config, f, indent=2)
+
+
+def draw_state(frame, points, current_zone, zone_idx, zone_sizes, zone_names):
+    vis = frame.copy()
+    colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]
+    total_zones = len(zone_names)
+
+    all_zones = []
+    offset = 0
+    for i, sz in enumerate(zone_sizes):
+        all_zones.append(points[offset : offset + sz])
+        offset += sz
+    if current_zone < total_zones:
+        all_zones.append(points[offset:])
+
+    for i, zone_pts in enumerate(all_zones):
+        if not zone_pts:
+            continue
+        color = colors[i % len(colors)]
+        if i < total_zones:
+            name = zone_names[i]
+            need = zone_sizes[i]
+        else:
+            name = "count_line"
+            need = LINE_POINTS
+
+        if len(zone_pts) > 1:
+            pts = np.array(zone_pts, dtype=np.int32)
+            if len(zone_pts) >= 3 and i < total_zones:
+                cv2.polylines(vis, [pts], True, color, 2)
+            else:
+                cv2.polylines(vis, [pts], False, color, 2)
+
+        for j, pt in enumerate(zone_pts):
+            cv2.circle(vis, pt, 6, color, -1)
+            cv2.putText(vis, str(j + 1), (pt[0] + 8, pt[1] - 8),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+        if i == current_zone:
+            status = f"[{name}] {len(zone_pts)}/{need} pts"
+        else:
+            status = f"{name}: {len(zone_pts)}/{need}"
+        y_pos = 30 + i * 25
+        cv2.putText(vis, status, (10, y_pos),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+
+    if current_zone < total_zones:
+        cur_name = zone_names[current_zone]
+    else:
+        cur_name = "count_line"
+    hint = f"Click {cur_name} point. s=save c=clear q=quit"
+    cv2.putText(vis, hint, (10, vis.shape[0] - 20),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+    return vis
+
+
+def run(video_path: str, preset_name: str, zones_path: str):
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        print(f"Error: cannot open {video_path}")
+        sys.exit(1)
+
+    ret, frame = cap.read()
+    cap.release()
+    if not ret:
+        print("Error: cannot read frame from video")
+        sys.exit(1)
+
+    h, w = frame.shape[:2]
+    print(f"Video: {w}x{h}")
+
+    points = []
+    all_sizes = list(ZONE_SIZES) + [LINE_POINTS]
+    all_names = list(ZONE_KEYS) + ["count_line"]
+    zone_idx = 0
+
+    def mouse_cb(event, x, y, flags, param):
+        nonlocal points, zone_idx
+        if event == cv2.EVENT_LBUTTONDOWN:
+            total = sum(all_sizes)
+            if len(points) < total:
+                points.append((x, y))
+                cur_total = sum(all_sizes[:zone_idx])
+                if zone_idx < len(all_sizes):
+                    cur_total += len(points) - sum(all_sizes[:zone_idx])
+                    if cur_total >= all_sizes[zone_idx] and zone_idx < len(all_sizes) - 1:
+                        zone_idx += 1
+
+    win = "Calibrate Zone"
+    cv2.namedWindow(win, cv2.WINDOW_NORMAL)
+    cv2.resizeWindow(win, min(w, 1280), min(h, 720))
+    cv2.setMouseCallback(win, mouse_cb)
+
+    while True:
+        vis = draw_state(frame, points, zone_idx, all_sizes, all_names)
+        cv2.imshow(win, vis)
+        key = cv2.waitKey(30) & 0xFF
+
+        if key in (ord("q"), 27):
+            cv2.destroyAllWindows()
+            print("Quit.")
+            return
+
+        if key == ord("c"):
+            points.clear()
+            zone_idx = 0
+            print("Cleared.")
+
+        if key == ord("s"):
+            total_needed = sum(all_sizes)
+            if len(points) < total_needed:
+                print(f"Need {total_needed} points, have {len(points)}. Continue clicking.")
+                continue
+
+            offset = 0
+            preset = {}
+            for i, sz in enumerate(ZONE_SIZES):
+                preset[ZONE_KEYS[i]] = [list(p) for p in points[offset : offset + sz]]
+                offset += sz
+
+            x_start = points[offset][0]
+            x_end = points[offset + 1][0]
+            line_y = min(points[offset][1], points[offset + 1][1])
+            preset["count_line"] = {"x_start": x_start, "x_end": x_end, "y": line_y}
+
+            config = load_zones(zones_path)
+            config.setdefault("presets", {})[preset_name] = preset
+            save_zones(config, zones_path)
+            print(f"Saved preset '{preset_name}' to {zones_path}")
+            cv2.destroyAllWindows()
+            return
+
+    cv2.destroyAllWindows()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Interactive zone calibration tool")
+    parser.add_argument("video_path", help="Path to video file")
+    parser.add_argument("--preset", default="calibrated", help="Preset name (default: calibrated)")
+    parser.add_argument("--zones", default="zones.json", help="Path to zones.json")
+    args = parser.parse_args()
+    run(args.video_path, args.preset, args.zones)
+
+
+if __name__ == "__main__":
+    main()