| """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" |
|
|
|
|
| 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] |
|
|
|
|
| |
| 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)) |
| |
| 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() |
|
|
|
|
| |
| 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 |
|
|
|
|
| |
| @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"] |
|
|
