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

tools/diffusion/prepare_patch_dataset.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+"""
+Prepare raster patch dataset and optional condition maps for diffusion training.
+
+Planned inputs:
+- --src_dirs: one or more directories containing PNG layout images
+- --out_dir: output root for images/ and conditions/
+- --size: patch size (e.g., 256)
+- --stride: sliding stride for patch extraction
+- --min_fg_ratio: minimum foreground ratio to keep a patch (0-1)
+- --make_conditions: flags to generate edge/skeleton/distance maps
+
+Current status: CLI skeleton and TODOs only.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Prepare patch dataset for diffusion training (skeleton)")
+    parser.add_argument("--src_dirs", type=str, nargs="+", help="Source PNG dirs for layouts")
+    parser.add_argument("--out_dir", type=str, required=True, help="Output root directory")
+    parser.add_argument("--size", type=int, default=256, help="Patch size")
+    parser.add_argument("--stride", type=int, default=256, help="Patch stride")
+    parser.add_argument("--min_fg_ratio", type=float, default=0.02, help="Min foreground ratio to keep a patch")
+    parser.add_argument("--make_edge", action="store_true", help="Generate edge map conditions (e.g., Sobel/Canny)")
+    parser.add_argument("--make_skeleton", action="store_true", help="Generate morphological skeleton condition")
+    parser.add_argument("--make_dist", action="store_true", help="Generate distance transform condition")
+    args = parser.parse_args()
+
+    out_root = Path(args.out_dir)
+    out_root.mkdir(parents=True, exist_ok=True)
+    (out_root / "images").mkdir(exist_ok=True)
+    (out_root / "conditions").mkdir(exist_ok=True)
+
+    # TODO: implement extraction loop over src_dirs, crop patches, filter by min_fg_ratio,
+    # and save into images/; generate optional condition maps into conditions/ mirroring filenames.
+    # Keep file naming consistent: images/xxx.png, conditions/xxx_edge.png, etc.
+
+    print("[TODO] Implement patch extraction and condition map generation.")
+
+
+if __name__ == "__main__":
+    main()

tools/diffusion/sample_layouts.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""
+Sample layout patches using a trained diffusion model (skeleton).
+
+Outputs raster PNGs into a target directory compatible with current training pipeline (no H pairing).
+
+Current status: CLI skeleton and TODOs only.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Sample layout patches from diffusion model (skeleton)")
+    parser.add_argument("--ckpt", type=str, required=True, help="Path to trained diffusion checkpoint or HF repo id")
+    parser.add_argument("--out_dir", type=str, required=True, help="Directory to write sampled PNGs")
+    parser.add_argument("--num", type=int, default=200)
+    parser.add_argument("--image_size", type=int, default=256)
+    parser.add_argument("--guidance", type=float, default=5.0)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--cond_dir", type=str, default=None, help="Optional condition maps directory")
+    parser.add_argument("--cond_types", type=str, nargs="*", default=None, help="e.g., edge skeleton dist")
+    args = parser.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # TODO: load pipeline from ckpt, set scheduler, handle conditions if provided,
+    # sample args.num images, save as PNG files into out_dir.
+
+    print("[TODO] Implement diffusion sampling and PNG saving.")
+
+
+if __name__ == "__main__":
+    main()

tools/diffusion/train_layout_diffusion.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+"""
+Train a diffusion model for layout patch generation (skeleton).
+
+Planned: fine-tune Stable Diffusion (or Latent Diffusion) with optional ControlNet edge/skeleton conditions.
+
+Dependencies to consider: diffusers, transformers, accelerate, torch, torchvision, opencv-python.
+
+Current status: CLI skeleton and TODOs only.
+"""
+from __future__ import annotations
+
+import argparse
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train diffusion model for layout patches (skeleton)")
+    parser.add_argument("--data_dir", type=str, required=True, help="Prepared dataset root (images/ + conditions/)")
+    parser.add_argument("--output_dir", type=str, required=True, help="Checkpoint output directory")
+    parser.add_argument("--image_size", type=int, default=256)
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--max_steps", type=int, default=100000)
+    parser.add_argument("--use_controlnet", action="store_true", help="Train with ControlNet conditioning")
+    parser.add_argument("--condition_types", type=str, nargs="*", default=["edge"], help="e.g., edge skeleton dist")
+    args = parser.parse_args()
+
+    # TODO: implement dataset/dataloader (images and optional conditions)
+    # TODO: load base pipeline (Stable Diffusion or Latent Diffusion) and optionally ControlNet
+    # TODO: set up optimizer, LR schedule, EMA, gradient accumulation, and run training loop
+    # TODO: save periodic checkpoints to output_dir
+
+    print("[TODO] Implement diffusion training loop and checkpoints.")
+
+
+if __name__ == "__main__":
+    main()