detect_lite.py

import argparse
import math
import os
from collections import deque

import cv2
import numpy as np
import onnxruntime as ort


# same motion-filter constants as detect.py
MIN_MOTION_RATIO  = 0.002
MIN_MOVEMENT_PX   = 5
MAX_TRACK_DIST    = 80
MAX_TRACK_AGE     = 12
TRACE_LEN         = 10
CONFIRM_HITS      = 2

# class index 0 = bird (flying), 1 = idle_bird — we only care about 0
BIRD_CLASS = 0


def get_args():
    p = argparse.ArgumentParser(description="Lightweight bird detector — ONNX only, no PyTorch")
    p.add_argument("--model",  default="models/bird_v5.onnx")
    p.add_argument("--conf",   type=float, default=0.45)
    p.add_argument("--iou",    type=float, default=0.45,  help="NMS IoU threshold")
    p.add_argument("--cam",    type=int,   default=0)
    p.add_argument("--width",  type=int,   default=1280)
    p.add_argument("--height", type=int,   default=720)
    p.add_argument("--imgsz",  type=int,   default=640)
    p.add_argument("--save",   action="store_true")
    p.add_argument("--no-gui", action="store_true")
    # test without a camera
    p.add_argument("--image",  default=None, help="run on a single image file instead of camera")
    p.add_argument("--video",  default=None, help="run on a video file instead of camera")
    return p.parse_args()


# ── ONNX inference ──────────────────────────────────────────────────────────

def load_model(path):
    if not os.path.exists(path):
        raise FileNotFoundError(f"Model not found: {path}")
    sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
    print(f"Loaded: {path}")
    print(f"ONNX Runtime {ort.__version__}  providers: {sess.get_providers()}")
    return sess


def preprocess(frame, imgsz):
    """Resize + normalize to (1, 3, imgsz, imgsz) float32."""
    img = cv2.resize(frame, (imgsz, imgsz))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    img = np.transpose(img, (2, 0, 1))          # HWC → CHW
    return np.expand_dims(img, axis=0)           # → (1, 3, H, W)


def postprocess(output, conf_thresh, iou_thresh, orig_w, orig_h, imgsz):
    """
    YOLOv8 ONNX output shape: (1, num_classes+4, 8400)
    Rows 0-3 = cx, cy, w, h  (in imgsz pixels)
    Rows 4+  = class scores
    """
    preds = output[0]                            # (num_classes+4, 8400)
    preds = preds.T                              # (8400, num_classes+4)

    boxes_cxcywh = preds[:, :4]
    class_scores = preds[:, 4:]

    # only keep flying bird class
    bird_scores = class_scores[:, BIRD_CLASS]
    mask = bird_scores >= conf_thresh
    if not mask.any():
        return []

    boxes_cxcywh = boxes_cxcywh[mask]
    scores       = bird_scores[mask]

    # cx,cy,w,h → x1,y1,x2,y2  (still in imgsz space)
    cx, cy, w, h = boxes_cxcywh.T
    x1 = cx - w / 2
    y1 = cy - h / 2
    x2 = cx + w / 2
    y2 = cy + h / 2

    # scale back to original frame
    sx = orig_w / imgsz
    sy = orig_h / imgsz
    x1, x2 = (x1 * sx).astype(int), (x2 * sx).astype(int)
    y1, y2 = (y1 * sy).astype(int), (y2 * sy).astype(int)

    # NMS using OpenCV (no extra deps)
    boxes_xywh = np.stack([x1, y1, x2 - x1, y2 - y1], axis=1).tolist()
    idxs = cv2.dnn.NMSBoxes(boxes_xywh, scores.tolist(), conf_thresh, iou_thresh)

    results = []
    for i in (idxs.flatten() if len(idxs) else []):
        results.append({
            "box":    (x1[i], y1[i], x2[i], y2[i]),
            "conf":   float(scores[i]),
            "center": ((x1[i] + x2[i]) // 2, (y1[i] + y2[i]) // 2),
        })
    return results


# ── camera ──────────────────────────────────────────────────────────────────

def open_cam(index, w, h):
    for backend, name in [(cv2.CAP_V4L2, "V4L2"), (cv2.CAP_ANY, "default")]:
        cap = cv2.VideoCapture(index, backend)
        if not cap.isOpened():
            continue
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, w)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"MJPG"))
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        ok, frame = cap.read()
        if ok and frame is not None:
            print(f"Camera {index} via {name} "
                  f"({int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))}x"
                  f"{int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))})")
            return cap
        cap.release()
    raise RuntimeError(f"Couldn't open camera {index}")


# ── motion / tracking helpers ────────────────────────────────────────────────

def dist(a, b):
    return math.hypot(a[0] - b[0], a[1] - b[1])


def diff_mask(prev, curr):
    d = cv2.absdiff(prev, curr)
    _, m = cv2.threshold(d, 20, 255, cv2.THRESH_BINARY)
    return cv2.dilate(m, None, iterations=2)


def motion_in_box(mask, box):
    x1, y1, x2, y2 = box
    if x2 <= x1 or y2 <= y1:
        return 0.0
    roi = mask[y1:y2, x1:x2]
    return cv2.countNonZero(roi) / float(roi.size) if roi.size else 0.0


def find_track(tracks, c, skip):
    best, best_d = None, MAX_TRACK_DIST
    for tid, t in tracks.items():
        if tid in skip:
            continue
        d = dist(t["center"], c)
        if d < best_d:
            best_d, best = d, tid
    return best


# ── main loop ────────────────────────────────────────────────────────────────

def run_on_image(args, sess, input_name):
    """Test on a single image file — no camera needed."""
    frame = cv2.imread(args.image)
    if frame is None:
        raise FileNotFoundError(f"Could not read image: {args.image}")

    fh, fw = frame.shape[:2]
    blob   = preprocess(frame, args.imgsz)
    raw    = sess.run(None, {input_name: blob})
    dets   = postprocess(raw[0], args.conf, args.iou, fw, fh, args.imgsz)

    result = frame.copy()
    for d in dets:
        x1, y1, x2, y2 = d["box"]
        cv2.rectangle(result, (x1, y1), (x2, y2), (0, 220, 0), 2)
        cv2.putText(result, f"bird {d['conf']:.2f}", (x1, max(20, y1 - 6)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 220, 0), 2)

    print(f"Detected {len(dets)} bird(s) in {args.image}")
    for i, d in enumerate(dets):
        print(f"  [{i+1}] conf={d['conf']:.2f}  box={d['box']}")

    cv2.imwrite("result.jpg", result)
    print("Saved annotated image → result.jpg")

    if not args.no_gui:
        cv2.imshow("Result", result)
        print("Press any key to close")
        cv2.waitKey(0)
        cv2.destroyAllWindows()


def main():
    args = get_args()
    sess = load_model(args.model)
    input_name = sess.get_inputs()[0].name

    # image mode — useful when no camera is connected
    if args.image:
        run_on_image(args, sess, input_name)
        return

    cap  = open_cam(args.cam, args.width, args.height)

    writer = None
    if args.save:
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        writer = cv2.VideoWriter("out.mp4", fourcc, 20, (args.width, args.height))
        print("Saving to out.mp4")

    tracks    = {}
    next_tid  = 1
    prev_gray = None
    total     = 0
    fnum      = 0

    print(f"Confidence: {args.conf}  |  Press q to quit")
    print("Green = confirmed flying bird | Orange = candidate")
    print("-" * 55)

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Camera read failed")
            break

        fh, fw = frame.shape[:2]
        gray   = cv2.GaussianBlur(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), (5, 5), 0)
        mask   = diff_mask(prev_gray, gray) if prev_gray is not None else None
        prev_gray = gray

        blob = preprocess(frame, args.imgsz)
        raw_out = sess.run(None, {input_name: blob})
        detections = postprocess(raw_out[0], args.conf, args.iou, fw, fh, args.imgsz)

        used   = set()
        active = set()
        for det in detections:
            mr  = motion_in_box(mask, det["box"]) if mask is not None else 0.0
            tid = find_track(tracks, det["center"], used)

            if tid is None:
                tid = next_tid
                next_tid += 1
                tracks[tid] = {
                    "center": det["center"], "box": det["box"],
                    "conf": det["conf"], "age": 0,
                    "trace": deque([det["center"]], maxlen=TRACE_LEN),
                    "hits": 1 if mr >= MIN_MOTION_RATIO else 0,
                    "ok": False,
                }
            else:
                t = tracks[tid]
                moving = mr >= MIN_MOTION_RATIO or dist(t["center"], det["center"]) >= MIN_MOVEMENT_PX
                t.update({"center": det["center"], "box": det["box"],
                           "conf": det["conf"], "age": 0})
                t["trace"].append(det["center"])
                t["hits"] = t["hits"] + 1 if moving else max(0, t["hits"] - 1)
                t["ok"]   = t["hits"] >= CONFIRM_HITS

            used.add(tid)
            active.add(tid)

        for tid in list(tracks):
            if tid not in active:
                tracks[tid]["age"] += 1
            if tracks[tid]["age"] > MAX_TRACK_AGE:
                del tracks[tid]

        out       = frame.copy()
        confirmed = 0
        for tid, t in tracks.items():
            if t["age"] > 0:
                continue
            x1, y1, x2, y2 = t["box"]
            if t["ok"]:
                confirmed += 1
                color = (0, 220, 0)
                lbl   = f"bird {t['conf']:.2f}"
                pts   = list(t["trace"])
                for i in range(1, len(pts)):
                    cv2.line(out, pts[i-1], pts[i], (0, 220, 220),
                             max(1, 4 - (len(pts) - i) // 3))
                cv2.circle(out, t["center"], 5, (0, 220, 220), -1)
            else:
                color = (0, 140, 255)
                lbl   = f"? {t['conf']:.2f}"

            cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
            cv2.putText(out, lbl, (x1, max(20, y1 - 6)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.65, color, 2)

        if confirmed:
            total += 1

        status = (f"Confirmed: {confirmed}  Raw: {len(detections)}  "
                  f"Triggers: {total}  conf>={args.conf:.2f}")
        cv2.putText(out, status, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 255, 0), 2)

        if not args.no_gui:
            cv2.imshow("Bird Detector (lite)", out)
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            fnum += 1
            if fnum % 30 == 0:
                print(status)

        if writer:
            writer.write(out)

    cap.release()
    if writer:
        writer.release()
    cv2.destroyAllWindows()
    print(f"\nDone. Total trigger frames: {total}")


if __name__ == "__main__":
    main()