maliang / scripts /res_test_edge.py

Add files using upload-large-folder tool

51ac7cc verified 16 days ago

5.31 kB

#!/usr/bin/env python3
"""res_test.py 的补充（分辨率决策的证据）：detail_ratio = |Laplacian| 能量比只看高频「多不多」，不看「在不在对的位置」——
没有重采样低通的 1280×704 会被奖励，哪怕多出来的高频是 VAE 伪影；crop PSNR 又偏爱模糊。这里在同一批 player crop、
同一条 VAE 往返（res_test.enc / dec / to_res 原样复用）上加三个看「对齐」的指标：
  ssim         灰度 SSIM（高斯窗 σ 1.5，标准常数）
  lap_corr     Laplacian(重建) 与 Laplacian(GT) 的 Pearson 相关：高频要在对的位置才高
  edge_f1      Canny(50,150) 边缘，1 px 容差的 F1（precision = 重建边缘落在 GT 边缘旁的比例；recall 反之）
psnr / detail_ratio 也重算一遍，对拍 outputs/char_diag/res_test.json（应相同）。写 outputs/char_diag/res_test_edge.json，
不动 res_test.json。默认 GC_RES（_ref 的 ids 是 832×480）。
    CUDA_VISIBLE_DEVICES=0 $PY scripts/res_test_edge.py
"""
import json
import os
import sys

import cv2
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import res_test as RT  # noqa: E402  （enc / dec / to_res / RES / NF / D 原样复用）

import geometry as GEO  # noqa: E402
import infer as INF  # noqa: E402
import model as GM  # noqa: E402
import sample as S  # noqa: E402

K3 = np.ones((3, 3), np.uint8)


def gray(a):
    return cv2.cvtColor(np.ascontiguousarray(a), cv2.COLOR_RGB2GRAY)


def ssim(a, b):
    a, b = gray(a).astype(np.float64), gray(b).astype(np.float64)
    c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
    blur = lambda x: cv2.GaussianBlur(x, (11, 11), 1.5)  # noqa: E731
    ma, mb = blur(a), blur(b)
    va, vb, cov = blur(a * a) - ma * ma, blur(b * b) - mb * mb, blur(a * b) - ma * mb
    return float((((2 * ma * mb + c1) * (2 * cov + c2)) / ((ma * ma + mb * mb + c1) * (va + vb + c2))).mean())


def lap_corr(a, b):
    la = cv2.Laplacian(gray(a).astype(np.float32), cv2.CV_32F).ravel()
    lb = cv2.Laplacian(gray(b).astype(np.float32), cv2.CV_32F).ravel()
    return float(np.corrcoef(la, lb)[0, 1])


def edge_f1(rc, gt):
    er, eg = cv2.Canny(gray(rc), 50, 150) > 0, cv2.Canny(gray(gt), 50, 150) > 0
    if not er.any() or not eg.any():
        return 0.0, 0.0, 0.0
    dr, dg = cv2.dilate(er.astype(np.uint8), K3) > 0, cv2.dilate(eg.astype(np.uint8), K3) > 0
    p, r = float((er & dg).sum() / er.sum()), float((eg & dr).sum() / eg.sum())
    return p, r, (2 * p * r / (p + r) if p + r > 0 else 0.0)


def main():
    assert GEO.RES_TAG == GEO.LEGACY_RES, "_ref 的 ids 是 832×480：用默认 GC_RES 跑"
    clips = json.load(open(f"{RT.D}/clips.json"))
    vae = INF.load_vae("cuda", S.MODEL_DIR)
    out = {k: [] for k in RT.RES}
    for cid in clips:
        ref = np.load(f"{RT.D}/_ref/{cid}.npz")
        src, ids = ref["gt_src"][:RT.NF], ref["ids"]
        for name, wh in RT.RES.items():
            x, x0, cw = (src, 0, 1280) if name == "1280x704" else RT.to_res(src, wh)
            rec = RT.dec(vae, RT.enc(vae, np.ascontiguousarray(x)))
            for t in range(8, RT.NF, 8):                       # 与 res_test.py 同一套 crop
                m = np.isin(ids[t], GM.PLAYER_PART_IDS)
                if m.sum() < 600:
                    continue
                ys, xs = np.nonzero(m)
                p = GEO.dst_to_src(np.array([[xs.min() - 10, ys.min() - 10], [xs.max() + 10, ys.max() + 10]], float))
                sx0, sy0 = int(max(0, p[0, 0])), int(max(0, p[0, 1]))
                sx1, sy1 = int(min(1280, p[1, 0])), int(min(704, p[1, 1]))
                gt_c = src[t, sy0:sy1, sx0:sx1]
                if name == "1280x704":
                    rc = rec[t, sy0:sy1, sx0:sx1]
                else:
                    up = cv2.resize(rec[t], (cw, 704), interpolation=cv2.INTER_CUBIC)
                    full = np.zeros_like(src[t]); full[:, x0:x0 + cw] = up
                    rc = full[sy0:sy1, sx0:sx1]
                mse = np.mean((rc.astype(np.float32) - gt_c.astype(np.float32)) ** 2)
                g = lambda a: float(np.abs(cv2.Laplacian(gray(a).astype(np.float32), cv2.CV_32F)).mean())  # noqa: E731
                ep, er, ef = edge_f1(rc, gt_c)
                out[name].append(dict(psnr=10 * np.log10(255 ** 2 / max(mse, 1e-6)), detail=g(rc) / max(g(gt_c), 1e-6),
                                      ssim=ssim(rc, gt_c), lap_corr=lap_corr(rc, gt_c), edge_p=ep, edge_r=er, edge_f1=ef))
        print(cid, {k: len(v) for k, v in out.items()}, flush=True)
    res = {k: {"n": len(v), **{m: round(float(np.mean([d[m] for d in v])), 4) for m in v[0]}} for k, v in out.items()}
    base = json.load(open(f"{RT.D}/res_test.json"))["vae_ceiling"]
    for k in res:                                              # 对拍 res_test.json（同一条路径 ⇒ 同数）
        res[k]["matches_res_test_json"] = (abs(res[k]["psnr"] - base[k]["player_crop_psnr_vs_native_gt"]) < 0.01
                                           and abs(res[k]["detail"] - base[k]["detail_ratio_vs_native_gt"]) < 0.001)
    print(json.dumps(res, indent=1), flush=True)
    json.dump({"vae_ceiling_alignment": res, "note": __doc__}, open(f"{RT.D}/res_test_edge.json", "w"), indent=1,
              ensure_ascii=False)


if __name__ == "__main__":
    sys.exit(main())

	#!/usr/bin/env python3
	"""res_test.py 的补充（分辨率决策的证据）：detail_ratio = \|Laplacian\| 能量比只看高频「多不多」，不看「在不在对的位置」——
	没有重采样低通的 1280×704 会被奖励，哪怕多出来的高频是 VAE 伪影；crop PSNR 又偏爱模糊。这里在同一批 player crop、
	同一条 VAE 往返（res_test.enc / dec / to_res 原样复用）上加三个看「对齐」的指标：
	ssim 灰度 SSIM（高斯窗 σ 1.5，标准常数）
	lap_corr Laplacian(重建) 与 Laplacian(GT) 的 Pearson 相关：高频要在对的位置才高
	edge_f1 Canny(50,150) 边缘，1 px 容差的 F1（precision = 重建边缘落在 GT 边缘旁的比例；recall 反之）
	psnr / detail_ratio 也重算一遍，对拍 outputs/char_diag/res_test.json（应相同）。写 outputs/char_diag/res_test_edge.json，
	不动 res_test.json。默认 GC_RES（_ref 的 ids 是 832×480）。

	CUDA_VISIBLE_DEVICES=0 $PY scripts/res_test_edge.py
	"""
	import json
	import os
	import sys

	import cv2
	import numpy as np

	sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	import res_test as RT # noqa: E402 （enc / dec / to_res / RES / NF / D 原样复用）

	import geometry as GEO # noqa: E402
	import infer as INF # noqa: E402
	import model as GM # noqa: E402
	import sample as S # noqa: E402

	K3 = np.ones((3, 3), np.uint8)


	def gray(a):
	return cv2.cvtColor(np.ascontiguousarray(a), cv2.COLOR_RGB2GRAY)


	def ssim(a, b):
	a, b = gray(a).astype(np.float64), gray(b).astype(np.float64)
	c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
	blur = lambda x: cv2.GaussianBlur(x, (11, 11), 1.5) # noqa: E731
	ma, mb = blur(a), blur(b)
	va, vb, cov = blur(a * a) - ma * ma, blur(b * b) - mb * mb, blur(a * b) - ma * mb
	return float((((2 * ma * mb + c1) * (2 * cov + c2)) / ((ma * ma + mb * mb + c1) * (va + vb + c2))).mean())


	def lap_corr(a, b):
	la = cv2.Laplacian(gray(a).astype(np.float32), cv2.CV_32F).ravel()
	lb = cv2.Laplacian(gray(b).astype(np.float32), cv2.CV_32F).ravel()
	return float(np.corrcoef(la, lb)[0, 1])


	def edge_f1(rc, gt):
	er, eg = cv2.Canny(gray(rc), 50, 150) > 0, cv2.Canny(gray(gt), 50, 150) > 0
	if not er.any() or not eg.any():
	return 0.0, 0.0, 0.0
	dr, dg = cv2.dilate(er.astype(np.uint8), K3) > 0, cv2.dilate(eg.astype(np.uint8), K3) > 0
	p, r = float((er & dg).sum() / er.sum()), float((eg & dr).sum() / eg.sum())
	return p, r, (2 * p * r / (p + r) if p + r > 0 else 0.0)


	def main():
	assert GEO.RES_TAG == GEO.LEGACY_RES, "_ref 的 ids 是 832×480：用默认 GC_RES 跑"
	clips = json.load(open(f"{RT.D}/clips.json"))
	vae = INF.load_vae("cuda", S.MODEL_DIR)
	out = {k: [] for k in RT.RES}
	for cid in clips:
	ref = np.load(f"{RT.D}/_ref/{cid}.npz")
	src, ids = ref["gt_src"][:RT.NF], ref["ids"]
	for name, wh in RT.RES.items():
	x, x0, cw = (src, 0, 1280) if name == "1280x704" else RT.to_res(src, wh)
	rec = RT.dec(vae, RT.enc(vae, np.ascontiguousarray(x)))
	for t in range(8, RT.NF, 8): # 与 res_test.py 同一套 crop
	m = np.isin(ids[t], GM.PLAYER_PART_IDS)
	if m.sum() < 600:
	continue
	ys, xs = np.nonzero(m)
	p = GEO.dst_to_src(np.array([[xs.min() - 10, ys.min() - 10], [xs.max() + 10, ys.max() + 10]], float))
	sx0, sy0 = int(max(0, p[0, 0])), int(max(0, p[0, 1]))
	sx1, sy1 = int(min(1280, p[1, 0])), int(min(704, p[1, 1]))
	gt_c = src[t, sy0:sy1, sx0:sx1]
	if name == "1280x704":
	rc = rec[t, sy0:sy1, sx0:sx1]
	else:
	up = cv2.resize(rec[t], (cw, 704), interpolation=cv2.INTER_CUBIC)
	full = np.zeros_like(src[t]); full[:, x0:x0 + cw] = up
	rc = full[sy0:sy1, sx0:sx1]
	mse = np.mean((rc.astype(np.float32) - gt_c.astype(np.float32)) ** 2)
	g = lambda a: float(np.abs(cv2.Laplacian(gray(a).astype(np.float32), cv2.CV_32F)).mean()) # noqa: E731
	ep, er, ef = edge_f1(rc, gt_c)
	out[name].append(dict(psnr=10 * np.log10(255 ** 2 / max(mse, 1e-6)), detail=g(rc) / max(g(gt_c), 1e-6),
	ssim=ssim(rc, gt_c), lap_corr=lap_corr(rc, gt_c), edge_p=ep, edge_r=er, edge_f1=ef))
	print(cid, {k: len(v) for k, v in out.items()}, flush=True)
	res = {k: {"n": len(v), **{m: round(float(np.mean([d[m] for d in v])), 4) for m in v[0]}} for k, v in out.items()}
	base = json.load(open(f"{RT.D}/res_test.json"))["vae_ceiling"]
	for k in res: # 对拍 res_test.json（同一条路径 ⇒ 同数）
	res[k]["matches_res_test_json"] = (abs(res[k]["psnr"] - base[k]["player_crop_psnr_vs_native_gt"]) < 0.01
	and abs(res[k]["detail"] - base[k]["detail_ratio_vs_native_gt"]) < 0.001)
	print(json.dumps(res, indent=1), flush=True)
	json.dump({"vae_ceiling_alignment": res, "note": __doc__}, open(f"{RT.D}/res_test_edge.json", "w"), indent=1,
	ensure_ascii=False)


	if __name__ == "__main__":
	sys.exit(main())