| |
| """teaser · Boss 换装用的视觉打分工具（img env：$M/envs/img/bin/python，GPU）。被 bossswap_f0.py / bossswap_metrics.py import。 |
| |
| * BossSeg：SAM2.1-hiera-large（transformers Sam2Model，本地 HF_HOME snapshot），提示 = 条件 Boss 掩码（膨胀 10 px）的外接框 + |
| 身体内核（非武器部件，腐蚀 3 px）最远点采样 4 个正点 + 武器中线上离身体最远的一点；玩家身体内核 2 个负点。 |
| 输出限制在 Boss 掩码膨胀 40 px 内（防止整块地面被分进去）。iou(seg, cond) = 生成画面里 Boss 轮廓与条件掩码的重合度。 |
| 注意：胶囊掩码只是近似轮廓（f0_realchar 实测真实角色 vs 胶囊 IoU 0.39–0.75），所以要和原版 Boss 同一指标比相对值。 |
| * BossClip：CLIP ViT-L/14，Boss 裁图对 16 种外观 + 原版描述做零样本分类，p(目标) 当「认不认得出是这个 Boss」。 |
| """ |
| from __future__ import annotations |
|
|
| import glob |
| import os |
|
|
| import numpy as np |
|
|
| os.environ.setdefault("HF_HOME", "/opt/dlami/nvme/zhiyangdeng/.cache/huggingface") |
| HUB = f"{os.environ['HF_HOME']}/hub" |
| PLAYER, BOSS = (10, 12, 13, 14, 15, 16), (11, 17, 18, 19, 20, 21) |
| WEAPONS = (16, 21) |
|
|
|
|
| def snap(repo: str) -> str: |
| p = sorted(glob.glob(f"{HUB}/models--{repo.replace('/', '--')}/snapshots/*")) |
| if not p: |
| raise FileNotFoundError(repo) |
| return p[-1] |
|
|
|
|
| def grow(mask: np.ndarray, r: int) -> np.ndarray: |
| import cv2 |
| if r == 0: |
| return mask.copy() |
| k = np.ones((2 * abs(r) + 1, 2 * abs(r) + 1), np.uint8) |
| f = cv2.dilate if r > 0 else cv2.erode |
| return f(mask.astype(np.uint8), k) > 0 |
|
|
|
|
| def fps_points(core: np.ndarray, k: int = 4) -> list[list[float]]: |
| ys, xs = np.nonzero(core) |
| P = np.stack([xs, ys], 1).astype(np.float32) |
| sel = [P[np.argmin(((P - P.mean(0)) ** 2).sum(1))]] |
| for _ in range(k - 1): |
| sel.append(P[np.argmax(np.min([((P - q) ** 2).sum(1) for q in sel], 0))]) |
| return [[float(x), float(y)] for x, y in sel] |
|
|
|
|
| def bbox(mask: np.ndarray, pad: float = 0.0, min_side: int = 1, W: int = 832, H: int = 480): |
| ys, xs = np.nonzero(mask) |
| x0, x1, y0, y1 = xs.min(), xs.max() + 1, ys.min(), ys.max() + 1 |
| w, h = x1 - x0, y1 - y0 |
| cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 |
| w, h = max(min_side, w * (1 + pad)), max(min_side, h * (1 + pad)) |
| return (int(max(0, cx - w / 2)), int(max(0, cy - h / 2)), int(min(W, cx + w / 2)), int(min(H, cy + h / 2))) |
|
|
|
|
| class BossSeg: |
| def __init__(self, device="cuda"): |
| import torch |
| from transformers import Sam2Model, Sam2Processor |
| p = snap("facebook/sam2.1-hiera-large") |
| self.torch, self.dev = torch, device |
| self.proc = Sam2Processor.from_pretrained(p) |
| self.model = Sam2Model.from_pretrained(p, dtype=torch.float32).to(device).eval() |
|
|
| def prompts(self, ids: np.ndarray): |
| mk = np.isin(ids, BOSS) |
| if mk.sum() < 300: |
| return None |
| body = mk & ~np.isin(ids, WEAPONS) |
| core = grow(body, -3) |
| core = core if core.sum() >= 10 else (body if body.sum() >= 10 else mk) |
| box = list(bbox(grow(mk, 10))) |
| pts = fps_points(core) |
| lab = [1] * len(pts) |
| wp = grow(mk & np.isin(ids, WEAPONS), -2) |
| if wp.sum() >= 60: |
| ys, xs = np.nonzero(wp) |
| c = np.array(pts[0]) |
| k = int(np.argmax((xs - c[0]) ** 2 + (ys - c[1]) ** 2)) |
| pts.append([float(xs[k]), float(ys[k])]) |
| lab.append(1) |
| pc = grow(np.isin(ids, PLAYER) & ~np.isin(ids, WEAPONS), -3) & ~grow(mk, 4) |
| if pc.sum() >= 30: |
| for q in fps_points(pc, 2): |
| pts.append(q) |
| lab.append(0) |
| return mk, box, pts, lab |
|
|
| def __call__(self, rgbs: list[np.ndarray], idss: list[np.ndarray]) -> list[np.ndarray | None]: |
| """每帧一个 Boss 掩码（看不到 Boss 的帧 = None）。批处理。""" |
| outs: list = [None] * len(rgbs) |
| todo = [] |
| for i, (rgb, ids) in enumerate(zip(rgbs, idss)): |
| pr = self.prompts(ids) |
| if pr is not None: |
| todo.append((i, rgb, pr)) |
| for s in range(0, len(todo), 8): |
| ch = todo[s:s + 8] |
| inp = self.proc(images=[c[1] for c in ch], input_boxes=[[c[2][1]] for c in ch], |
| input_points=[[c[2][2]] for c in ch], input_labels=[[c[2][3]] for c in ch], |
| return_tensors="pt").to(self.dev) |
| with self.torch.no_grad(): |
| res = self.model(**inp, multimask_output=False) |
| ms = self.proc.post_process_masks(res.pred_masks.cpu(), inp["original_sizes"]) |
| for (i, _, pr), m in zip(ch, ms): |
| outs[i] = (m[0, 0].numpy() > 0) & grow(pr[0], 40) |
| return outs |
|
|
|
|
| def iou(a: np.ndarray, b: np.ndarray) -> float: |
| u = (a | b).sum() |
| return float((a & b).sum() / u) if u else 0.0 |
|
|
|
|
| class BossClip: |
| def __init__(self, labels: dict[str, str], device="cuda"): |
| """labels：键 → 文本描述（含原版 Boss）。""" |
| import torch |
| from transformers import CLIPModel, CLIPProcessor |
| p = snap("openai/clip-vit-large-patch14") |
| self.torch, self.dev = torch, device |
| self.model = CLIPModel.from_pretrained(p, dtype=torch.float16).to(device).eval() |
| self.proc = CLIPProcessor.from_pretrained(p) |
| self.keys = list(labels) |
| txt = [f"a top-down video game screenshot of {labels[k]}" for k in self.keys] |
| with torch.no_grad(): |
| t = self.proc(text=txt, return_tensors="pt", padding=True, truncation=True).to(device) |
| e = self.model.get_text_features(**t) |
| e = e if isinstance(e, torch.Tensor) else e.pooler_output |
| self.temb = torch.nn.functional.normalize(e.float(), dim=-1) |
|
|
| def probs(self, crops: list[np.ndarray]) -> np.ndarray: |
| """crops：RGB uint8 → [N, K] softmax 概率（CLIP logit scale 100）。""" |
| from PIL import Image |
| torch = self.torch |
| out = [] |
| for s in range(0, len(crops), 32): |
| im = [Image.fromarray(c) for c in crops[s:s + 32]] |
| with torch.no_grad(): |
| x = self.proc(images=im, return_tensors="pt").to(self.dev) |
| e = self.model.get_image_features(pixel_values=x["pixel_values"].half()) |
| e = e if isinstance(e, torch.Tensor) else e.pooler_output |
| e = torch.nn.functional.normalize(e.float(), dim=-1) |
| out.append(torch.softmax(100 * e @ self.temb.T, -1).cpu().numpy()) |
| return np.concatenate(out, 0) |
|
|
|
|
| def boss_crop(rgb: np.ndarray, mask: np.ndarray, pad: float = 0.25, min_side: int = 96) -> np.ndarray: |
| x0, y0, x1, y1 = bbox(mask, pad=pad, min_side=min_side) |
| return np.ascontiguousarray(rgb[y0:y1, x0:x1]) |
|
|
|
|
| def lab(img: np.ndarray) -> np.ndarray: |
| import cv2 |
| return cv2.cvtColor(img, cv2.COLOR_RGB2LAB).astype(np.float32) * np.array([100 / 255, 1, 1], np.float32) |
|
|
|
|
| FFMPEG = f"/opt/dlami/nvme/zhiyangdeng/maliang/.venv/lib/python3.12/site-packages/imageio_ffmpeg/binaries/ffmpeg-linux-x86_64-v7.0.2" |
|
|
|
|
| def read_video(path: str, w: int = 832, h: int = 480) -> np.ndarray: |
| """mp4 → RGB uint8 [F,h,w,3]（ffmpeg 解码；尺寸须已知，不依赖 PyAV）。""" |
| import subprocess |
| p = subprocess.run([FFMPEG, "-loglevel", "error", "-i", path, "-f", "rawvideo", "-pix_fmt", "rgb24", "-"], |
| capture_output=True, check=True) |
| return np.frombuffer(p.stdout, np.uint8).reshape(-1, h, w, 3) |
|
|
