swarm / companies/sparrow/seed/main.py

#!/usr/bin/env python3
"""Live two-panel viewer — the product, run demo-time by the boss.

    pixi run demo [--source scene|webcam]

WINDOW 1 "Sparrow — overlay":  live scene with tracked features and short
motion tails drawn on top, plus a metric strip.
WINDOW 2 "Sparrow — sparse top-down map": the sparse map panel, updating
live from map.snapshot(): points + the camera trail.
"""
import argparse

import cv2
import numpy as np

import core
import scene
from overlay import draw_overlay
from topmap import draw_topmap

FPS = 30.0


def frame_source(kind):
    if kind == "webcam":
        cap = cv2.VideoCapture(0)
        while True:
            ok, f = cap.read()
            if not ok:
                break
            yield f
        return

    t = 0.0
    while True:
        yield scene.render(t)
        t += 1 / FPS


def metric_strip(info, side="bottom"):
    h = 22
    return f"fps≈{info.get('fps', 0):4.0f}  features {info.get('features', 0):3d}  lost {info.get('lost_rate', 0):4.2f}  map {info.get('map', 0):4d}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", default="scene", choices=["scene", "webcam"])
    a = ap.parse_args()

    pipe = core.Pipeline()
    t_prev = cv2.getTickCount()
    trail = []
    for frame in frame_source(a.source):
        res = pipe.on_frame(frame)
        dt = (cv2.getTickCount() - t_prev) / cv2.getTickFrequency()
        t_prev = cv2.getTickCount()
        fps = 1.0 / max(dt, 1e-3)
        snap = pipe.map.snapshot()
        info = {"fps": fps, "features": res["features"], "lost": res["lost"],
                "map": snap.get("count", 0)}

        # camera trail for the top-down panel (image x ≈ scene x for the
        # seed — real pose estimation is T3+, engineers, you know the way)
        cam_xy = snap.get("camera_xy")
        if cam_xy:
            trail.append(tuple(cam_xy))
        trail[:] = trail[-250:]

        vis = draw_overlay(frame.copy(), snap, info)
        top = draw_topmap(snap, trail, (640, 400))
        strip_rect = (0, vis.shape[0] - 24, vis.shape[1], vis.shape[0])
        cv2.rectangle(vis, strip_rect[:2], (strip_rect[2], strip_rect[3]), (18, 18, 18), -1)
        cv2.putText(vis, metric_strip(info), (8, vis.shape[0] - 8),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (220, 255, 220), 1)

        cv2.imshow("Sparrow — overlay", vis)
        cv2.imshow("Sparrow — sparse top-down map", top)
        if cv2.waitKey(1) == 27:
            break
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()