app_function.py

import gc
import os
import sys
import tempfile
import time
import warnings

import cv2
import gradio as gr
import h5py
import numpy as np
import rawpy
import scipy.io
import spaces
import torch
import yaml
import zipfile
from bm3d import bm3d
from dataclasses import dataclass, field
from typing import Optional

# 优先使用 pyarmor 加密后的 dist/isp_algos.py；不存在或加载失败时回退到源码 private/isp_algos.py
try:
    sys.path.insert(0, "./dist")
    from isp_algos import VST, inverse_VST, ddim, BiasLUT, SimpleNLF
except Exception as e:
    print(f"[WARN] 无法加载 dist/isp_algos.py: {e}，回退到 private/isp_algos.py")
    from private.isp_algos import VST, inverse_VST, ddim, BiasLUT, SimpleNLF

from utils import bayer2rggb, rggb2bayer, FastISP
from utils import big_image_split, big_image_merge, log, get_host_with_dir, rawread, load_weights
from archs import *

# ─────────────────────────────────────────────
# 并发上限：最多同时持有数据的用户数
MAX_CONCURRENT_USERS = 3
# ─────────────────────────────────────────────


# ══════════════════════════════════════════════
#  UserState：每个浏览器 Tab 独立持有一份
# ══════════════════════════════════════════════
@dataclass
class UserState:
    """每个用户 Session 私有的数据与参数，存放在 gr.State 中。"""

    # 图像数据（numpy，可被 pickle 序列化）
    raw_data:       Optional[np.ndarray] = None
    denoised_data:  Optional[np.ndarray] = None
    denoised_npy:   Optional[np.ndarray] = None
    denoised_rgb:   Optional[np.ndarray] = None   # uint8 RGB numpy，不放 PIL
    mask_data:      Optional[np.ndarray] = None

    # 处理参数字典
    p: dict = field(default_factory=lambda: {
        "ratio":    1.0,
        "ispgain":  1.0,
        "h":        2160,
        "w":        3840,
        "bl":       64.0,
        "wp":       1023.0,
        "gain":     0.0,
        "sigma":    0.0,
        "wb":       [2.0, 1.0, 2.0],
        "ccm":      None,        # None → 用 np.eye(3)
        "scale":    959.0,
        "ransac":   False,
        "ddim_mode": False,
        "clip":     False,
        "sigsnr":   1.03,
        "epoch":    10,
        "sigma_t":  0.8,
        "eta_t":    0.85,
        "patch_size": 1024,
    })

    # 是否持有有效数据（用于并发计数）
    has_data: bool = False

    # 用户是否通过 YAML 手动指定了 wb/ccm（避免被 RAW 元数据覆盖）
    manual_wb:  bool = False
    manual_ccm: bool = False

    def get_ccm(self) -> np.ndarray:
        return self.p["ccm"] if self.p["ccm"] is not None else np.eye(3)

    def update_param(self, param: str, value):
        if param in ("h", "w"):
            self.p[param] = int(value)
        else:
            self.p[param] = float(value)
        if param in ("wp", "bl"):
            self.p["scale"] = self.p["wp"] - self.p["bl"]

    def clear_images(self):
        self.raw_data      = None
        self.denoised_data = None
        self.denoised_npy  = None
        self.denoised_rgb  = None
        self.mask_data     = None
        self.has_data      = False
        gc.collect()


# ══════════════════════════════════════════════
#  ModelService：全局单例，只持有模型和 bias_lut
# ══════════════════════════════════════════════
class ModelService:
    """
    持有共享的模型权重和 BiasLUT。
    纯只读推理，不保存任何用户数据。
    """

    def __init__(self):
        self.device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.yond      = None   # YOND_anytest 实例（含 .net .args 等）
        self.bias_lut  = None
        self._args     = {}     # 缓存 yaml args，供外部读取 pipeline 参数

    # ── 并发计数 ──────────────────────────────
    _active_user_count = 0   # 类变量，简单计数

    @classmethod
    def increment_users(cls) -> bool:
        """尝试占一个用户槽，成功返回 True，满员返回 False。"""
        if cls._active_user_count >= MAX_CONCURRENT_USERS:
            return False
        cls._active_user_count += 1
        return True

    @classmethod
    def decrement_users(cls):
        cls._active_user_count = max(0, cls._active_user_count - 1)

    @classmethod
    def current_users(cls) -> int:
        return cls._active_user_count

    # ── 模型管理 ──────────────────────────────
    @property
    def is_loaded(self) -> bool:
        return self.yond is not None and getattr(self.yond, "net", None) is not None

    @property
    def args(self) -> dict:
        return self._args

    def load_config(self, config_path: str):
        self.yond = YOND_anytest(config_path, self.device)
        self._args = self.yond.args
        model_path = f"{self.yond.fast_ckpt}/{self.yond.yond_name}_last_model.pth"
        self._load_model(model_path)
        gr.Success(f"配置加载成功: {config_path}")
        gr.Info(f"当前设备: {self.device}")

    def _load_model(self, model_path: str):
        self.yond.load_model(model_path)
        self.bias_lut = BiasLUT(lut_path="checkpoints/bias_lut_2d.npy")
        if self.bias_lut is None:
            raise RuntimeError("BiasLUT 加载失败")
        gr.Success(f"模型加载成功: {model_path}")

    def unload(self):
        if self.yond is not None:
            del self.yond
            self.yond     = None
            self.bias_lut = None
            self._args    = {}
        torch.cuda.empty_cache()
        gr.Success("GPU 已释放，如需继续请重新加载配置")

    # ── 推理（ZeroGPU 装饰器保留在此） ────────
    @spaces.GPU
    def denoise(self, raw_vst: np.ndarray, patch_size: int, nsr: float, p: dict) -> np.ndarray:
        """
        VST 域去噪。
        raw_vst : (H, W, 4) float32，已归一化
        返回    : (H, W, 4) float32，去噪结果（未逆变换）
        """
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.yond.net = self.yond.net.to(device)

        t_raw = torch.from_numpy(raw_vst).float().to(device).permute(2, 0, 1)[None]

        if "guided" in self.yond.arch:
            t = torch.tensor(
                nsr * p["sigsnr"], dtype=t_raw.dtype, device=device
            ).view(-1, 1, 1, 1)

        target_size   = patch_size
        overlap_ratio = 1 / 8
        raw_inp, metadata = big_image_split(t_raw, target_size, overlap_ratio)
        raw_dn = torch.zeros_like(raw_inp[:, :4])

        with torch.no_grad():
            if p["ddim_mode"]:
                for i in range(raw_inp.shape[0]):
                    print(f"Patch: {i+1}/{len(raw_dn)}")
                    raw_dn[i] = ddim(
                        raw_inp[i][None].clip(None, 2),
                        self.yond.net, t,
                        epoch=p["epoch"], sigma_t=p["sigma_t"],
                        eta=p["eta_t"], sigma_corr=1.00,
                    )
            else:
                for i in range(raw_inp.shape[0]):
                    raw_dn[i] = self.yond.net(
                        raw_inp[i][None].clip(None, 2), t
                    ).clamp(0, None)

        raw_dn = big_image_merge(raw_dn, metadata, blend_mode="avg")
        return raw_dn[0].permute(1, 2, 0).detach().cpu().numpy()


# ── 全局单例 ──────────────────────────────────
model_service = ModelService()


# ══════════════════════════════════════════════
#  YAML 参数覆盖：白名单与校验
# ══════════════════════════════════════════════
# 允许用户通过上传 YAML 覆盖 UserState.p 中的可调字段。
# 键名必须在此白名单内，否则会被忽略并提示 Warning。
PARAM_SCHEMA = {
    # 颜色校正
    "wb":       {"type": list,  "desc": "白平衡增益，支持 1x3 [R,G,B] 或 1x4 [R,G1,B,G2]"},
    "ccm":      {"type": list,  "desc": "3x3 颜色校正矩阵"},

    # RAW 元数据
    "bl":       {"type": (int, float), "desc": "黑电平"},
    "wp":       {"type": (int, float), "desc": "白点"},
    "ratio":    {"type": (int, float), "desc": "数字增益 (DGain)"},
    "ispgain":  {"type": (int, float), "desc": "ISP 预览增益（仅可视化）"},

    # 噪声参数
    "gain":     {"type": (int, float), "desc": "系统增益 K"},
    "sigma":    {"type": (int, float), "desc": "读出噪声水平 σ"},
    "sigsnr":   {"type": (int, float), "desc": "信噪比缩放系数"},

    # 去噪流程控制
    "ddim_mode":{"type": bool, "desc": "是否使用 DDIM 采样"},
    "epoch":    {"type": int,   "desc": "DDIM 迭代步数"},
    "sigma_t":  {"type": (int, float), "desc": "DDIM sigma_t"},
    "eta_t":    {"type": (int, float), "desc": "DDIM eta_t"},
    "patch_size":{"type": int,  "desc": "去噪分块大小"},
}


def _validate_wb(wb_raw):
    """把用户输入的 wb 校验并规范化为 4 元 [R, G1, B, G2]。"""
    if not isinstance(wb_raw, (list, tuple)):
        raise ValueError("wb 必须是 list 或 tuple")
    if len(wb_raw) not in (3, 4):
        raise ValueError("wb 长度必须是 3 ([R,G,B]) 或 4 ([R,G1,B,G2])")
    try:
        wb = [float(x) for x in wb_raw]
    except Exception as e:
        raise ValueError(f"wb 元素必须可转为 float: {e}") from e
    if any(x <= 0 for x in wb):
        raise ValueError("wb 增益必须全部为正数")
    if len(wb) == 3:
        # [R, G, B] → [R, G, B, G]
        wb = [wb[0], wb[1], wb[2], wb[1]]
    return wb


def _validate_ccm(ccm_raw):
    """把用户输入的 ccm 校验为 3x3 float32 numpy 数组。"""
    if not isinstance(ccm_raw, (list, tuple)):
        raise ValueError("ccm 必须是 list 或 tuple")
    if len(ccm_raw) != 3:
        raise ValueError("ccm 必须是 3x3 矩阵（外层长度为 3）")
    for row in ccm_raw:
        if not isinstance(row, (list, tuple)) or len(row) != 3:
            raise ValueError("ccm 每一行必须是长度为 3 的 list/tuple")
    try:
        ccm = np.array(ccm_raw, dtype=np.float32).reshape(3, 3)
    except Exception as e:
        raise ValueError(f"ccm 无法转为 3x3 float32 矩阵: {e}") from e
    return ccm


def _validate_param_override(params: dict) -> dict:
    """
    校验并规范化用户 YAML 中的参数覆盖。
    返回可直接 update 到 state.p 的新字典。
    """
    validated = {}
    for key, value in params.items():
        if key not in PARAM_SCHEMA:
            gr.Warning(f"忽略未识别的参数: {key}")
            continue

        schema = PARAM_SCHEMA[key]
        expected_type = schema["type"]
        if not isinstance(value, expected_type):
            raise ValueError(f"参数 {key} 类型错误: 期望 {expected_type}，得到 {type(value)}")

        if key == "wb":
            validated[key] = _validate_wb(value)
        elif key == "ccm":
            validated[key] = _validate_ccm(value)
        else:
            validated[key] = value

    # 基础合理性检查
    if "bl" in validated and "wp" in validated:
        if validated["bl"] >= validated["wp"]:
            raise ValueError("黑电平 bl 必须小于白点 wp")
    if "bl" in validated and "wp" not in validated:
        # 仅更新 bl 时不需要检查，后续 update_param 会自动重算 scale
        pass

    return validated


# ══════════════════════════════════════════════
#  业务函数（纯函数：接收 state，返回 state）
# ══════════════════════════════════════════════

def _vst_denoise_pipeline(
    lr_raw: np.ndarray,
    state: UserState,
    patch_size: int,
) -> tuple[np.ndarray, np.ndarray]:
    """
    VST → Denoise → InvVST 完整流水线。
    返回 (denoised_rggb [0,1], denoised_bayer)
    """
    p = state.p
    lr_raw_np  = lr_raw * p["scale"]
    bias_base  = np.maximum(lr_raw_np, 0)
    bias       = model_service.bias_lut.get_lut(bias_base, K=p["gain"], sigGs=p["sigma"])

    raw_vst    = VST(lr_raw_np, p["sigma"], gain=p["gain"]) - bias
    lower      = VST(0,          p["sigma"], gain=p["gain"])
    upper      = VST(p["scale"], p["sigma"], gain=p["gain"])
    nsr        = 1.0 / (upper - lower)
    raw_vst    = (raw_vst - lower) / (upper - lower)

    raw_dn     = model_service.denoise(raw_vst, patch_size, nsr, p)

    raw_dn     = raw_dn * (upper - lower) + lower
    denoised   = inverse_VST(raw_dn, p["sigma"], gain=p["gain"]) / p["scale"]
    return denoised, rggb2bayer(denoised)


def _generate_preview(state: UserState) -> np.ndarray:
    p = state.p
    processed = (state.raw_data - p["bl"]) / p["scale"]
    rgb = FastISP(
        bayer2rggb(processed) * p["ratio"] * p["ispgain"],
        p["wb"], state.get_ccm(),
    )
    return (rgb.clip(0, 1) * 255).astype(np.uint8)


def _visualize_mask(state: UserState) -> np.ndarray:
    from matplotlib import pyplot as plt
    mask = state.mask_data
    if mask.ndim != 2:
        raise gr.Error("掩模必须是 2D 数组")
    cmap = plt.cm.viridis
    lut  = (cmap(np.linspace(0, 1, 256))[:, :3] * 255).astype(np.uint8)
    idx  = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
    rgb  = cv2.resize(lut[idx], (state.p["w"], state.p["h"]), interpolation=cv2.INTER_LINEAR)
    return rgb


def _generate_result(state: UserState) -> np.ndarray:
    p   = state.p
    rgb = FastISP(state.denoised_data * p["ispgain"], p["wb"], state.get_ccm())
    arr = (rgb.clip(0, 1) * 255).astype(np.uint8)
    state.denoised_rgb = arr
    return arr


# ──────────────────────────────────────────────
#  对外暴露的业务函数（供 app.py 绑定）
# ──────────────────────────────────────────────

def load_config(config_path: str, state: UserState):
    """加载模型配置（全局），同时将 pipeline 参数写入用户 state。"""
    try:
        model_service.load_config(config_path)
        args = model_service.args
        if "pipeline" in args:
            state.p.update(args["pipeline"])
        else:
            state.p.update({"epoch": 10, "sigma_t": 0.8, "eta_t": 0.85})
        return state
    except Exception as e:
        raise gr.Error(f"配置加载失败: {e}")


def load_params_yaml(file_path: str, state: UserState):
    """用户上传 YAML 覆盖当前 state.p 中的可调参数，并刷新预览。"""
    if file_path is None:
        raise gr.Error("请先上传参数 YAML 文件")

    try:
        with open(file_path, "r", encoding="utf-8") as f:
            params = yaml.safe_load(f)
    except Exception as e:
        raise gr.Error(f"YAML 读取失败: {e}")

    if not isinstance(params, dict):
        raise gr.Error("YAML 顶层必须是一个字典（key-value 映射）")

    try:
        validated = _validate_param_override(params)
    except ValueError as e:
        raise gr.Error(f"参数校验失败: {e}")

    # 标记用户手动指定的颜色参数
    if "wb" in validated:
        state.manual_wb = True
    if "ccm" in validated:
        state.manual_ccm = True

    # 应用覆盖
    state.p.update(validated)

    # 若 bl/wp 被修改，需要同步更新 scale
    state.p["scale"] = state.p["wp"] - state.p["bl"]

    gr.Success(f"参数覆盖成功: {list(validated.keys())}")

    # 若已有 RAW 数据，刷新预览
    if state.raw_data is not None:
        preview = _generate_preview(state)
        return preview, state
    return None, state


def download_param_template(state: UserState) -> str:
    """
    基于当前 state.p 生成参数模板 YAML（默认全部注释）。
    用户取消注释对应行即可覆盖该参数，保留注释则继续使用当前值。
    """
    p = state.p
    ccm = p["ccm"] if p["ccm"] is not None else np.eye(3, dtype=np.float32)
    ccm = ccm.tolist() if isinstance(ccm, np.ndarray) else ccm

    sections = [
        ("颜色校正", [
            ("wb", p["wb"], "白平衡增益，支持 1x3 [R,G,B] 或 1x4 [R,G1,B,G2]"),
            ("ccm", ccm, "3x3 颜色校正矩阵（从相机色彩空间 → sRGB）"),
        ]),
        ("RAW 元数据", [
            ("bl", p["bl"], "黑电平"),
            ("wp", p["wp"], "白点"),
            ("ratio", p["ratio"], "数字增益 (DGain)"),
            ("ispgain", p["ispgain"], "ISP 预览增益（仅可视化）"),
        ]),
        ("噪声参数", [
            ("gain", p["gain"], "系统增益 K"),
            ("sigma", p["sigma"], "读出噪声水平 σ"),
            ("sigsnr", p["sigsnr"], "信噪比缩放系数"),
        ]),
        ("去噪流程控制", [
            ("ddim_mode", p["ddim_mode"], "是否使用 DDIM 采样"),
            ("epoch", p["epoch"], "DDIM 迭代步数"),
            ("sigma_t", p["sigma_t"], "DDIM sigma_t"),
            ("eta_t", p["eta_t"], "DDIM eta_t"),
            ("patch_size", p["patch_size"], "去噪分块大小"),
        ]),
    ]

    lines = [
        "# YOND WebUI 参数覆盖文件",
        "# 上传后会被安全地合并到当前运行参数中。",
        "# 只支持下方列出的字段，其他字段会被忽略。",
        "# 每行默认已注释：取消注释即可覆盖对应参数，保留注释则继续使用当前值。",
        "",
    ]

    for section_title, params in sections:
        lines.append(f"# ---------- {section_title} ----------")
        for key, value, desc in params:
            lines.append(f"# {desc}")
            if key == "ccm":
                lines.append(f"# {key}:")
                for row in value:
                    lines.append(f"#   - {row}")
            else:
                lines.append(f"# {key}: {value}")
            lines.append("")
    lines.append("")

    content = "\n".join(str(line) for line in lines)
    with tempfile.NamedTemporaryFile(mode="w", suffix=".yml", delete=False, encoding="utf-8") as f:
        f.write(content)
        return f.name


def process_image(file_path: str, h, w, bl, wp, ratio, ispgain, state: UserState):
    """读取 RAW 文件，更新 state，返回预览图和元数据。"""
    gr.Info("正在可视化图像")

    # ── 并发限流 ──────────────────────────────
    if not state.has_data:
        if not ModelService.increment_users():
            raise gr.Error(
                f"当前已有 {MAX_CONCURRENT_USERS} 位用户在使用，请稍后再试"
            )
        state.has_data = True

    # ── 更新基础参数 ──────────────────────────
    for k, v in [("h", h), ("w", w), ("bl", bl), ("wp", wp), ("ratio", ratio), ("ispgain", ispgain)]:
        state.update_param(k, v)

    # 重置图像数据（保留用户手动指定的 wb/ccm）
    state.raw_data      = None
    state.denoised_data = None
    state.mask_data     = None
    if not state.manual_wb:
        state.p["wb"] = [2, 1, 2]
    if not state.manual_ccm:
        state.p["ccm"] = None

    try:
        ext = file_path.lower().rsplit(".", 1)[-1]

        if ext in ("arw", "dng", "nef", "cr2"):
            with rawpy.imread(str(file_path)) as raw:
                state.raw_data = raw.raw_image_visible.astype(np.float32)
                h_r, w_r = state.raw_data.shape
                bl_r = float(raw.black_level_per_channel[0])
                wp_r = float(raw.white_level)
                updates = {
                    "h": h_r, "w": w_r,
                    "bl": bl_r, "wp": wp_r, "scale": wp_r - bl_r,
                }
                # 仅当用户未手动指定时才用 RAW 元数据覆盖 wb/ccm
                if not state.manual_wb:
                    wb = np.array(raw.camera_whitebalance) / raw.camera_whitebalance[1]
                    updates["wb"] = wb.tolist()
                if not state.manual_ccm:
                    ccm = raw.color_matrix[:3, :3].astype(np.float32)
                    updates["ccm"] = ccm
                state.p.update(updates)

        elif ext in ("raw", "npy"):
            try:
                state.raw_data = np.fromfile(file_path, dtype=np.uint16).reshape(
                    state.p["h"], state.p["w"]
                ).astype(np.float32)
            except Exception as e:
                gr.Info(f"默认参数读取失败: {e}，尝试魔↗术↘技↘巧")
                info = rawread(file_path)
                state.raw_data = info["raw"].astype(np.float32)
                state.p.update({
                    "h": info["h"], "w": info["w"],
                    "bl": info["bl"], "wp": info["wp"],
                    "scale": info["wp"] - info["bl"],
                })
                gr.Success("基于魔↗术↘技↘巧，参数已更新")

        elif ext == "mat":
            with h5py.File(file_path, "r") as f:
                state.raw_data = np.array(f["x"]).astype(np.float32) * state.p["scale"] + state.p["bl"]
            state.p.update({
                "h": state.raw_data.shape[0], "w": state.raw_data.shape[1],
            })

        else:
            raise gr.Error("不支持的格式")

        if state.p.get("clip"):
            state.raw_data = state.raw_data.clip(state.p["bl"], state.p["wp"])

        preview = _generate_preview(state)
        p = state.p
        return preview, p["h"], p["w"], p["bl"], p["wp"], state

    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"图像处理失败: {e}")


def update_image(bl, wp, ratio, ispgain, state: UserState):
    """仅更新渲染参数，重新生成预览，不重新读取文件。"""
    if state.raw_data is None:
        raise gr.Error("请先加载图像")
    gr.Info("更新图像参数...")
    for k, v in [("bl", bl), ("wp", wp), ("ratio", ratio), ("ispgain", ispgain)]:
        state.update_param(k, v)
    state.denoised_data = None
    state.mask_data     = None
    preview = _generate_preview(state)
    return preview, state


def estimate_noise(double_est: bool, ransac: bool, patch_size: int, state: UserState):
    """噪声估计，double_est=True 时先去噪再精估。"""
    if not model_service.is_loaded:
        raise gr.Error("请先加载模型")
    if state.raw_data is None:
        raise gr.Error("请先加载图像")

    gr.Info("正在估计噪声...")
    p = state.p
    p["ransac"] = ransac

    processed = (state.raw_data - p["bl"]) / p["scale"]
    lr_raw    = bayer2rggb(processed) * p["ratio"]

    # 粗估计
    reg, state.mask_data = SimpleNLF(
        rggb2bayer(lr_raw), k=19, eps=1e-3,
        setting={"mode": "self", "thr_mode": "score2", "ransac": p["ransac"]},
    )
    p["gain"]  = reg[0] * p["scale"]
    p["sigma"] = float(np.sqrt(max(reg[1], 0))) * p["scale"]

    if double_est:
        log("使用精估计")
        if state.denoised_npy is None:
            log("先去噪再估计")
            state.denoised_data, state.denoised_npy = _vst_denoise_pipeline(lr_raw, state, patch_size)

        reg, state.mask_data = SimpleNLF(
            rggb2bayer(lr_raw), state.denoised_npy, k=13,
            setting={"mode": "collab", "thr_mode": "score3", "ransac": p["ransac"]},
        )
        p["gain"]  = reg[0] * p["scale"]
        p["sigma"] = float(np.sqrt(max(reg[1], 0))) * p["scale"]

    mask_img = _visualize_mask(state)
    gain_out  = round(p["gain"],  2)
    sigma_out = round(p["sigma"], 2)
    log(f"噪声估计完成: gain={gain_out}, sigma={sigma_out}")
    gr.Success(f"噪声估计完成: gain={gain_out:.2f}, sigma={sigma_out:.2f}")
    return mask_img, gain_out, sigma_out, state


def enhance_image(gain, sigma, sigsnr, ddim_mode, patch_size, state: UserState):
    """图像去噪增强。"""
    if not model_service.is_loaded:
        raise gr.Error("请先加载模型")
    if state.raw_data is None:
        raise gr.Error("请先加载图像")

    gr.Info("正在增强图像...")
    p = state.p
    p["ddim_mode"] = ddim_mode
    for k, v in [("gain", gain), ("sigma", sigma), ("sigsnr", sigsnr)]:
        state.update_param(k, v)

    processed = (state.raw_data - p["bl"]) / p["scale"]
    lr_raw    = bayer2rggb(processed) * p["ratio"]

    state.denoised_data, state.denoised_npy = _vst_denoise_pipeline(lr_raw, state, patch_size)

    result = _generate_result(state)
    gr.Success("图像增强完成")
    return result, state


def save_result_npy(state: UserState) -> str:
    if state.denoised_npy is None:
        raise gr.Error("请先进行图像增强")
    with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
        np.save(f.name, state.denoised_npy.astype(np.float32))
        return f.name


def save_result_png(state: UserState) -> str:
    if state.denoised_rgb is None:
        raise gr.Error("请先进行图像增强")
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
        cv2.imwrite(f.name, state.denoised_rgb[:, :, ::-1])
        return f.name


def save_input_npy(state: UserState) -> str:
    if state.raw_data is None:
        raise gr.Error("请先加载图像")
    with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
        np.save(f.name, state.raw_data.astype(np.float32))
        return f.name


def save_input_png(state: UserState) -> str:
    if state.raw_data is None:
        raise gr.Error("请先加载图像")
    preview = _generate_preview(state)
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
        cv2.imwrite(f.name, preview[:, :, ::-1])
        return f.name


def save_mask_npy(state: UserState) -> str:
    if state.mask_data is None:
        raise gr.Error("请先进行噪声估计")
    with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as f:
        np.save(f.name, state.mask_data.astype(np.float32))
        return f.name


def save_mask_png(state: UserState) -> str:
    if state.mask_data is None:
        raise gr.Error("请先进行噪声估计")
    mask_img = _visualize_mask(state)
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
        cv2.imwrite(f.name, mask_img[:, :, ::-1])
        return f.name


def download_zip(selected_items: list, selected_formats: list, state: UserState) -> str:
    """
    打包选中的图像/掩模为 ZIP。
    selected_items:  ["Input (Noisy)", "Output (Denoised)", "Mask"]
    selected_formats: ["NPY", "PNG"]
    """
    if not selected_items:
        raise gr.Error("请至少选择一项下载内容")
    if not selected_formats:
        raise gr.Error("请至少选择一种格式")

    files_to_zip = []
    label_map = {
        "Input (Noisy)": "input",
        "Output (Denoised)": "output",
        "Mask": "mask",
    }

    for item in selected_items:
        base = label_map[item]
        if item == "Input (Noisy)":
            if "NPY" in selected_formats:
                files_to_zip.append((save_input_npy(state), f"{base}.npy"))
            if "PNG" in selected_formats:
                files_to_zip.append((save_input_png(state), f"{base}.png"))
        elif item == "Output (Denoised)":
            if "NPY" in selected_formats:
                files_to_zip.append((save_result_npy(state), f"{base}.npy"))
            if "PNG" in selected_formats:
                files_to_zip.append((save_result_png(state), f"{base}.png"))
        elif item == "Mask":
            if "NPY" in selected_formats:
                files_to_zip.append((save_mask_npy(state), f"{base}.npy"))
            if "PNG" in selected_formats:
                files_to_zip.append((save_mask_png(state), f"{base}.png"))

    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f:
        zip_path = f.name
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for src_path, arcname in files_to_zip:
            zf.write(src_path, arcname)

    return zip_path


def release_user(state: UserState):
    """用户主动或超时释放时调用，归还并发槽并清空所有状态。"""
    if state.has_data:
        ModelService.decrement_users()
    state.clear_images()
    # Release GPU 时一并清空手动覆盖标记，下次为新会话
    state.manual_wb  = False
    state.manual_ccm = False
    gr.Info(f"资源已释放（当前用户数: {ModelService.current_users()}/{MAX_CONCURRENT_USERS}）")
    return state


# ══════════════════════════════════════════════
#  YOND_anytest / YONDParser（保持不变）
# ══════════════════════════════════════════════
class YONDParser:
    def __init__(self, yaml_path="runfiles/Gaussian/gru32n_paper_noclip.yml"):
        self.runfile = yaml_path
        self.mode    = "eval"
        self.debug   = False
        self.nofig   = False
        self.nohost  = False
        self.gpu     = 0


class YOND_anytest:
    def __init__(self, yaml_path, device):
        self.device = device
        self.parser = YONDParser(yaml_path)
        self._init()

    def _init(self):
        with open(self.parser.runfile, "r", encoding="utf-8") as f:
            self.args = yaml.load(f.read(), Loader=yaml.FullLoader)
        self.mode = self.args["mode"] if self.parser.mode is None else self.parser.mode
        if self.parser.debug:
            self.args["num_workers"] = 0
            warnings.warn("Debug 模式：仅使用主进程")
        if "clip" not in self.args["dst"]:
            self.args["dst"]["clip"] = False
        self.save_plot  = not self.parser.nofig
        self.args["dst"]["mode"] = self.mode
        self.hostname, self.hostpath, self.multi_gpu = get_host_with_dir()
        self.yond_dir   = self.args["checkpoint"]
        if not self.parser.nohost:
            for key in self.args:
                if "dst" in key:
                    self.args[key]["root_dir"] = f"{self.hostpath}/{self.args[key]['root_dir']}"
        self.dst        = self.args["dst"]
        self.arch       = self.args["arch"]
        self.pipe       = self.args["pipeline"]
        if self.pipe["bias_corr"] == "none":
            self.pipe["bias_corr"] = None
        self.yond_name  = self.args["model_name"]
        self.method_name = self.args["method_name"]
        self.fast_ckpt  = self.args["fast_ckpt"]
        self.sample_dir = os.path.join(self.args["result_dir"], self.method_name)
        os.makedirs(self.sample_dir, exist_ok=True)
        os.makedirs("./logs", exist_ok=True)

    def load_model(self, model_path):
        self.net = globals()[self.arch["name"]](self.arch)
        model    = torch.load(model_path, map_location="cpu")
        self.net = load_weights(self.net, model, by_name=False)