"""sample_user.py 的 CPU 单测：首帧尺寸规则、ids 读取 / 严格校验 / 源分辨率重采样、prompts / actions、条件占比、clip 素材导出。 PYTHONPATH=code/gc25d:code/arope .venv/bin/python -m pytest -q code/gc25d/tests/test_sample_user.py 最后一个用例导出一条真实 val clip（数据不在、或数据根的 res ≠ GC_RES 就跳过）。形状取自 geometry（GC_RES）。 """ from __future__ import annotations import json import os import types import cv2 import numpy as np import pytest import geometry as GEO import infer as INF import model as GM import sample as S import sample_user as U import text as TX CLIP = "08_Gothic_Cathedral_Quarter__P9C_default_s8946__003552" H, W = GEO.DST_H, GEO.DST_W SRC_MODE = "as_is" if GEO.IDENTITY else "resample_rgb" # 1280×704：源 = 训练分辨率 ⇒ 原样收 def _ids(shape=(81, H, W), seed=0, values=(0, 1, 3, 5, 10, 12, 16, 11, 21, 255)): rng = np.random.default_rng(seed) return np.asarray(values, np.uint8)[rng.integers(0, len(values), shape)] def _png(path, rgb): assert cv2.imwrite(str(path), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) return str(path) # ---------------------------------------------------------------- 首帧 def test_first_frame_sizes(tmp_path): rng = np.random.default_rng(1) a = rng.integers(0, 256, (H, W, 3), dtype=np.uint8) img, info = U.load_first_frame(_png(tmp_path / "a.png", a)) assert info["resize"] == "as_is" and np.array_equal(img, a) src = rng.integers(0, 256, (704, 1280, 3), dtype=np.uint8) img, info = U.load_first_frame(_png(tmp_path / "s.png", src)) assert info["resize"] == SRC_MODE and np.array_equal(img, GEO.resample_rgb(src)) img, info = U.load_first_frame(_png(tmp_path / "o.png", rng.integers(0, 256, (600, 1000, 3), dtype=np.uint8))) assert info["resize"] == "center_crop_area" and img.shape == (H, W, 3) and info["src_size"] == [1000, 600] # ---------------------------------------------------------------- ids def test_ids_npz_npy_keys(tmp_path): ids = _ids() np.savez_compressed(tmp_path / "k.npz", ids=ids, other=np.zeros(3)) np.savez_compressed(tmp_path / "one.npz", whatever=ids) np.savez_compressed(tmp_path / "two.npz", a=ids, b=ids) np.save(tmp_path / "x.npy", ids) for f in ("k.npz", "one.npz", "x.npy"): out, info = U.load_ids(str(tmp_path / f), "parts") assert np.array_equal(out, ids) and not info["resampled"] with pytest.raises(ValueError, match="不止一个"): U.load_ids(str(tmp_path / "two.npz"), "parts") def test_ids_strict(tmp_path): ids = _ids() np.save(tmp_path / "f.npy", ids.astype(np.int32)) with pytest.raises(ValueError, match="uint8"): U.load_ids(str(tmp_path / "f.npy"), "parts") np.save(tmp_path / "s.npy", ids[:80]) with pytest.raises(ValueError, match=rf"\[81,{H},{W}\]"): U.load_ids(str(tmp_path / "s.npy"), "parts") for bad in (7, 9, 22, 254): b = ids.copy() b[40, 100, 200:203] = bad np.save(tmp_path / "b.npy", b) with pytest.raises(ValueError, match=rf"{bad}: 3"): U.load_ids(str(tmp_path / "b.npy"), "parts") def test_ids_source_res_and_ids_vocab(tmp_path): src = _ids((81, 704, 1280), seed=3) np.savez(tmp_path / "src.npz", ids=src) out, info = U.load_ids(str(tmp_path / "src.npz"), "parts") assert info["resampled"] == (not GEO.IDENTITY) and info["src_shape"] == [81, 704, 1280] assert np.array_equal(out, GEO.resample_ids(src)) # ids 词表的 ckpt：parts 先折叠（与 clip 模式 sample.model_ids 相同），结果 ⊆ ids 词表 out_ids, _ = U.load_ids(str(tmp_path / "src.npz"), "ids") assert np.array_equal(out_ids, GM.collapse_parts(GEO.resample_ids(src))) assert set(np.unique(out_ids).tolist()) <= GM.valid_raw_ids("ids") assert np.array_equal(S.model_ids(out_ids, "ids", "cond"), out_ids) assert (S.model_ids(out, "parts", "unknown") == 255).all() # ---------------------------------------------------------------- prompts def _tt(): return types.SimpleNamespace(table={p: None for p in TX.all_prompts() + [TX.NULL_PROMPT]}) def test_prompts_json_and_actions(tmp_path): p = TX.cell_prompts(INF.parse_actions("D:10,A:10")) a = types.SimpleNamespace(prompts=None, actions="D:10,A:10") prompts, blocks, info = U.load_prompts(a, _tt()) assert prompts == p and blocks.shape == (21, 5) and info["source"] == "actions" assert all("moving right" in x for x in prompts[:11]) and all("moving left" in x for x in prompts[11:]) (tmp_path / "p.json").write_text(json.dumps(p)) prompts2, blocks2, info2 = U.load_prompts(types.SimpleNamespace(prompts=str(tmp_path / "p.json"), actions=None), _tt()) assert prompts2 == p and blocks2 is None and info2["source"] == "prompts_json" (tmp_path / "short.json").write_text(json.dumps(p[:20])) with pytest.raises(ValueError, match="21"): U.load_prompts(types.SimpleNamespace(prompts=str(tmp_path / "short.json"), actions=None), _tt()) (tmp_path / "bad.json").write_text(json.dumps(p[:20] + ["a cat"])) with pytest.raises(KeyError): U.load_prompts(types.SimpleNamespace(prompts=str(tmp_path / "bad.json"), actions=None), _tt()) with pytest.raises(ValueError): U.load_prompts(types.SimpleNamespace(prompts=None, actions="D:10,A:9"), _tt()) def test_cond_stats(): ids = _ids((5, 48, 64), seed=4) st = U.cond_stats(ids) assert abs(sum(st["class_frac"].values()) - 1) < 1e-9 and abs(sum(st["class_frac_frame0"].values()) - 1) < 1e-9 assert set(st["class_frac"]) <= set(U.CLASS_NAMES.values()) assert abs(st["unknown_frac_per_frame"]["mean"] - (ids == 255).mean()) < 1e-9 # ---------------------------------------------------------------- 真实 clip 导出 @pytest.mark.skipif(not (os.path.isfile(os.path.join(S.DATA_ROOT, "cond", CLIP + ".npz")) and GEO.root_res(S.DATA_ROOT) == GEO.RES_TAG), reason="数据不在（或 res ≠ GC_RES）") def test_export_clip_roundtrip(tmp_path): p = U.main(["--export_clip", CLIP, "--out", str(tmp_path)]) c = S.load_clip(S.DATA_ROOT, CLIP) ids, _ = U.load_ids(p["ids.npz"], "parts") assert np.array_equal(ids, c["ids"]) f0, i0 = U.load_first_frame(p["frame0.png"]) fs, i1 = U.load_first_frame(p["frame0_src.png"]) assert i0["resize"] == "as_is" and i1["resize"] == SRC_MODE and np.array_equal(f0, fs) with open(p["prompts.json"]) as fh: assert json.load(fh) == c["prompts"] assert TX.cell_prompts(INF.parse_actions(p["actions.npy"])) == c["prompts"]