File size: 7,571 Bytes
68e2686 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | #!/usr/bin/env python3
"""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)
|