#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""selfieseg.py — MediaPipe Selfie Segmentation（Google Meet 同款）每帧人像分割。

为什么是它（2026-08-13，替换掉异步 BiRefNet+背景门控的复杂堆叠）：
- 专为 selfie/视频会议人像训练：手、手臂、手持物在训练分布里 → 不丢"看似断开的手"；
- 454KB / 256×256 输入 → 5090 上 ~1-2ms，**每帧同步跑，零滞后** → 拖尾在机理上不存在；
- 这是 Zoom/Meet 级产品验证过的"webcam 换背景"正解；BiRefNet 42ms+ 只能异步 → 旧掩膜拖尾，
  背景门控又被"肤色 vs 肤色"和"首帧含人"两个盲点反噬（实测更糟）。
模型：huggingface qualcomm/MediaPipe-Selfie-Segmentation（BSD/AI-Hub 发布）。
🔴 加载/推理失败不致命：pipeline 退回纯 RVM。
"""
import os

import cv2
import numpy as np

# 🔴 qualcomm 仓库已把 .onnx 移去 AI Hub（仓库里只剩 release_assets.json，resolve 404）→
#    用 onnx-community 的导出（transformers.js 团队维护，fp32 主档）。
URL = ("https://huggingface.co/onnx-community/mediapipe_selfie_segmentation/resolve/main/"
       "onnx/model.onnx")


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


def ensure_selfieseg(path):
    if os.path.exists(path) and os.path.getsize(path) > 100_000:
        return path
    import urllib.request
    _log("下载 MediaPipe-Selfie-Segmentation (454KB) …")
    urllib.request.urlretrieve(URL, path)
    if os.path.getsize(path) < 100_000:
        raise RuntimeError("selfie-seg 下载不完整")
    return path


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 SelfieSeg:
    """mask(bgr, (h,w)) → float32 [0,1] 人像置信度（1=人）。自适应 NHWC/NCHW 与输出布局。"""

    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)
        i = self.sess.get_inputs()[0]
        self.in_name = i.name
        sh = list(i.shape)
        self.nhwc = (len(sh) == 4 and sh[-1] in (1, 3))          # [1,H,W,3]=NHWC
        if self.nhwc:
            self.h = sh[1] if isinstance(sh[1], int) else 256
            self.w = sh[2] if isinstance(sh[2], int) else 256
        else:                                                     # [1,3,H,W]=NCHW
            self.h = sh[2] if isinstance(sh[2], int) else 256
            self.w = sh[3] if isinstance(sh[3], int) else 256
        self.out_name = self.sess.get_outputs()[0].name
        # GPU 直通用：输出形状（动态维按输入尺寸补实）+ 复用缓冲
        osh = list(self.sess.get_outputs()[0].shape)
        self._oshape = [d if isinstance(d, int) and d > 0 else v
                        for d, v in zip(osh, [1, 1, self.h, self.w])]
        self._obuf = None
        self.gpu_io = "CUDAExecutionProvider" in self.sess.get_providers()
        _log(f"就绪 in={sh} out={osh} nhwc={self.nhwc} gpu_io={self.gpu_io} "
             f"providers={self.sess.get_providers()[0]}")

    # ---------------- GPU 直通（与 matseg 同一套做法，见那边顶部注释） ----------------
    @torch_no_grad
    def mask_t(self, x_gpu, out_hw):
        """x_gpu: `to_gpu(bgr)` 的输出（RGB float NCHW [0,1]，cuda）——**正是 mediapipe 口径**，
        连减均值都不用，直接 resize 就能喂。返回 cuda tensor [H,W] float32 [0,1]。

        🔴 那两行 stream 同步不能省：torch 与 ORT 跑在不同 CUDA stream 上，不同步会读到
           复用输出缓冲里**上一帧**的内容（RMBG 那边实测踩过，症状=拖尾幻影）。
        """
        import torch
        import torch.nn.functional as F
        xin = F.interpolate(x_gpu, size=(self.h, self.w), mode="bilinear",
                            align_corners=False).contiguous()
        if self.nhwc:
            xin = xin.permute(0, 2, 3, 1).contiguous()
        if self._obuf is None:
            self._obuf = torch.empty(tuple(self._oshape), 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())
        torch.cuda.synchronize()
        self.sess.run_with_iobinding(b)
        b.synchronize_outputs()
        y = self._obuf
        # 与 CPU 路径同口径：多通道取前景通道；超出 [0,1] 视作 logits 过 sigmoid
        y = y[:, -1:] if y.shape[1] in (1, 2) else y[:, :1]
        if float(y.min()) < 0.0 or float(y.max()) > 1.0:
            y = torch.sigmoid(y)
        return F.interpolate(y, size=tuple(out_hw), mode="bilinear",
                             align_corners=False)[0, 0]

    def mask(self, bgr, out_hw):
        a = cv2.resize(bgr, (self.w, self.h))
        x = a[:, :, ::-1].astype(np.float32) / 255.0              # RGB [0,1]（mediapipe 口径）
        x = x[None] if self.nhwc else x.transpose(2, 0, 1)[None]
        x = np.ascontiguousarray(x)
        y = np.squeeze(self.sess.run([self.out_name], {self.in_name: x})[0]).astype(np.float32)
        if y.ndim == 3:                                           # 多通道：取前景通道（softmax 后）
            if y.shape[0] in (1, 2):                              # CHW
                y = y[-1]
            else:                                                 # HWC
                y = y[..., -1]
        if y.min() < 0.0 or y.max() > 1.0:                        # logits → sigmoid
            y = 1.0 / (1.0 + np.exp(-y))
        return cv2.resize(y, (out_hw[1], out_hw[0]))
