添加数据增强方案以及扩散生成模型的想法

tools/validate_h_consistency.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""
+Validate homography consistency produced by ICLayoutTrainingDataset.
+For random samples, we check that cv2.warpPerspective(original, H) ≈ transformed.
+Saves visual composites and prints basic metrics (MSE / PSNR).
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+# Ensure project root is on sys.path when running as a script
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from data.ic_dataset import ICLayoutTrainingDataset
+
+
+def tensor_to_u8_img(t: torch.Tensor) -> np.ndarray:
+    """Convert 1xHxW or 3xHxW float tensor in [0,1] to uint8 HxW or HxWx3."""
+    if t.dim() != 3:
+        raise ValueError(f"Expect 3D tensor, got {t.shape}")
+    if t.size(0) == 1:
+        arr = (t.squeeze(0).cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
+    elif t.size(0) == 3:
+        arr = (t.permute(1, 2, 0).cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
+    else:
+        raise ValueError(f"Unexpected channels: {t.size(0)}")
+    return arr
+
+
+def mse(a: np.ndarray, b: np.ndarray) -> float:
+    diff = a.astype(np.float32) - b.astype(np.float32)
+    return float(np.mean(diff * diff))
+
+
+def psnr(a: np.ndarray, b: np.ndarray) -> float:
+    m = mse(a, b)
+    if m <= 1e-8:
+        return float('inf')
+    return 10.0 * np.log10((255.0 * 255.0) / m)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Validate homography consistency")
+    parser.add_argument("--dir", dest="image_dir", type=str, required=True, help="PNG images directory")
+    parser.add_argument("--out", dest="out_dir", type=str, default="validate_h_out", help="Output directory for composites")
+    parser.add_argument("--n", dest="num", type=int, default=8, help="Number of samples to validate")
+    parser.add_argument("--patch", dest="patch_size", type=int, default=256)
+    parser.add_argument("--elastic", dest="use_elastic", action="store_true")
+    args = parser.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Use no photometric/Sobel transform here to compare raw grayscale content
+    ds = ICLayoutTrainingDataset(
+        args.image_dir,
+        patch_size=args.patch_size,
+        transform=None,
+        scale_range=(1.0, 1.0),
+        use_albu=args.use_elastic,
+        albu_params={"prob": 0.5},
+    )
+
+    n = min(args.num, len(ds))
+    if n == 0:
+        print("[WARN] Empty dataset.")
+        return
+
+    mses = []
+    psnrs = []
+
+    for i in range(n):
+        patch_t, trans_t, H2x3_t = ds[i]
+        # Convert to uint8 arrays
+        patch_u8 = tensor_to_u8_img(patch_t)
+        trans_u8 = tensor_to_u8_img(trans_t)
+        if patch_u8.ndim == 3:
+            patch_u8 = cv2.cvtColor(patch_u8, cv2.COLOR_BGR2GRAY)
+        if trans_u8.ndim == 3:
+            trans_u8 = cv2.cvtColor(trans_u8, cv2.COLOR_BGR2GRAY)
+
+        # Reconstruct 3x3 H
+        H2x3 = H2x3_t.numpy()
+        H = np.vstack([H2x3, [0.0, 0.0, 1.0]]).astype(np.float32)
+
+        # Warp original with H
+        warped = cv2.warpPerspective(patch_u8, H, (patch_u8.shape[1], patch_u8.shape[0]))
+
+        # Metrics
+        m = mse(warped, trans_u8)
+        p = psnr(warped, trans_u8)
+        mses.append(m)
+        psnrs.append(p)
+
+        # Composite image: [orig | warped | transformed | absdiff]
+        diff = cv2.absdiff(warped, trans_u8)
+        comp = np.concatenate([
+            patch_u8, warped, trans_u8, diff
+        ], axis=1)
+        out_path = out_dir / f"sample_{i:03d}.png"
+        cv2.imwrite(out_path.as_posix(), comp)
+        print(f"[OK] sample {i}: MSE={m:.2f}, PSNR={p:.2f} dB -> {out_path}")
+
+    print(f"\nSummary: MSE avg={np.mean(mses):.2f} ± {np.std(mses):.2f}, PSNR avg={np.mean(psnrs):.2f} dB")
+
+
+if __name__ == "__main__":
+    main()