MAI/helpers.py

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)))