maliang / code /gc25d /geometry.py
teawhite's picture
Add files using upload-large-folder tool
bcdd8fb verified
Raw History Blame Contribute Delete
18.8 kB
"""gc25d 几何:源 1280×704 → 训练分辨率(默认 832×480,SPEC §1;env GC_RES="<W>x<H>" 可选,如 1280x704),各组件共用,
纯 numpy(resample_rgb / 画 marker 用 cv2)。分辨率的唯一出处:其它模块的 H / W / latent / 池化格形状一律从这里取。
分辨率(GC_RES,import 时解析一次;默认 "832x480" = 原行为,逐位相同)
------------------------------------------------------------------
W、H 是 32 的倍数(latent 16× ⇒ token 网格 H/32 × W/32 为整数),384 ≤ H ≤ 704、W ≤ 1280,宽高比不超过源(只做水平居中
裁剪)。H ≥ 384 ⇔ 缩小倍数 < 2:中心对齐最近邻的 _IDX 表才盖得到首末行 / 列(更小的 H import 时的断言会炸,parse_res 先拒)。
CROP_W = min(1280, 2·round(352·W/H))(= 704·W/H 取最近的偶数 ⇒ 裁剪恰好居中),CROP_X0 = (1280 − CROP_W)//2:
832×480 ⇒ 裁 x ∈ [30, 1250)(1220 px)再缩放;
1280×704 ⇒ 不裁不缩(= 原始帧 / 原始 ids.png,_IDX_X / _IDX_Y 是恒等映射,resample_* 走恒等快路径 IDENTITY)。
RES_TAG = "<W>x<H>":写进 index/prep_args.json、index 行、ckpt metadata 的 "res";缺失 ⇒ LEGACY_RES = "832x480"。
LAT_H, LAT_W = H/16, W/16(latent 网格);POOL_H, POOL_W = H/4, W/4(gc_sem_cond / UNKNOWN 遮挡的 4×4 池化格)。
子进程(accelerate / DataLoader worker / spawn)继承环境变量 ⇒ 同一次运行里各进程分辨率一致。
坐标约定(与 SPEC §1 逐字一致;括号里是 832×480 的数):
像素中心在整数坐标上,左上原点,x 向右、y 向下。
裁剪源列 x ∈ [CROP_X0, CROP_X0 + CROP_W)([30, 1250),1220 px)、全部 704 行,再缩放到 W×H:
x_dst = (x_src − CROP_X0 + 0.5)·W/CROP_W − 0.5 ((x_src − 30 + 0.5)·832/1220 − 0.5)
y_dst = (y_src + 0.5)·H/704 − 0.5 ((y_src + 0.5)·480/704 − 0.5)
玩家恒在源 (639.5, 351.5) ⇒ dst = 画幅几何中心(832×480:(415.5, 239.5);1280×704:(639.5, 351.5))。
sx = W/CROP_W,sy = H/704(0.681967 / 0.681818,0.02 % 的长宽比差异忽略)。
投影(common.py::proj 同一实现,本数据上验证到 0.0002 px):world_to_clip 是行主序、**行向量**右乘的 4×4,
c = [X, Y, Z, 1] @ M,ndc = c[:3]/c[3],x_src = (ndc_x+1)/2·1280 − 0.5,y_src = (1−ndc_y)/2·704 − 0.5。
ids 用「中心对齐最近邻」(不是 cv2.INTER_NEAREST,后者有半像素偏):
src_x = CROP_X0 + floor((x_dst + 0.5)·CROP_W/W),src_y = floor((y_dst + 0.5)·704/H)
这正是 dst 像素中心反投到源后四舍五入的那个源像素;用整数算术求 floor,避免浮点在 .5 边界上抖。
RGB 用 cv2.INTER_AREA:非整数倍缩小时按面积覆盖加权,同样是中心对齐的(QA 实测 RGB 边缘与 id 边界峰值偏移 0)。
pose-free marker(SPEC §3.2):实心椭圆,中心 = actor position_cm 的投影(dst px),
a = R·ppc,b = (R·|sin(pitch)| + Hh·cos(pitch))·ppc,ppc = (1280/ortho_width_cm)·sx
即世界里半径 R、半高 Hh 的竖直圆柱在正交俯视下的屏幕外接椭圆。用 cv2.ellipse(LINE_8) 的定点(shift)坐标画,
中心可以落在半像素上(玩家恰在画幅中心),画出来的像素集关于该中心对称。
「椭圆与画幅相交」在 cond 里落实为 marker_drawn:几何相交且 draw_marker 真的画出 ≥ 1 个像素(只擦到画幅外缘半个像素、
一个像素中心都没盖住的椭圆算不可见;修复前全量有 150 帧 boss 这样被标成可见却没有 11 号像素)。
"""
from __future__ import annotations
import json
import math
import os
import struct
import cv2
import numpy as np
SRC_W, SRC_H = 1280, 704
RES_ENV = "GC_RES"
DEFAULT_RES = LEGACY_RES = "832x480" # 默认分辨率;也是没写 "res" 的旧数据根 / 旧 ckpt 的分辨率
MIN_H = SRC_H // 2 + 32 # 384:H 的下限(缩小倍数 < 2,见 parse_res)
def crop_w_for(w: int, h: int) -> int:
"""W×H 的水平裁剪宽度(源 px)= 704·W/H 取最近的偶数(≤ 1280):偶数 ⇒ 裁剪关于源中心 x = 639.5 对称(1152×640:1268 而非 1267)。"""
return min(SRC_W, 2 * round(SRC_H * w / (2 * h)))
def parse_res(s: str) -> tuple[int, int]:
""""<W>x<H>" → (W, H)。W、H 为 32 的正倍数,384 ≤ H ≤ 704、W ≤ 1280,宽高比 ≤ 源的 1280:704(只做水平裁剪)。"""
try:
w, h = (int(v) for v in str(s).strip().lower().split("x"))
except ValueError:
raise ValueError(f"{RES_ENV}={s!r}:应形如 '832x480' / '1280x704'") from None
if w <= 0 or h <= 0 or w % 32 or h % 32 or h > SRC_H or w > SRC_W:
raise ValueError(f"{RES_ENV}={s!r}:W、H 须为 32 的正倍数且 H ≤ {SRC_H}、W ≤ {SRC_W}")
if w * SRC_H > h * SRC_W:
raise ValueError(f"{RES_ENV}={s!r}:宽高比 {w / h:.4f} 超过源 {SRC_W}:{SRC_H}(需要竖直裁剪,未实现)")
if 2 * h <= SRC_H or 2 * w <= crop_w_for(w, h): # 缩小倍数 ≥ 2 ⇒ _IDX 取不到首末行 / 列
raise ValueError(f"{RES_ENV}={s!r}:H 须 ≥ {MIN_H}(缩小倍数 < 2,中心对齐最近邻才盖得到源的首末行 / 列)")
return w, h
DST_W, DST_H = parse_res(os.environ.get(RES_ENV) or DEFAULT_RES)
RES_TAG = f"{DST_W}x{DST_H}"
CROP_W = crop_w_for(DST_W, DST_H) # 832×480 ⇒ 1220,1280×704 ⇒ 1280
CROP_X0 = (SRC_W - CROP_W) // 2 # 832×480 ⇒ 30,1280×704 ⇒ 0(CROP_W 偶数 ⇒ 恰好居中)
SX = DST_W / CROP_W # 832×480:0.681967…
SY = DST_H / SRC_H # 832×480:0.681818…
SCALE = np.array([SX, SY], np.float64)
IDENTITY = CROP_X0 == 0 and CROP_W == SRC_W and (DST_W, DST_H) == (SRC_W, SRC_H) # 1280×704:不裁不缩
LAT_H, LAT_W = DST_H // 16, DST_W // 16 # latent 网格(832×480 ⇒ 30×52,1280×704 ⇒ 44×80)
POOL_H, POOL_W = DST_H // 4, DST_W // 4 # 4×4 池化格(832×480 ⇒ 120×208,1280×704 ⇒ 176×320)
TOK_PER_CELL = (LAT_H // 2) * (LAT_W // 2) # DiT patch (1,2,2) 之后每个 latent cell 的 token 数(390 / 880)
AREA_VS_480 = (DST_W * DST_H) / (832 * 480) # 按像素数定的阈值随分辨率缩放(832×480 ⇒ 1.0)
PLAYER_SRC = (639.5, 351.5)
# = src_to_dst(PLAYER_SRC),先乘后除(整数比例下精确:832×480 ⇒ (415.5, 239.5),与旧常数逐位相同)
PLAYER_DST = ((PLAYER_SRC[0] - CROP_X0 + 0.5) * DST_W / CROP_W - 0.5, (PLAYER_SRC[1] + 0.5) * DST_H / SRC_H - 0.5)
PLAYER_ID, BOSS_ID, UNKNOWN_ID = 10, 11, 255
STATIC_IDS = (0, 1, 2, 3, 4, 5, 6)
RAW_IDS = STATIC_IDS + (PLAYER_ID, BOSS_ID)
# 中心对齐最近邻的源索引表(整数算术:floor((2i+1)·CROP_W / (2·W));1280×704 ⇒ 恒等)
_IDX_X = CROP_X0 + ((2 * np.arange(DST_W, dtype=np.int64) + 1) * CROP_W) // (2 * DST_W)
_IDX_Y = ((2 * np.arange(DST_H, dtype=np.int64) + 1) * SRC_H) // (2 * DST_H)
assert _IDX_X.min() == CROP_X0 and _IDX_X.max() == CROP_X0 + CROP_W - 1, (_IDX_X.min(), _IDX_X.max())
assert _IDX_Y.min() == 0 and _IDX_Y.max() == SRC_H - 1, (_IDX_Y.min(), _IDX_Y.max())
# marker 定点坐标的小数位数(cv2 上限 16;8 位 = 1/256 px)
MARKER_SHIFT = 8
# ------------------------------------------------------------------ 分辨率一致性闸 ----
def root_res(root: str) -> str:
"""数据根的分辨率 = <root>/index/prep_args.json 的 "res";没有这个文件 / 键(旧数据根)⇒ LEGACY_RES。"""
try:
with open(os.path.join(root, "index", "prep_args.json"), encoding="utf-8") as fh:
return str(json.load(fh).get("res") or LEGACY_RES)
except FileNotFoundError:
return LEGACY_RES
def check_root_res(root: str, what: str = "data_root") -> str:
"""数据根的 res ≠ 当前 RES_TAG ⇒ RuntimeError(点名 GC_RES)。返回 res。"""
r = root_res(root)
if r != RES_TAG:
raise RuntimeError(f"{what} {root} 是 {r} 的数据(index/prep_args.json 的 res,缺失 = {LEGACY_RES}),"
f"本进程 {RES_ENV}={RES_TAG}:设 {RES_ENV}={r},或换成 {RES_TAG} 的数据根")
return r
def check_run_dir_res(run_dir: str, what: str = "--output") -> str | None:
"""蒸馏 run 目录里已有的 config.json 的 "res"(缺失 = LEGACY_RES)≠ RES_TAG ⇒ RuntimeError;没有 config.json(新目录)⇒ None。
runner 在打开 train.log 之前调:别往另一个分辨率的 run 里追加日志 / 产物。"""
p = os.path.join(run_dir, "config.json")
if not os.path.isfile(p):
return None
with open(p, encoding="utf-8") as fh:
r = str(json.load(fh).get("res") or LEGACY_RES)
if r != RES_TAG:
raise RuntimeError(f"{what} {run_dir} 是 {r} 的 run 目录(config.json 的 res,缺失 = {LEGACY_RES}),"
f"本进程 {RES_ENV}={RES_TAG}:换一个 {RES_TAG} 的输出目录,或设 {RES_ENV}={r}")
return r
def ckpt_res(path: str) -> str:
"""只读 safetensors 头的 metadata "res"(不碰数据);缺失(旧 ckpt)⇒ LEGACY_RES。"""
with open(path, "rb") as fh:
hdr = json.loads(fh.read(struct.unpack("<Q", fh.read(8))[0]))
return str((hdr.get("__metadata__") or {}).get("res") or LEGACY_RES)
def check_ckpt_res(path: str, what: str = "ckpt") -> str:
"""ckpt 记录的 res ≠ 当前 RES_TAG ⇒ RuntimeError(点名 GC_RES)。返回 res。"""
r = ckpt_res(path)
if r != RES_TAG:
raise RuntimeError(f"{what} {path} 是在 {r} 上训练 / 导出的(metadata res,缺失 = {LEGACY_RES}),"
f"本进程 {RES_ENV}={RES_TAG}:设 {RES_ENV}={r},或换成 {RES_TAG} 的 ckpt")
return r
# ------------------------------------------------------------------ 坐标映射 ----
def src_to_dst(xy) -> np.ndarray:
"""源像素坐标 [..., 2] → dst 像素坐标 [..., 2](float64),像素中心约定。"""
xy = np.asarray(xy, np.float64)
if xy.shape[-1] != 2:
raise ValueError(f"src_to_dst 期望最后一维为 2,得到 {xy.shape}")
out = np.empty_like(xy)
out[..., 0] = (xy[..., 0] - CROP_X0 + 0.5) * SX - 0.5
out[..., 1] = (xy[..., 1] + 0.5) * SY - 0.5
return out
def dst_to_src(xy) -> np.ndarray:
"""dst 像素坐标 [..., 2] → 源像素坐标 [..., 2](src_to_dst 的精确逆)。"""
xy = np.asarray(xy, np.float64)
if xy.shape[-1] != 2:
raise ValueError(f"dst_to_src 期望最后一维为 2,得到 {xy.shape}")
out = np.empty_like(xy)
out[..., 0] = (xy[..., 0] + 0.5) / SX - 0.5 + CROP_X0
out[..., 1] = (xy[..., 1] + 0.5) / SY - 0.5
return out
def proj(M, P) -> np.ndarray:
"""世界点 → 源像素坐标。
M: [4,4] 或 [N,4,4](world_to_clip,行主序、行向量);P: [3] 或 [N,3](cm)。
返回 [2](M、P 都不带 N 时)或 [N,2]。与 analysis/canvas_consistency/scripts/common.py::proj 同一公式。
"""
M = np.asarray(M, np.float64)
P = np.asarray(P, np.float64)
single = M.ndim == 2 and P.ndim == 1
if M.ndim == 2:
M = M[None]
if M.shape[1:] != (4, 4):
raise ValueError(f"proj: M 形状应为 [N,4,4],得到 {M.shape}")
if P.ndim == 1:
if P.shape != (3,):
raise ValueError(f"proj: P 形状应为 [3] 或 [N,3],得到 {P.shape}")
P = np.broadcast_to(P, (M.shape[0], 3))
if M.shape[0] == 1 and P.shape[0] > 1:
M = np.broadcast_to(M, (P.shape[0], 4, 4))
if P.shape != (M.shape[0], 3):
raise ValueError(f"proj: M {M.shape} 与 P {P.shape} 的 N 不一致")
h = np.concatenate([P, np.ones((P.shape[0], 1))], 1)
c = np.einsum("ni,nij->nj", h, M)
ndc = c[:, :3] / c[:, 3:4]
out = np.stack([(ndc[:, 0] + 1) * 0.5 * SRC_W - 0.5, (1 - ndc[:, 1]) * 0.5 * SRC_H - 0.5], 1)
return out[0] if single else out
def static_origin_src(M) -> np.ndarray:
"""org[t] = proj(M[t], (0,0,0)):世界原点的屏幕位置(源 px)。静态内容在帧 t 的 p 处 ⇒ 在帧 u 的 p + org[u] − org[t]。"""
return proj(M, np.zeros(3))
def frame_offset_dst(org_src: np.ndarray) -> np.ndarray:
"""org [T,2](源 px)→ frame_offset_px [T,2](dst px)= (org[t] − org[0])·(sx, sy),静态内容相对帧 0 的位移。"""
org_src = np.asarray(org_src, np.float64)
if org_src.ndim != 2 or org_src.shape[1] != 2:
raise ValueError(f"frame_offset_dst: org 形状应为 [T,2],得到 {org_src.shape}")
return ((org_src - org_src[0]) * SCALE).astype(np.float32)
# ------------------------------------------------------------------ 重采样 ----
def resample_ids(ids_src: np.ndarray) -> np.ndarray:
"""uint8 [..., 704, 1280] → uint8 [..., H, W],中心对齐最近邻(numpy fancy indexing)。
IDENTITY(1280×704)⇒ 直接拷贝一份(与查表结果逐位相同,tests 核过)。"""
ids_src = np.asarray(ids_src)
if ids_src.dtype != np.uint8 or ids_src.shape[-2:] != (SRC_H, SRC_W):
raise ValueError(f"resample_ids 期望 uint8 [..., {SRC_H}, {SRC_W}],得到 {ids_src.dtype} {ids_src.shape}")
if IDENTITY:
return np.array(ids_src, order="C", copy=True)
return np.ascontiguousarray(ids_src[..., _IDX_Y[:, None], _IDX_X[None, :]])
def resample_rgb(rgb_src: np.ndarray) -> np.ndarray:
"""uint8 [704, 1280, 3] → uint8 [H, W, 3]:裁 x∈[CROP_X0, CROP_X0+CROP_W)(832×480:[30,1250))后 cv2.INTER_AREA 缩放。
IDENTITY(1280×704)⇒ 直接拷贝一份(= 同尺寸 INTER_AREA,tests 核过)。"""
rgb_src = np.asarray(rgb_src)
if rgb_src.dtype != np.uint8 or rgb_src.shape != (SRC_H, SRC_W, 3):
raise ValueError(f"resample_rgb 期望 uint8 [{SRC_H}, {SRC_W}, 3],得到 {rgb_src.dtype} {rgb_src.shape}")
if IDENTITY:
return np.array(rgb_src, order="C", copy=True)
crop = np.ascontiguousarray(rgb_src[:, CROP_X0:CROP_X0 + CROP_W])
out = cv2.resize(crop, (DST_W, DST_H), interpolation=cv2.INTER_AREA)
assert out.shape == (DST_H, DST_W, 3) and out.dtype == np.uint8, out.shape
return out
# ------------------------------------------------------------------ 可视化 ----
PALETTE = np.full((256, 3), 255, np.uint8) # 不在词表里的值 = 白(显眼)
PALETTE[0] = (135, 190, 235) # sky / bg
PALETTE[1] = (140, 110, 70) # terrain
PALETTE[2] = (200, 200, 190) # paving
PALETTE[3] = (200, 90, 60) # architecture
PALETTE[4] = (110, 110, 120) # rock
PALETTE[5] = (60, 160, 60) # vegetation
PALETTE[6] = (230, 190, 40) # prop
PALETTE[PLAYER_ID] = (255, 0, 255)
PALETTE[BOSS_ID] = (255, 0, 0)
PALETTE[UNKNOWN_ID] = (0, 0, 0) # 255 另外叠灰色斜纹,见 colorize_ids
HATCH_GREY = np.array([96, 96, 96], np.uint8)
def colorize_ids(ids: np.ndarray) -> np.ndarray:
"""uint8 [..., H, W] → uint8 [..., H, W, 3](固定调色板;255 = 黑底灰色斜纹,周期 8 px)。"""
ids = np.asarray(ids)
if ids.dtype != np.uint8 or ids.ndim < 2:
raise ValueError(f"colorize_ids 期望 uint8 [..., H, W],得到 {ids.dtype} {ids.shape}")
out = PALETTE[ids]
unk = ids == UNKNOWN_ID
if unk.any():
h, w = ids.shape[-2:]
yy, xx = np.mgrid[0:h, 0:w]
stripe = ((xx + yy) // 4) % 2 == 1
out[unk & np.broadcast_to(stripe, ids.shape)] = HATCH_GREY
return out
# ------------------------------------------------------------------ marker ----
def px_per_cm_dst(ortho_width_cm: float) -> float:
"""dst 像素 / 世界 cm:(1280/ortho_width_cm)·sx(832×480:sx = 832/1220)。"""
if not ortho_width_cm > 0:
raise ValueError(f"ortho_width_cm 必须 > 0,得到 {ortho_width_cm}")
return (SRC_W / float(ortho_width_cm)) * SX
def marker_axes(R_cm: float, Hh_cm: float, pitch_deg: float, ortho_width_cm: float) -> tuple[float, float]:
"""(a, b) dst px:a = R·ppc,b = (R·|sin(pitch)| + Hh·cos(pitch))·ppc(pitch 取绝对值,度)。"""
if not (R_cm > 0 and Hh_cm > 0):
raise ValueError(f"marker 尺寸必须为正:R={R_cm}, Hh={Hh_cm}")
p = math.radians(abs(float(pitch_deg)))
ppc = px_per_cm_dst(ortho_width_cm)
return R_cm * ppc, (R_cm * abs(math.sin(p)) + Hh_cm * math.cos(p)) * ppc
def ellipse_intersects_frame(cx: float, cy: float, a: float, b: float) -> bool:
"""轴对齐椭圆与 W×H 画幅(像素区域 [−0.5, W−0.5]×[−0.5, H−0.5])是否相交。
精确判定:取画幅矩形里离椭圆中心最近的点,看它是否落在椭圆内(矩形凸、椭圆凸,这样判是精确的)。
"""
if not (np.isfinite(cx) and np.isfinite(cy)):
return False
qx = min(max(cx, -0.5), DST_W - 0.5)
qy = min(max(cy, -0.5), DST_H - 0.5)
return ((qx - cx) / a) ** 2 + ((qy - cy) / b) ** 2 <= 1.0
def draw_marker(ids_dst: np.ndarray, cx: float, cy: float, a: float, b: float, cls: int) -> np.ndarray:
"""在 uint8 [H, W] 上原地画实心轴对齐椭圆(cv2.ellipse,LINE_8,定点 shift=8),返回同一数组。"""
if ids_dst.dtype != np.uint8 or ids_dst.shape != (DST_H, DST_W) or not ids_dst.flags.c_contiguous:
raise ValueError(f"draw_marker 期望 C 连续 uint8 [{DST_H}, {DST_W}],得到 {ids_dst.dtype} {ids_dst.shape}")
s = 1 << MARKER_SHIFT
centre = (int(round(cx * s)), int(round(cy * s)))
axes = (max(1, int(round(a * s))), max(1, int(round(b * s))))
cv2.ellipse(ids_dst, centre, axes, 0, 0, 360, int(cls), thickness=-1, lineType=cv2.LINE_8, shift=MARKER_SHIFT)
return ids_dst
def marker_mask(cx: float, cy: float, a: float, b: float) -> np.ndarray:
"""单个 marker 的 bool [H, W] 掩码(与 draw_marker 像素级一致,QA / 增广复用)。"""
m = np.zeros((DST_H, DST_W), np.uint8)
draw_marker(m, cx, cy, a, b, 1)
return m.astype(bool)
def marker_drawn(cx: float, cy: float, a: float, b: float) -> bool:
"""draw_marker 是否会在 W×H 上画出 ≥ 1 个像素(= cond 里 visible 标记的定义)。
ellipse_intersects_frame 判的是与像素**区域** [−0.5, W−0.5]×[−0.5, H−0.5] 相交;只擦到最外那半个像素的椭圆
(例如底边在 y = −0.34)一个像素中心都盖不住,cv2.ellipse 什么也不画。所以先做几何相交(快速排除),再用
与 draw_marker 同一套定点光栅化的掩码确认真的有像素——visible 与 cond 里的像素逐帧一致。
"""
if not ellipse_intersects_frame(cx, cy, a, b):
return False
return bool(marker_mask(cx, cy, a, b).any())