#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""matseg.py — 鲁棒人像分割「前景 floor」，治 RVM 抠像把举起的手/胳膊误判成背景而消失。

思路（语义引导抠像 arxiv 2210.05210 的工程化）：RVM 给软 alpha（发丝边缘好但手置信度低会掉），
叠一个语义鲁棒的人像分割（RMBG-1.4 / IS-Net，训练数据含全身，手/胳膊都稳）做 floor，
final = max(RVM软alpha, 人像分割) → RVM 掉的整块前景被分割兜住、永不消失；边缘仍以 RVM 软边为主。

🔴 §11.5b 定案：分割**每帧同步**跑，不隔帧（隔帧会滞后→拖尾幻影）。所以它的单帧成本
   直接进预算，必须压到最低——这就是下面 GPU 路径存在的理由。

🔴 加载/推理失败不致命：pipeline 退回纯 RVM（只是手可能掉），由上层 use_seg_floor 开关兜。

────────────────────────────────────────────────────────────────────────
🔴 2026-08-31 GPU 直通改造（ai002 上实测 24.53ms → 11.39ms，快 2.15×）

原实现每帧在 CPU 上做 numpy 预处理再喂 ORT，实测分段（真实 process 内）：
    pre(CPU numpy) 7.09ms + run 14.34ms + post(CPU) 2.07ms = 23.49ms
三个浪费点，全部有实测支撑：
  ① 预处理在 CPU：cv2.resize 到 1024² + float 转换 + transpose，实测 7.32ms；
     同样的事在 GPU 上只要 **0.20ms**（36×）。
  ② `transpose(2,0,1)` 产出**非连续数组**，ORT 内部要再整理一遍：
     实测 run 非连续 13.94ms vs 连续 11.92ms —— 白吃 2ms。
  ③ 管线里 `x = to_gpu(img)` 已经在 GPU 上备好了 **RGB float NCHW [0,1]**，
     而 RMBG 要的正是 RGB 减 0.5 —— 从 x 出发只需 GPU resize + sub_(0.5)，
     **零 CPU 操作、零额外上传**；旧路径却从 CPU 的 bgr 重新走一遍。

新路径 `mask_np(x_gpu, out_hw)`：GPU resize → sub 0.5 → **ORT IO Binding 直接吃 GPU 指针**
（run_with_iobinding，输入输出都不过 CPU）→ GPU 上归一化 + resize → 最后一次 D2H 出口。
数值一致性实测：相关系数 1.000000、平均差 0.000012（maxdiff 0.071 落在边缘过渡带的极少数
像素上，源于 cv2.INTER_LINEAR 与 torch bilinear 的采样点差异）。

🔴 模型输入尺寸**写死 1024²**（`['batch_size',3,1024,1024]`），喂 768/512 直接
   INVALID_ARGUMENT。所以「降分辨率提速」这条路在当前权重上走不通，别再试。
"""
import os

import cv2
import numpy as np

RMBG_URLS = [
    "https://huggingface.co/briaai/RMBG-1.4/resolve/main/onnx/model.onnx",
]


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


def ensure_rmbg(path):
    if os.path.exists(path) and os.path.getsize(path) > 1_000_000:
        return path
    import urllib.request
    for u in RMBG_URLS:
        try:
            _log("下载 RMBG-1.4 …")
            urllib.request.urlretrieve(u, path)
            if os.path.getsize(path) > 1_000_000:
                _log(f"RMBG 就绪 {os.path.getsize(path)/1e6:.0f}MB")
                return path
        except Exception as e:
            _log("  失败:", e)
    raise RuntimeError("RMBG 下载失败")


def torch_no_grad(fn):
    """推理路径不建计算图（省开销，且避免 .numpy() 撞上 requires_grad）。
    写成装饰器而非顶层 `@torch.no_grad()`，是因为本模块把 torch 留作懒导入。"""
    def wrap(*a, **kw):
        import torch
        with torch.no_grad():
            return fn(*a, **kw)
    return wrap


class HumanSeg:
    """RMBG-1.4 人像分割。

    两条入口，输出语义完全一致（float32 [0,1] 前景概率图，1=人）：
      · `mask(bgr, (h,w))`        —— 旧的 CPU 路径，保留作回落与数值基准
      · `mask_np(x_gpu, (h,w))`   —— GPU 直通路径（快 2.15×），x_gpu = to_gpu(bgr) 的输出
    """

    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)
        self.model_path = model_path
        if "CUDAExecutionProvider" not in self.sess.get_providers():
            _log("⚠️ RMBG 未跑在 CUDA 上:", self.sess.get_providers())
        inp = self.sess.get_inputs()[0]
        self.in_name = inp.name
        sh = inp.shape
        self.size = sh[2] if isinstance(sh[2], int) and sh[2] > 0 else 1024
        self.out_name = self.sess.get_outputs()[0].name
        # GPU 直通需要 CUDA EP；没有就只能走 CPU 路径（IO Binding 到 cuda 会直接报错）
        self.gpu_io = "CUDAExecutionProvider" in self.sess.get_providers()
        self._obuf = None          # 输出缓冲，懒分配后复用（每帧重新分配会拖慢）
        self._warned = False

    # ---------------- 旧 CPU 路径（回落 / 数值基准） ----------------
    def mask(self, bgr, out_hw):
        a = cv2.resize(bgr, (self.size, self.size))
        x = (a[:, :, ::-1].astype(np.float32) / 255.0 - 0.5)     # RGB, mean0.5 std1
        # 🔴 transpose 产出非连续数组，ORT 要再整理一遍（实测白吃 2ms）→ 显式连续化
        x = np.ascontiguousarray(x.transpose(2, 0, 1)[None])
        y = np.squeeze(self.sess.run(None, {self.in_name: x})[0]).astype(np.float32)
        mn, mx = float(y.min()), float(y.max())
        y = (y - mn) / (mx - mn + 1e-6)
        return cv2.resize(y, (out_hw[1], out_hw[0]))

    # ---------------- GPU 直通路径 ----------------
    @torch_no_grad
    def mask_t(self, x_gpu, out_hw):
        """x_gpu: `to_gpu(bgr)` 的输出（RGB float NCHW，值域 [0,1]，在 cuda 上）。
        返回 cuda tensor [H,W] float32 [0,1]。调用方若要 numpy 用 `mask_np`。"""
        import torch
        import torch.nn.functional as F
        xin = F.interpolate(x_gpu, size=(self.size, self.size),
                            mode="bilinear", align_corners=False)
        xin = (xin - 0.5).contiguous()          # RMBG 归一化就是 RGB/255 - 0.5
        if self._obuf is None:
            self._obuf = torch.empty((1, 1, self.size, self.size),
                                     dtype=torch.float32, device=xin.device)
        b = self.sess.io_binding()
        b.bind_input(self.in_name, "cuda", 0, np.float32, tuple(xin.shape), xin.data_ptr())
        b.bind_output(self.out_name, "cuda", 0, np.float32,
                      tuple(self._obuf.shape), self._obuf.data_ptr())
        # 🔴🔴 跨 stream 竞态，这两行缺一不可（漏掉会得到**上一帧的掩膜**）：
        #   torch 在它自己的 CUDA stream 上算 interpolate/sub，ORT 在**另一个 stream** 上跑推理。
        #   不同步的话 ORT 会读到还没写完的输入、我会读到还没写完的输出缓冲（缓冲是复用的，
        #   于是拿到上一次的内容）。
        #   🔴 这个 bug 极其隐蔽：连续帧内容相近时数值几乎一样、单张图测一致性也测不出来
        #   （warmup 已把缓冲填成同样内容）；**只有换场景才暴露**——表现正是「拖尾幻影」。
        #   实测证据：8 张难例里 t1 的 IoU 掉到 0.287、maxdiff 1.0，且它的前景占比与上一张
        #   t0 完全相同（16.67%）= 直接returned 上一张的结果。加同步后 8/8 全部 ≥0.998。
        torch.cuda.synchronize()          # 输入侧：确保 torch 那边算完了 ORT 才读
        self.sess.run_with_iobinding(b)
        b.synchronize_outputs()           # 输出侧：确保 ORT 写完了我才读 _obuf
        y = self._obuf
        mn, mx = y.min(), y.max()
        y = (y - mn) / (mx - mn + 1e-6)
        return F.interpolate(y, size=tuple(out_hw), mode="bilinear",
                             align_corners=False)[0, 0]

    def mask_np(self, x_gpu, out_hw, bgr=None):
        """GPU 直通 + 出口一次 D2H，返回与 `mask()` 同语义的 numpy。
        失败自动回落 CPU 路径（需要 bgr；没给就只能把异常抛上去）。"""
        if self.gpu_io:
            try:
                return self.mask_t(x_gpu, out_hw).cpu().numpy()
            except Exception as e:
                if not self._warned:        # 只吼一次，别每帧刷屏
                    _log("⚠️ GPU 直通失败，本进程改走 CPU 路径:", type(e).__name__, e)
                    self._warned = True
                self.gpu_io = False
        if bgr is None:
            raise RuntimeError("GPU 直通不可用且未提供 bgr 回落输入")
        return self.mask(bgr, out_hw)
