#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""birefnet.py — BiRefNet 每帧鲁棒人像抠像（主引擎，可切换取代 RVM 递归软抠像）。

治本：RVM 靠递归记忆做软抠像，快速移动/周边/贴脸的手因递归置信度跟不上而"丢失"（半透明/消失）。
BiRefNet 每帧从零做高分辨率二分语义分割，**无时序记忆 → 不会丢手**；有 matting 变体给软边。
授权 MIT（可商用，区别于 RMBG-2.0 的非商用）。模型 emrikol/birefnet-matting-onnx（源 ZhengPeng7/BiRefNet-matting）。

I/O（实测口径）：输入 `input_image` [1,3,1024,1024] float32，BGR→RGB /255 再 ImageNet 归一化；
输出 `output_image` [1,1,1024,1024] logits，过 sigmoid → [0,1] 软 alpha，再 resize 回原尺寸。
🔴 该导出用 grid_sample 替形变卷积，留意 ORT CUDA 支持；加载后自测单次耗时，>100ms=落 CPU 要处理。
🔴 加载/推理失败不致命：pipeline 退回 RVM，由上层 matte_engine 开关兜。
"""
import os

import cv2
import numpy as np

# 🔴 变体权衡（都 MIT 可商用）：matting@1024 质量最好但实测 142ms/帧(5fps 不可用)；
#    lite=轻骨干、保 1024 分辨率(细手指要分辨率)、有 fp16 → 速度/鲁棒最佳平衡，做默认。
#    512=分辨率减半更快但细节差；出问题按 config/tune 切。
VARIANTS = {
    "lite_fp16":  ("birefnet_lite_fp16.onnx",  "onnx-community/BiRefNet_lite-ONNX",    "onnx/model_fp16.onnx"),
    "lite":       ("birefnet_lite.onnx",       "onnx-community/BiRefNet_lite-ONNX",    "onnx/model.onnx"),
    "full_fp16":  ("birefnet_full_fp16.onnx",  "onnx-community/BiRefNet-ONNX",         "onnx/model_fp16.onnx"),
    "b512_fp16":  ("birefnet_512_fp16.onnx",   "onnx-community/BiRefNet_512x512-ONNX", "onnx/model_fp16.onnx"),
    "matting":    ("birefnet-matting.onnx",    "emrikol/birefnet-matting-onnx",       "birefnet-matting.onnx"),
}
# 🔴 真实耗时(RTX5090/ORT-CUDA,720p入): full_fp16=122ms lite_fp16=70.8ms b512_fp16=42ms —
#    任何变体同步每帧跑都塞不进15fps预算(~25ms)。b512 走异步旁路线程(~24Hz刷新)是唯一可行解。
DEFAULT_VARIANT = "b512_fp16"
_MEAN = np.array([0.485, 0.456, 0.406], np.float32)
_STD = np.array([0.229, 0.224, 0.225], np.float32)


def _log(*a):
    print("[birefnet]", *a, flush=True)


def ensure_birefnet(assets_dir, variant=DEFAULT_VARIANT):
    fname, repo, sub = VARIANTS[variant]
    path = os.path.join(assets_dir, fname)
    if os.path.exists(path) and os.path.getsize(path) > 20_000_000:
        return path
    import urllib.request
    url = f"https://huggingface.co/{repo}/resolve/main/{sub}"
    _log(f"下载 BiRefNet[{variant}] {url} …")
    urllib.request.urlretrieve(url, path)
    if os.path.getsize(path) < 20_000_000:
        raise RuntimeError("BiRefNet 下载不完整")
    _log(f"BiRefNet[{variant}] 就绪 {os.path.getsize(path)/1e6:.0f}MB")
    return path


class BiRefNetMatte:
    """每帧鲁棒 alpha。alpha(bgr, (h,w)) → float32 [0,1]（1=前景/人）。自适应 fp16/fp32 输入。"""

    def __init__(self, model_path, providers=None):
        import onnxruntime as ort
        providers = providers or ["CUDAExecutionProvider", "CPUExecutionProvider"]
        so = ort.SessionOptions(); so.log_severity_level = 3
        self.sess = ort.InferenceSession(model_path, so, providers=providers)
        if "CUDAExecutionProvider" not in self.sess.get_providers():
            _log("⚠️ BiRefNet 未跑在 CUDA 上:", self.sess.get_providers())
        i = self.sess.get_inputs()[0]
        self.in_name = i.name
        self.out_name = self.sess.get_outputs()[0].name
        sh = i.shape
        self.size = sh[2] if isinstance(sh[2], int) and sh[2] > 0 else 1024
        # fp16 导出输入 dtype 是 float16 → 必须喂 fp16，否则报类型错
        self.in_dtype = np.float16 if "float16" in (i.type or "") else np.float32

    def alpha(self, bgr, out_hw):
        a = cv2.resize(bgr, (self.size, self.size))
        x = a[:, :, ::-1].astype(np.float32) / 255.0        # BGR→RGB /255
        x = (x - _MEAN) / _STD                               # ImageNet 归一化
        x = np.ascontiguousarray(x.transpose(2, 0, 1)[None].astype(self.in_dtype))  # NCHW，按需 fp16
        y = np.squeeze(self.sess.run([self.out_name], {self.in_name: x})[0]).astype(np.float32)
        if y.min() < 0.0 or y.max() > 1.0:                   # 是 logits 才过 sigmoid（兼容已 sigmoid 的导出）
            y = 1.0 / (1.0 + np.exp(-y))
        return cv2.resize(y, (out_hw[1], out_hw[0]))
