From 08f488f0d8fa59d4add9b24038bd036da45cc030 Mon Sep 17 00:00:00 2001 From: Jiao77 Date: Mon, 20 Oct 2025 21:14:03 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=95=B0=E6=8D=AE=E5=A2=9E?= =?UTF-8?q?=E5=BC=BA=E6=96=B9=E6=A1=88=E4=BB=A5=E5=8F=8A=E6=89=A9=E6=95=A3?= =?UTF-8?q?=E7=94=9F=E6=88=90=E6=A8=A1=E5=9E=8B=E7=9A=84=E6=83=B3=E6=B3=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 123 ++++++- benchmark_grid.json | 92 +++++ configs/base_config.yaml | 21 ++ data/ic_dataset.py | 61 +++- docs/NextStep.md | 317 ++++++++++-------- docs/description/Performance_Benchmark.md | 47 +++ docs/feature_work.md | 178 +++++++++- .../03_Stage3_Integration_Optimization.md | 25 ++ docs/todos/04_Stage4_Advanced_Features.md | 35 ++ models/rord.py | 51 ++- tests/benchmark_attention.py | 91 +++++ tests/benchmark_grid.py | 96 ++++++ tools/diffusion/prepare_patch_dataset.py | 46 +++ tools/diffusion/sample_layouts.py | 38 +++ tools/diffusion/train_layout_diffusion.py | 37 ++ tools/generate_synthetic_layouts.py | 90 +++++ tools/layout2png.py | 160 +++++++++ tools/preview_dataset.py | 68 ++++ tools/smoke_test.py | 76 +++++ tools/synth_pipeline.py | 169 ++++++++++ tools/validate_h_consistency.py | 117 +++++++ train.py | 155 ++++++++- 22 files changed, 1903 insertions(+), 190 deletions(-) create mode 100644 benchmark_grid.json create mode 100644 tests/benchmark_attention.py create mode 100644 tests/benchmark_grid.py create mode 100644 tools/diffusion/prepare_patch_dataset.py create mode 100644 tools/diffusion/sample_layouts.py create mode 100644 tools/diffusion/train_layout_diffusion.py create mode 100644 tools/generate_synthetic_layouts.py create mode 100644 tools/layout2png.py create mode 100644 tools/preview_dataset.py create mode 100644 tools/smoke_test.py create mode 100644 tools/synth_pipeline.py create mode 100644 tools/validate_h_consistency.py diff --git a/README.md b/README.md index 6a9b4ff..4def268 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,29 @@ ![Python Version](https://img.shields.io/badge/Python-3.12-blue) ![License](https://img.shields.io/badge/License-Apache%202.0-orange.svg) +## ⚡ Quick Start（含合成数据与H校验） + +```bash +# 一键生成→渲染→预览→H校验→写回配置（开启合成混采与 Elastic） +uv run python tools/synth_pipeline.py \ + --out_root data/synthetic \ + --num 50 \ + --dpi 600 \ + --config configs/base_config.yaml \ + --ratio 0.3 \ + --enable_elastic \ + --validate_h --validate_n 6 +``` + +提示：zsh 下使用反斜杠续行时，确保每行末尾只有一个 `\` 且下一行不要粘连参数（避免如 `6uv` 这样的粘连）。 + +可选：为 KLayout 渲染指定图层配色/线宽/背景（示例：金属层绿色、过孔红色、黑底） +```bash +uv run python tools/layout2png.py \ + --in data/synthetic/gds --out data/synthetic/png --dpi 800 \ + --layermap '1/0:#00FF00,2/0:#FF0000' --line_width 2 --bgcolor '#000000' +``` + ## 📖 描述 本项目实现了 **RoRD (Rotation-Robust Descriptors)** 模型，这是一种先进的局部特征匹配方法，专用于集成电路（IC）版图的识别。 @@ -70,7 +93,9 @@ RoRD-Layout-Recognation/ ├── match.py # 模板匹配脚本（FPN / 滑窗 + NMS） ├── tests/ │ ├── benchmark_fpn.py # FPN vs 滑窗性能对标 -│ └── benchmark_backbones.py # 多骨干 A/B 前向基准 +│ ├── benchmark_backbones.py # 多骨干 A/B 前向基准 +│ ├── benchmark_attention.py # 注意力 none/se/cbam A/B 基准 +│ └── benchmark_grid.py # 三维基准：Backbone × Attention × Single/FPN ├── config.py # 兼容旧流程的 YAML 读取 shim ├── pyproject.toml └── README.md @@ -81,6 +106,11 @@ RoRD-Layout-Recognation/ - **YAML 配置中心**：所有路径与超参数集中存放在 `configs/*.yaml`，通过 `utils.config_loader.load_config` 统一解析；CLI 的 `--config` 参数可切换实验配置，`to_absolute_path` 则保证相对路径相对配置文件解析。 - **旧配置兼容**：`config.py` 现在仅作为兼容层，将 YAML 配置转换成原有的 Python 常量，便于逐步迁移历史代码。 - **损失与数据解耦**：`losses.py` 汇总几何感知损失，`data/ic_dataset.py` 与 `utils/data_utils.py` 分离数据准备逻辑，便于扩展新的采样策略或损失项。 + +# 5. 运行 A/B 基准（骨干、注意力、三维网格） +PYTHONPATH=. uv run python tests/benchmark_backbones.py --device cpu --image-size 512 --runs 5 +PYTHONPATH=. uv run python tests/benchmark_attention.py --device cpu --image-size 512 --runs 10 --backbone resnet34 --places backbone_high desc_head +PYTHONPATH=. uv run python tests/benchmark_grid.py --device cpu --image-size 512 --runs 3 --backbones vgg16 resnet34 efficientnet_b0 --attentions none se cbam --places backbone_high desc_head - **日志体系**：`logging` 配置节配合 TensorBoard 集成，`train.py`、`evaluate.py`、`match.py` 可统一写入 `log_dir/子任务/experiment_name`。 - **模型配置扩展**： - `model.backbone.name`: `vgg16 | resnet34 | efficientnet_b0` @@ -350,6 +380,7 @@ uv run python match.py --config configs/base_config.yaml --no_nms \ 可参考以下文档与脚本复现并查看最新结果： - CPU 多骨干 A/B 基准（512×512，5 次）：见 `docs/description/Performance_Benchmark.md` +- 三维基准（Backbone × Attention × Single/FPN）：见 `docs/description/Performance_Benchmark.md` 与 `tests/benchmark_grid.py` - FPN vs 滑窗对标脚本：`tests/benchmark_fpn.py` - 多骨干 A/B 基准脚本：`tests/benchmark_backbones.py` @@ -357,4 +388,92 @@ uv run python match.py --config configs/base_config.yaml --no_nms \ ## 📄 许可协议 -本项目根据 [Apache License 2.0](LICENSE.txt) 授权。 \ No newline at end of file +本项目根据 [Apache License 2.0](LICENSE.txt) 授权。 + +--- + +## 🧪 合成数据一键流程与常见问题 + +### 一键命令 +```bash +uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42 +uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600 +uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic +uv run python train.py --config configs/base_config.yaml +``` + +或使用单脚本一键执行（含配置写回）： +```bash +uv run python tools/synth_pipeline.py --out_root data/synthetic --num 200 --dpi 600 \ + --config configs/base_config.yaml --ratio 0.3 --enable_elastic +``` + +### YAML 关键片段 +```yaml +synthetic: + enabled: true + png_dir: data/synthetic/png + ratio: 0.3 + +augment: + elastic: + enabled: true + alpha: 40 + sigma: 6 + alpha_affine: 6 + prob: 0.3 +``` + +### 参数建议 +- DPI：600–900；图形极细时可到 1200（注意磁盘占用与 IO）。 +- ratio：数据少取 0.3–0.5；中等 0.2–0.3；数据多 0.1–0.2。 +- Elastic：alpha=40, sigma=6, prob=0.3 为安全起点。 + +### FAQ +- 找不到 `klayout`：安装系统级 KLayout 并加入 PATH；或使用回退（gdstk+SVG）。 +- `cairosvg`/`gdstk` 报错：升级版本、确认写权限、检查输出目录存在。 +- 训练集为空：检查 `paths.layout_dir` 与 `synthetic.png_dir` 是否存在且包含 .png；若 syn 目录为空将自动仅用真实数据。 + +--- + +## 🧪 合成数据管线与可视化 + +### 1) 生成合成 GDS +```bash +uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42 +``` + +### 2) 批量转换 GDS → PNG +```bash +uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600 +``` + +若本机未安装 KLayout，将自动回退到 gdstk+SVG 路径；图像外观可能与 KLayout 有差异。 + +### 3) 开启训练混采 +在 `configs/base_config.yaml` 中设置： +```yaml +synthetic: + enabled: true + png_dir: data/synthetic/png + ratio: 0.3 +``` + +### 4) 预览训练对（目检增强/H 一致性） +```bash +uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic +``` + +### 5) 开启/调整 Elastic 变形 +```yaml +augment: + elastic: + enabled: true + alpha: 40 + sigma: 6 + alpha_affine: 6 + prob: 0.3 + photometric: + brightness_contrast: true + gauss_noise: true +``` diff --git a/benchmark_grid.json b/benchmark_grid.json new file mode 100644 index 0000000..493d645 --- /dev/null +++ b/benchmark_grid.json @@ -0,0 +1,92 @@ +[ + { + "backbone": "vgg16", + "attention": "none", + "places": "backbone_high,desc_head", + "single_ms_mean": 351.6519069671631, + "single_ms_std": 1.8778125281542124, + "fpn_ms_mean": 719.3304697672526, + "fpn_ms_std": 3.949980966745213, + "runs": 3 + }, + { + "backbone": "vgg16", + "attention": "se", + "places": "backbone_high,desc_head", + "single_ms_mean": 349.7585455576579, + "single_ms_std": 1.9950684383137551, + "fpn_ms_mean": 721.4130560557047, + "fpn_ms_std": 2.7448351792281374, + "runs": 3 + }, + { + "backbone": "vgg16", + "attention": "cbam", + "places": "backbone_high,desc_head", + "single_ms_mean": 354.4490337371826, + "single_ms_std": 1.4903953036396786, + "fpn_ms_mean": 744.7629769643148, + "fpn_ms_std": 29.3233387791729, + "runs": 3 + }, + { + "backbone": "resnet34", + "attention": "none", + "places": "backbone_high,desc_head", + "single_ms_mean": 90.98696708679199, + "single_ms_std": 0.41179110533866975, + "fpn_ms_mean": 117.2173023223877, + "fpn_ms_std": 0.40632490569423124, + "runs": 3 + }, + { + "backbone": "resnet34", + "attention": "se", + "places": "backbone_high,desc_head", + "single_ms_mean": 90.78375498453777, + "single_ms_std": 0.4705899743190883, + "fpn_ms_mean": 115.90576171875, + "fpn_ms_std": 1.3081578935341862, + "runs": 3 + }, + { + "backbone": "resnet34", + "attention": "cbam", + "places": "backbone_high,desc_head", + "single_ms_mean": 96.49538993835449, + "single_ms_std": 3.17170034860506, + "fpn_ms_mean": 111.08938852945964, + "fpn_ms_std": 1.0126843546619573, + "runs": 3 + }, + { + "backbone": "efficientnet_b0", + "attention": "none", + "places": "backbone_high,desc_head", + "single_ms_mean": 40.451606114705406, + "single_ms_std": 1.5293525027201111, + "fpn_ms_mean": 127.30161348978679, + "fpn_ms_std": 0.08508800981401025, + "runs": 3 + }, + { + "backbone": "efficientnet_b0", + "attention": "se", + "places": "backbone_high,desc_head", + "single_ms_mean": 46.480417251586914, + "single_ms_std": 0.2622188910897682, + "fpn_ms_mean": 142.35099156697592, + "fpn_ms_std": 6.611047958580852, + "runs": 3 + }, + { + "backbone": "efficientnet_b0", + "attention": "cbam", + "places": "backbone_high,desc_head", + "single_ms_mean": 47.10610707600912, + "single_ms_std": 0.47150733957171853, + "fpn_ms_mean": 150.99199612935385, + "fpn_ms_std": 12.465987661773038, + "runs": 3 + } +] \ No newline at end of file diff --git a/configs/base_config.yaml b/configs/base_config.yaml index 176fbf6..a2b3c26 100644 --- a/configs/base_config.yaml +++ b/configs/base_config.yaml @@ -51,3 +51,24 @@ paths: val_ann_dir: "path/to/val/annotations" template_dir: "path/to/templates" model_path: "path/to/save/model_final.pth" + +# 数据增强与合成数据配置（可选） +augment: + elastic: + enabled: false + alpha: 40 + sigma: 6 + alpha_affine: 6 + prob: 0.3 + photometric: + brightness_contrast: true + gauss_noise: true + +synthetic: + enabled: false + png_dir: "data/synthetic/png" + ratio: 0.0 # 0~1，训练时混合的合成样本比例 + diffusion: + enabled: false + png_dir: "data/synthetic_diff/png" + ratio: 0.0 # 0~1，训练时混合的扩散样本比例 diff --git a/data/ic_dataset.py b/data/ic_dataset.py index 049ab54..607bdaa 100644 --- a/data/ic_dataset.py +++ b/data/ic_dataset.py @@ -1,6 +1,6 @@ import os import json -from typing import Tuple +from typing import Tuple, Optional import cv2 import numpy as np @@ -70,6 +70,8 @@ class ICLayoutTrainingDataset(Dataset): patch_size: int = 256, transform=None, scale_range: Tuple[float, float] = (1.0, 1.0), + use_albu: bool = False, + albu_params: Optional[dict] = None, ) -> None: self.image_dir = image_dir self.image_paths = [ @@ -80,6 +82,28 @@ class ICLayoutTrainingDataset(Dataset): self.patch_size = patch_size self.transform = transform self.scale_range = scale_range + # 可选的 albumentations 管道 + self.albu = None + if use_albu: + try: + import albumentations as A # 延迟导入，避免环境未安装时报错 + p = albu_params or {} + elastic_prob = float(p.get("prob", 0.3)) + alpha = float(p.get("alpha", 40)) + sigma = float(p.get("sigma", 6)) + alpha_affine = float(p.get("alpha_affine", 6)) + use_bc = bool(p.get("brightness_contrast", True)) + use_noise = bool(p.get("gauss_noise", True)) + transforms_list = [ + A.ElasticTransform(alpha=alpha, sigma=sigma, alpha_affine=alpha_affine, p=elastic_prob), + ] + if use_bc: + transforms_list.append(A.RandomBrightnessContrast(p=0.5)) + if use_noise: + transforms_list.append(A.GaussNoise(var_limit=(5.0, 20.0), p=0.3)) + self.albu = A.Compose(transforms_list) + except Exception: + self.albu = None def __len__(self) -> int: return len(self.image_paths) @@ -102,22 +126,27 @@ class ICLayoutTrainingDataset(Dataset): patch = image.crop((x, y, x + crop_size, y + crop_size)) patch = patch.resize((self.patch_size, self.patch_size), Image.Resampling.LANCZOS) - # 亮度/对比度增强 - if np.random.random() < 0.5: - brightness_factor = np.random.uniform(0.8, 1.2) - patch = patch.point(lambda px: int(np.clip(px * brightness_factor, 0, 255))) - - if np.random.random() < 0.5: - contrast_factor = np.random.uniform(0.8, 1.2) - patch = patch.point(lambda px: int(np.clip(((px - 128) * contrast_factor) + 128, 0, 255))) - - if np.random.random() < 0.3: - patch_np = np.array(patch, dtype=np.float32) - noise = np.random.normal(0, 5, patch_np.shape) - patch_np = np.clip(patch_np + noise, 0, 255) - patch = Image.fromarray(patch_np.astype(np.uint8)) - + # photometric/elastic（在几何 H 之前） patch_np_uint8 = np.array(patch) + if self.albu is not None: + patch_np_uint8 = self.albu(image=patch_np_uint8)["image"] + patch = Image.fromarray(patch_np_uint8) + else: + # 原有轻量光度增强 + if np.random.random() < 0.5: + brightness_factor = np.random.uniform(0.8, 1.2) + patch = patch.point(lambda px: int(np.clip(px * brightness_factor, 0, 255))) + + if np.random.random() < 0.5: + contrast_factor = np.random.uniform(0.8, 1.2) + patch = patch.point(lambda px: int(np.clip(((px - 128) * contrast_factor) + 128, 0, 255))) + + if np.random.random() < 0.3: + patch_np = np.array(patch, dtype=np.float32) + noise = np.random.normal(0, 5, patch_np.shape) + patch_np = np.clip(patch_np + noise, 0, 255) + patch = Image.fromarray(patch_np.astype(np.uint8)) + patch_np_uint8 = np.array(patch) # 随机旋转与镜像（8个离散变换） theta_deg = int(np.random.choice([0, 90, 180, 270])) diff --git a/docs/NextStep.md b/docs/NextStep.md index 5350d64..f44eb17 100644 --- a/docs/NextStep.md +++ b/docs/NextStep.md @@ -1,173 +1,200 @@ -# 下一步工作计划 (NextStep) +## 一、数据策略与增强 (Data Strategy & Augmentation) -**最后更新**: 2025-10-20 -**范围**: 仅聚焦于 `feature_work.md` 的第二部分「模型架构 (Model Architecture)」的落地执行计划 -**上下文**: 核心功能已完成，本文档将模型架构优化转化为可执行的工程计划，便于直接实施与验收。 +> 目标：提升模型的鲁棒性和泛化能力，减少对大量真实数据的依赖。 -> 参考来源：`docs/feature_work.md` 第二部分；更宏观的阶段规划见 `docs/todos/` +- [x] 引入弹性变形 (Elastic Transformations) + - ✔️ 价值：模拟芯片制造中可能出现的微小物理形变，使模型对非刚性变化更鲁棒。 + - 🧭 关键原则（与当前数据管线一致）： + - 现有自监督训练数据集 `ICLayoutTrainingDataset` 会返回 (original, rotated, H)；其中 H 是两张 patch 间的单应关系，用于 loss 监督。 + - 非刚性弹性变形若只对其中一张或在生成 H 之后施加，会破坏几何约束，导致 H 失效。 + - 因此，Elastic 需在“生成 homography 配对之前”对基础 patch 施加；随后对该已变形的 patch 再执行旋转/镜像与单应计算，这样 H 仍严格成立。 + - 📝 执行计划： + 1) 依赖核对 + - `pyproject.toml` 已包含 `albumentations>=2.0.8`，无需新增依赖；确保环境安装齐全。 + 2) 集成位置与方式 + - 在 `data/ic_dataset.py` 的 `ICLayoutTrainingDataset.__getitem__` 中，裁剪并缩放得到 `patch` 后，转换为 `np.ndarray`，对其调用 `albumentations` 管道（包含 `A.ElasticTransform`）。 + - 将变形后的 `patch_np_uint8` 作为“基准图”，再按现有逻辑计算旋转/镜像与 `homography`，生成 `transformed_patch`，从而确保 H 有效。 + 3) 代码改动清单（建议） + - `data/ic_dataset.py` + - 顶部新增：`import albumentations as A` + - `__init__` 新增可选参数：`use_albu: bool=False`、`albu_params: dict|None=None` + - 在 `__init__` 构造 `self.albu = A.Compose([...])`（当 `use_albu` 为 True 时），包含： + - `A.ElasticTransform(alpha=40, sigma=6, alpha_affine=6, p=0.3)` + - （可选）`A.RandomBrightnessContrast(p=0.5)`、`A.GaussNoise(var_limit=(5.0, 20.0), p=0.3)` 以替代当前手写的亮度/对比度与噪声逻辑（减少重复）。 + - 在 `__getitem__`：裁剪与缩放后，若启用 `self.albu`：`patch_np_uint8 = self.albu(image=patch_np_uint8)["image"]`，随后再计算旋转/镜像与 `homography`。 + - 注意：保持输出张量与当前 `utils.data_utils.get_transform()` 兼容（单通道→三通道→Normalize）。 + - `configs/base_config.yaml` + - 新增配置段： + - `augment.elastic.enabled: true|false` + - `augment.elastic.alpha: 40` + - `augment.elastic.sigma: 6` + - `augment.elastic.alpha_affine: 6` + - `augment.elastic.prob: 0.3` + - （可选）`augment.photometric.*` 开关与参数 + - `train.py` + - 从配置读取上述参数，并将 `use_albu` 与 `albu_params` 通过 `ICLayoutTrainingDataset(...)` 传入（不影响现有 `get_transform()`）。 + 4) 参数与默认值建议 + - 起始：`alpha=40, sigma=6, alpha_affine=6, p=0.3`；根据训练收敛与可视化效果微调。 + - 若发现描述子对局部形变敏感，可逐步提高 `alpha` 或 `p`；若训练不稳定则降低。 + 5) 验证与可视化 + - 在 `tests/benchmark_grid.py` 或新增简单可视化脚本中，采样 16 个 (original, rotated) 对，叠加可视化 H 变换后的网格，确认几何一致性未破坏。 + - 训练前 1000 个 batch：记录 `loss_det/loss_desc` 曲线，确认未出现异常发散。 ---- +- [x] 创建合成版图数据生成器 + - ✔️ 价值：解决真实版图数据获取难、数量少的问题，通过程序化生成大量多样化的训练样本。 + - 📝 执行计划： + 1) 新增脚本 `tools/generate_synthetic_layouts.py` + - 目标：使用 `gdstk` 程序化生成包含不同尺寸、密度与单元类型的 GDSII 文件。 + - 主要能力： + - 随机生成“标准单元”模版（如若干矩形/多边形组合）、金属走线、过孔阵列； + - 支持多层（layer/datatype）与规则化阵列（row/col pitch）、占空比（density）控制； + - 形状参数与布局由随机种子控制，支持可重复性。 + - CLI 设计（示例）： + - `--out-dir data/synthetic/gds`、`--num-samples 1000`、`--seed 42` + - 版图规格：`--width 200um --height 200um --grid 0.1um` + - 多样性开关：`--cell-types NAND,NOR,INV --metal-layers 3 --density 0.1-0.6` + - 关键实现要点： + - 使用 `gdstk.Library()` 与 `gdstk.Cell()` 组装基本单元； + - 通过 `gdstk.Reference` 和阵列生成放置； + - 生成完成后 `library.write_gds(path)` 落盘。 + 2) 批量转换 GDSII → PNG（训练用） + - 现状核对：仓库中暂无 `tools/layout2png.py`；计划新增该脚本（与本项一并交付）。 + - 推荐实现 A（首选）：使用 `klayout` 的 Python API（`pya`）以无头模式加载 GDS，指定层映射与缩放，导出为高分辨率 PNG： + - 脚本 `tools/layout2png.py` 提供 CLI：`--in data/synthetic/gds --out data/synthetic/png --dpi 600 --layers 1/0:gray,2/0:blue ...` + - 支持目录批量与单文件转换；可配置画布背景、线宽、边距。 + - 替代实现 B：导出 SVG 再用 `cairosvg` 转 PNG（依赖已在项目中），适合无 klayout 环境的场景。 + - 输出命名规范：与 GDS 同名，如 `chip_000123.gds → chip_000123.png`。 + 3) 数据目录与元数据 + - 目录结构建议： + - `data/synthetic/gds/`、`data/synthetic/png/`、`data/synthetic/meta/` + - 可选：为每个样本生成 `meta/*.json`，记录层数、单元类型分布、密度等，用于后续分析/分层采样。 + 4) 与训练集集成 + - `configs/base_config.yaml` 新增： + - `paths.synthetic_dir: data/synthetic/png` + - `training.use_synthetic_ratio: 0.0~1.0`（混合采样比例；例如 0.3 表示 30% 合成样本） + - 在 `train.py` 中： + - 若 `use_synthetic_ratio>0`，构建一个 `ICLayoutTrainingDataset` 指向合成 PNG 目录； + - 实现简单的比例采样器或 `ConcatDataset + WeightedRandomSampler` 以按比例混合真实与合成样本。 + 5) 质量与稳健性检查 + - 可视化抽样：随机展示若干 PNG，检查层次颜色、对比度、线宽是否清晰； + - 分布对齐：统计真实数据与合成数据的连线长度分布、拓扑度量（如节点度、环路数量），做基础分布对齐； + - 训练烟雾测试：仅用 100～200 个合成样本跑 1～2 个 epoch，确认训练闭环无错误、loss 正常下降。 + 6) 基准验证与复盘 + - 在 `tests/benchmark_grid.py` 与 `tests/benchmark_backbones.py` 增加一组“仅真实 / 真实+合成”的对照实验； + - 记录 mAP/匹配召回/描述子一致性等指标，评估增益； + - 产出 `docs/Performance_Benchmark.md` 的对比表格。 -## 🔴 模型架构优化（Feature Work 第二部分） +### 验收标准 (Acceptance Criteria) -目标：在保证现有精度的前提下，提升特征提取效率与推理速度；为后续注意力机制与多尺度策略提供更强的特征基础。 +- Elastic 变形： + - [ ] 训练数据可视化（含 H 网格叠加）无几何错位； + - [ ] 训练前若干 step loss 无异常尖峰，长期收敛不劣于 baseline； + - [ ] 可通过配置无缝开/关与调参。 +- 合成数据： + - [ ] 能批量生成带多层元素的 GDS 文件并成功转为 PNG； + - [ ] 训练脚本可按设定比例混合采样真实与合成样本； + - [ ] 在小规模对照实验中，验证指标有稳定或可解释的变化（不劣化）。 -### 总体验收标准（全局） -- [ ] 训练/验证流程在新骨干和注意力方案下均可跑通，无崩溃/NaN。 -- [ ] 在代表性验证集上，最终指标（IoU/mAP）不低于当前 VGG-16 基线；若下降需给出改进措施或回滚建议。 -- [ ] 推理时延或显存占用至少一种维度优于基线，或达到“相当 + 结构可扩展”的工程收益。 -- [ ] 关键改动均通过配置开关控制，可随时回退。 +### 风险与规避 (Risks & Mitigations) ---- +- 非刚性变形破坏 H 的风险：仅在生成 homography 前对基准 patch 施加 Elastic，或在两图上施加相同变形但更新 H′=f∘H∘f⁻¹（当前计划采用前者，简单且稳定）。 +- GDS → PNG 渲染差异：优先使用 `klayout`，保持工业级渲染一致性；无 `klayout` 时使用 SVG→PNG 备选路径。 +- 合成分布与真实分布不匹配：通过密度与单元类型分布约束进行对齐，并在训练中控制混合比例渐进提升。 -## 2.1 实验更现代的骨干网络（Backbone） +### 里程碑与时间估算 (Milestones & ETA) -优先级：🟠 中 | 预计工期：~1 周 | 产出：可切换的 backbone 实现 + 对照报告 +## 二、实现状态与使用说明（2025-10-20 更新） -### 设计要点（小合约） -- 输入：与现有 `RoRD` 一致的图像张量 B×C×H×W。 -- 输出：供检测头/描述子头使用的中高层特征张量；通道数因骨干不同而异（VGG:512、ResNet34:512、Eff-B0:1280）。 -- 约束：不改变下游头部的接口形状（头部输入通道需根据骨干进行对齐适配）。 -- 失败模式：通道不匹配/梯度不通/预训练权重未正确加载/收敛缓慢。 +- Elastic 变形已按计划集成： + - 开关与参数：见 `configs/base_config.yaml` 下的 `augment.elastic` 与 `augment.photometric`； + - 数据集实现：`data/ic_dataset.py` 中 `ICLayoutTrainingDataset`； + - 可视化验证：`tools/preview_dataset.py --dir --n 8 --elastic`。 -### 配置扩展（YAML） -在 `configs/base_config.yaml` 增加（或确认存在）： +- 合成数据生成与渲染： + - 生成 GDS：`tools/generate_synthetic_layouts.py --out-dir data/synthetic/gds --num 100 --seed 42`； + - 转换 PNG：`tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600`； + - 训练混采：在 `configs/base_config.yaml` 设置 `synthetic.enabled: true`、`synthetic.png_dir: data/synthetic/png`、`synthetic.ratio: 0.3`。 -```yaml -model: - backbone: - name: "vgg16" # 可选：vgg16 | resnet34 | efficientnet_b0 - pretrained: true - # 用于选择抽取的特征层（按不同骨干约定名称） - feature_layers: - vgg16: ["relu3_3", "relu4_3"] - resnet34: ["layer3", "layer4"] - efficientnet_b0: ["features_5", "features_7"] +- 训练脚本： + - `train.py` 已接入真实/合成混采（ConcatDataset + WeightedRandomSampler），验证集仅用真实数据； + - TensorBoard 文本摘要记录数据构成（mix 开关、比例、样本量）。 + +注意：若未安装 KLayout，可自动回退 gdstk+SVG 路径；显示效果可能与 KLayout 存在差异。 + +- D1：Elastic 集成 + 可视化验证（代码改动与测试） +- D2：合成生成器初版（GDS 生成 + PNG 渲染脚本） +- D3：训练混合采样接入 + 小规模基准 +- D4：参数扫与报告更新（Performance_Benchmark.md） + +### 一键流水线（生成 → 渲染 → 预览 → 训练） + +1) 生成 GDS（合成版图） +```bash +uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42 ``` -### 代码改动建议 -- 文件：`models/rord.py` - 1) 在 `__init__` 中根据 `cfg.model.backbone.name` 动态构建骨干： - - vgg16（现状保持） - - resnet34：从 `torchvision.models.resnet34(weights=IMAGENET1K_V1)` 构建；保存 `layer3/layer4` 输出。 - - efficientnet_b0：从 `torchvision.models.efficientnet_b0(weights=IMAGENET1K_V1)` 构建；保存末两段 `features` 输出。 - 2) 为不同骨干提供统一的“中间层特征导出”接口（注册 forward hook 或显式调用子模块）。 - 3) 依据所选骨干的输出通道，调整检测头与描述子头的输入通道（如使用 1×1 conv 过渡层以解耦通道差异）。 - 4) 保持现有前向签名与返回数据结构不变（训练/推理兼容）。 +2) 渲染 PNG（KLayout 优先，自动回退 gdstk+SVG） +```bash +uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600 +``` -### 进展更新（2025-10-20） -- 已完成：在 `models/rord.py` 集成多骨干选择（`vgg16`/`resnet34`/`efficientnet_b0`），并实现统一的中间层抽取函数 `_extract_c234`（可后续重构为 `build_backbone`/`extract_features` 明确接口）。 -- 已完成：FPN 通用化，基于 C2/C3/C4 构建 P2/P3/P4，按骨干返回正确的 stride。 -- 已完成：单图前向 Smoke Test（三种骨干，单尺度与 FPN）均通过。 -- 已完成：CPU 环境 A/B 基准（单尺度 vs FPN）见 `docs/description/Performance_Benchmark.md`。 -- 待完成：GPU 环境基准（速度/显存）、基于真实数据的精度评估与收敛曲线对比。 +3) 预览训练对（核验增强/H 一致性） +```bash +uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic +``` -### 落地步骤（Checklist） -- [x] 在 `models/rord.py` 增加/落地骨干构建与中间层抽取逻辑（当前通过 `_extract_c234` 实现）。 -- [x] 接入 ResNet-34：返回等价中高层特征（layer2/3/4，通道≈128/256/512）。 -- [x] 接入 EfficientNet-B0：返回 `features[2]/[3]/[6]`（约 24/40/192），FPN 以 1×1 横向连接对齐到 `fpn_out_channels`。 -- [x] 头部适配：单尺度头使用骨干高层通道数；FPN 头统一使用 `fpn_out_channels`。 -- [ ] 预训练权重：支持 `pretrained=true` 加载；补充权重加载摘要打印（哪些层未命中）。 -- [x] 单图 smoke test：前向通过、无 NaN（三种骨干，单尺度与 FPN）。 - -### 评测与选择（A/B 实验） -- [ ] 在固定数据与超参下，比较 vgg16/resnet34/efficientnet_b0： - - 收敛速度（loss 曲线 0-5 epoch） - - 推理速度（ms / 2048×2048）与显存（GB）[CPU 初步结果已产出，GPU 待复测；见 `docs/description/Performance_Benchmark.md`] - - 验证集 IoU/mAP（真实数据集待跑） -- [ ] 形成表格与可视化图，给出选择结论与原因（CPU 版初稿已在报告中给出观察）。 -- [ ] 若新骨干在任一关键指标明显受损，则暂缓替换，仅保留为可切换实验选项。 - -### 验收标准（2.1） -- [ ] 三种骨干方案均可训练与推理（当前仅验证推理，训练与收敛待验证）； -- [ ] 最终入选骨干在 IoU/mAP 不低于 VGG 的前提下，带来显著的速度/显存优势之一； -- [x] 切换完全配置化（无需改代码）。 - -### 风险与回滚（2.1） -- 通道不匹配导致维度错误 → 在进入头部前统一使用 1×1 conv 适配； -- 预训练权重与自定义层名不一致 → 显式映射并记录未加载层； -- 收敛变慢 → 暂时提高训练轮数、调学习率/BN 冻结策略；不达标即回滚 `backbone.name=vgg16`。 - ---- - -## 2.2 集成注意力机制（CBAM / SE-Net） - -优先级：🟠 中 | 预计工期：~7–10 天 | 产出：注意力增强的 RoRD 变体 + 对照报告 - -### 模块选择与嵌入位置 -- 方案 A：CBAM（通道注意 + 空间注意），插入至骨干高层与两类头部之前； -- 方案 B：SE-Net（通道注意），轻量但仅通道维，插入多个阶段以增强稳定性； -- 建议：先实现 CBAM，保留 SE 作为备选开关。 - -### 配置扩展（YAML） +4) 在 YAML 中开启混采与 Elastic（示例） ```yaml -model: - attention: +synthetic: + enabled: true + png_dir: data/synthetic/png + ratio: 0.3 + +augment: + elastic: enabled: true - type: "cbam" # 可选：cbam | se | none - places: ["backbone_high", "det_head", "desc_head"] - # 可选超参：reduction、kernel_size 等 - reduction: 16 - spatial_kernel: 7 + alpha: 40 + sigma: 6 + alpha_affine: 6 + prob: 0.3 ``` -### 代码改动建议 -- 文件：`models/rord.py` - 1) 实现 `CBAM` 与 `SEBlock` 模块（或从可靠实现迁移），提供简洁 forward。 - 2) 在 `__init__` 中依据 `cfg.model.attention` 决定在何处插入： - - backbone 高层输出后（增强高层语义的判别性）； - - 检测头、描述子头输入前（分别强化不同任务所需特征）。 - 3) 注意保持张量尺寸不变；若引入残差结构，保证与原路径等价时可退化为恒等映射。 +5) 开始训练 +```bash +uv run python train.py --config configs/base_config.yaml +``` -### 落地步骤（Checklist） -- [ ] 实现 `CBAM`：通道注意（MLP/Avg+Max Pool）+ 空间注意（7×7 conv）。 -- [ ] 实现 `SEBlock`：Squeeze（全局池化）+ Excitation（MLP, reduction）。 -- [ ] 在 `RoRD` 中用配置化开关插拔注意力，默认关闭。 -- [ ] 在进入检测/描述子头前分别测试开启/关闭注意力的影响。 -- [ ] 记录注意力图（可选）：导出中间注意图用于可视化对比。 +可选：使用单脚本一键执行（含配置写回） +```bash +uv run python tools/synth_pipeline.py --out_root data/synthetic --num 200 --dpi 600 \ + --config configs/base_config.yaml --ratio 0.3 --enable_elastic +``` -### 训练与评估 -- [ ] 以入选骨干为基线，分别开启 `cbam` 与 `se` 进行对照； -- [ ] 记录：训练损失、验证 IoU/mAP、推理时延/显存； -- [ ] 观察注意力图是否集中在关键几何（边角/交点/突变）； -- [ ] 若带来过拟合迹象（验证下降），尝试减弱注意力强度或减少插入位置。 +### 参数建议与经验 -### 验收标准（2.2） -- [ ] 模型在开启注意力后稳定训练，无数值异常； -- [ ] 指标不低于无注意力基线；若提升则量化收益； -- [ ] 配置可一键关闭以回退。 +- 渲染 DPI：600–900 通常足够，图形极细时可提高到 1200（注意磁盘与 IO）。 +- 混采比例 synthetic.ratio： + - 数据少（<500 张）可取 0.3–0.5； + - 数据中等（500–2000 张）建议 0.2–0.3； + - 数据多（>2000 张）建议 0.1–0.2 以免分布偏移。 +- Elastic 强度：从 alpha=40, sigma=6 开始；若描述子对局部形变敏感，可小步上调 alpha 或 prob。 -### 风险与回滚（2.2） -- 注意力导致过拟合或梯度不稳 → 降低 reduction、减少插入点、启用正则； -- 推理时延上升明显 → 对注意力路径进行轻量化（如仅通道注意或更小 kernel）。 +### 质量检查清单（建议在首次跑通后执行） ---- +- 预览拼图无明显几何错位（orig/rot 对应边界对齐合理）。 +- 训练日志包含混采信息（real/syn 样本量、ratio、启停状态）。 +- 若开启 Elastic，训练初期 loss 无异常尖峰，长期收敛不劣于 baseline。 +- 渲染 PNG 与 GDS 在关键层上形态一致（优先使用 KLayout）。 -## 工程与度量配套 +### 常见问题与排查（FAQ） -### 实验记录（建议） -- 在 TensorBoard 中新增： - - `arch/backbone_name`、`arch/attention_type`（Text/Scalar）； - - `train/loss_total`、`eval/iou_metric`、`eval/map`； - - 推理指标：`infer/ms_per_image`、`infer/vram_gb`。 - -### 对照报告模板（最小集） -- 数据集与配置摘要（随机种子、批大小、学习率、图像尺寸）。 -- 三个骨干 + 注意力开关的结果表（速度/显存/IoU/mAP）。 -- 结论与落地选择（保留/关闭/待进一步实验）。 - ---- - -## 排期与里程碑（建议） -- M1（1 天）：骨干切换基础设施与通道适配层；单图 smoke 测试。 -- M2（2–3 天）：ResNet34 与 EfficientNet-B0 接入与跑通； -- M3（1–2 天）：A/B 评测与结论； -- M4（3–4 天）：注意力模块接入、训练对照、报告输出。 - ---- - -## 相关参考 -- 源文档：`docs/feature_work.md` 第二部分（模型架构） -- 阶段规划：`docs/todos/` -- 配置系统：`configs/base_config.yaml` +- klayout: command not found + - 方案A：安装系统级 KLayout 并确保可执行文件在 PATH； + - 方案B：暂用 gdstk+SVG 回退（外观可能略有差异）。 +- cairosvg 报错或 SVG 不生成 + - 升级 `cairosvg` 与 `gdstk`；确保磁盘有写入权限；检查 `.svg` 是否被安全软件拦截。 +- gdstk 版本缺少 write_svg + - 尝试升级 gdstk；脚本已做 library 与 cell 双路径兼容，仍失败则优先使用 KLayout。 +- 训练集为空或样本过少 + - 检查 `paths.layout_dir` 与 `synthetic.png_dir` 是否存在且包含 .png；ratio>0 但 syn 目录为空会自动回退仅真实数据。 diff --git a/docs/description/Performance_Benchmark.md b/docs/description/Performance_Benchmark.md index 42ce6bb..0d2841e 100644 --- a/docs/description/Performance_Benchmark.md +++ b/docs/description/Performance_Benchmark.md @@ -17,6 +17,53 @@ - 备注：本次测试在 CPU 上进行，`gpu_mem_mb` 始终为 0。 +## 注意力 A/B（CPU，resnet34，512×512，runs=10，places=backbone_high+desc_head） + +| Attention | Single Mean ± Std | FPN Mean ± Std | +|-----------|-------------------:|----------------:| +| none | 97.57 ± 0.55 | 124.57 ± 0.48 | +| se | 101.48 ± 2.13 | 123.12 ± 0.50 | +| cbam | 119.80 ± 2.38 | 123.11 ± 0.71 | + +观察： +- 单尺度路径对注意力类型更敏感，CBAM 开销相对更高，SE 较轻； +- FPN 路径耗时在本次设置下差异很小（可能因注意力仅在 `backbone_high/desc_head`，且 FPN 头部计算占比较高）。 + +复现实验： +```zsh +PYTHONPATH=. uv run python tests/benchmark_attention.py \ + --device cpu --image-size 512 --runs 10 \ + --backbone resnet34 --places backbone_high desc_head +``` + +## 三维基准（Backbone × Attention × Single/FPN） + +环境：CPU，输入 1×3×512×512，重复 3 次，places=backbone_high,desc_head。 + +| Backbone | Attention | Single Mean ± Std (ms) | FPN Mean ± Std (ms) | +|------------------|-----------|-----------------------:|--------------------:| +| vgg16 | none | 351.65 ± 1.88 | 719.33 ± 3.95 | +| vgg16 | se | 349.76 ± 2.00 | 721.41 ± 2.74 | +| vgg16 | cbam | 354.45 ± 1.49 | 744.76 ± 29.32 | +| resnet34 | none | 90.99 ± 0.41 | 117.22 ± 0.41 | +| resnet34 | se | 90.78 ± 0.47 | 115.91 ± 1.31 | +| resnet34 | cbam | 96.50 ± 3.17 | 111.09 ± 1.01 | +| efficientnet_b0 | none | 40.45 ± 1.53 | 127.30 ± 0.09 | +| efficientnet_b0 | se | 46.48 ± 0.26 | 142.35 ± 6.61 | +| efficientnet_b0 | cbam | 47.11 ± 0.47 | 150.99 ± 12.47 | + +复现实验： + +```zsh +PYTHONPATH=. uv run python tests/benchmark_grid.py \ + --device cpu --image-size 512 --runs 3 \ + --backbones vgg16 resnet34 efficientnet_b0 \ + --attentions none se cbam \ + --places backbone_high desc_head +``` + +运行会同时输出控制台摘要并保存 JSON：`benchmark_grid.json`。 + ## 观察与解读 - vgg16 明显最慢，FPN 额外的横向/上采样代价在 CPU 上更突出（>2×）。 - resnet34 在单尺度上显著快于 vgg16，FPN 增幅较小（约 +25%）。 diff --git a/docs/feature_work.md b/docs/feature_work.md index d7d6651..454e04a 100644 --- a/docs/feature_work.md +++ b/docs/feature_work.md @@ -1,5 +1,41 @@ # 后续工作 +## 新增功能汇总（2025-10-20） + +- 数据增强：集成 `albumentations` 的 ElasticTransform（配置在 `augment.elastic`），并保持几何配对的 H 正确性。 +- 合成数据：新增 `tools/generate_synthetic_layouts.py`（GDS 生成）与 `tools/layout2png.py`（GDS→PNG 批量转换）。 +- 训练混采：`train.py` 接入真实/合成混采，按 `synthetic.ratio` 使用加权采样；验证集仅使用真实数据。 +- 可视化：`tools/preview_dataset.py` 快速导出训练对的拼图图，便于人工质检。 + +## 立即可做的小改进 + +- 在 `layout2png.py` 增加图层配色与线宽配置（读取 layermap 或命令行参数）。 +- 为 `ICLayoutTrainingDataset` 添加随机裁剪失败时的回退逻辑（极小图像）。 +- 增加最小单元测试：验证 ElasticTransform 下 H 的 warp 一致性（采样角点/网格点）。 +- 在 README 增加一键命令合集（生成合成数据 → 渲染 → 预览 → 训练）。 + +## 一键流程与排查（摘要） + +**一键命令**： +```bash +uv run python tools/generate_synthetic_layouts.py --out_dir data/synthetic/gds --num 200 --seed 42 +uv run python tools/layout2png.py --in data/synthetic/gds --out data/synthetic/png --dpi 600 +uv run python tools/preview_dataset.py --dir data/synthetic/png --out preview.png --n 8 --elastic +uv run python train.py --config configs/base_config.yaml +``` + +或使用单脚本一键执行（含配置写回）： +```bash +uv run python tools/synth_pipeline.py --out_root data/synthetic --num 200 --dpi 600 \ + --config configs/base_config.yaml --ratio 0.3 --enable_elastic +``` + +**参数建议**：DPI=600–900；ratio=0.2–0.3（首训）；Elastic 从 alpha=40/sigma=6 起步。 + +**FAQ**： +- 找不到 klayout：安装后确保在 PATH；无则使用回退渲染（外观可能有差异）。 +- SVG/PNG 未生成：检查写权限与版本（cairosvg/gdstk），或优先用 KLayout。 + 本文档整合了 RoRD 项目的优化待办清单和训练需求，用于规划未来的开发和实验工作。 --- @@ -12,18 +48,70 @@ > *目标：提升模型的鲁棒性和泛化能力，减少对大量真实数据的依赖。* -- [ ] **引入弹性变形 (Elastic Transformations)** +- [x] **引入弹性变形 (Elastic Transformations)** - **✔️ 价值**: 模拟芯片制造中可能出现的微小物理形变，使模型对非刚性变化更鲁棒。 - **📝 执行方案**: 1. 添加 `albumentations` 库作为项目依赖。 2. 在 `train.py` 的 `ICLayoutTrainingDataset` 类中，集成 `A.ElasticTransform` 到数据增强管道中。 -- [ ] **创建合成版图数据生成器** +- [x] **创建合成版图数据生成器** - **✔️ 价值**: 解决真实版图数据获取难、数量少的问题，通过程序化生成大量多样化的训练样本。 - **📝 执行方案**: 1. 创建一个新脚本，例如 `tools/generate_synthetic_layouts.py`。 2. 利用 `gdstk` 库 编写函数，程序化地生成包含不同尺寸、密度和类型标准单元的 GDSII 文件。 3. 结合 `tools/layout2png.py` 的逻辑，将生成的版图批量转换为 PNG 图像，用于扩充训练集。 +- [ ] **基于扩散生成的版图数据生成器（研究型）** + - **🎯 目标**: 使用扩散模型（Diffusion）生成具备“曼哈顿几何特性”的版图切片（raster PNG），作为现有程序化合成的补充来源，进一步提升数据多样性与风格覆盖。 + - **📦 产物**: + - 推理脚本（计划）: `tools/diffusion/sample_layouts.py` + - 训练脚本（计划）: `tools/diffusion/train_layout_diffusion.py` + - 数据集打包与统计工具（计划）: `tools/diffusion/prepare_patch_dataset.py` + - **🧭 范围界定**: + - 优先生成单层的二值/灰度光栅图像（256–512 像素方形 patch）。 + - 短期不追求多层/DRC 严格约束的工业可制造性；定位为数据增强来源，而非版图设计替代。 + - **🛤️ 技术路线**: + - 路线 A（首选，工程落地快）: 基于 HuggingFace diffusers 的 Latent Diffusion/Stable Diffusion 微调；输入为 1 通道灰度（训练时复制到 3 通道或改 UNet 首层），输出为版图样式图像。 + - 路线 B（结构引导）: 加入 ControlNet/T2I-Adapter 条件，如 Sobel/Canny/直方结构图、粗草图（Scribble）、程序化几何草图，以控制生成的总体连通性与直角占比。 + - 路线 C（两阶段）: 先用程序化生成器输出“草图/骨架”（低细节），再用扩散模型进行“风格化/细化”。 + - **🧱 数据表示与条件**: + - Raster 表示：PNG（二值/灰度），可预生成条件图：Sobel、Canny、距离变换、形态学骨架等。 + - 条件输入建议：`[image (target-like), edge_map, skeleton]` 的任意子集；PoC 以 edge_map 为主。 + - **🧪 训练配置（建议起点）**: + - 图像尺寸：256（PoC），后续 384/512。 + - 批大小：8–16（依显存），学习率 1e-4，训练步数 100k–300k。 + - 数据来源：`data/**/png` 聚合 + 程序合成数据 `data/synthetic/png`；采样时按风格/密度分层均衡。 + - 预处理：随机裁剪非空 patch、二值阈值均衡、弱摄影增强（噪声/对比度）控制在小幅度范围。 + - **🧰 推理与后处理**: + - 采样参数：采样步数 30–100、guidance scale 3–7、seed 固定以便复现。 + - 后处理：Otsu/固定阈值二值化，形态学开闭/细化，断点连接（morphology bridge），可选矢量化（`gdstk` 轮廓化）回写 GDS。 + - **📈 评估指标**: + - 结构统计对齐：水平/垂直边比例、连通组件面积分布、线宽分布、密度直方图与真实数据 KL 距离。 + - 规则近似性：形态学开闭后碎片率、连通率、冗余孤立像素占比。 + - 训练收益：将扩散样本混入 `train.py`，对 IoU/mAP/收敛轮数的提升幅度（与仅程序合成相比）。 + - **🔌 与现有管线集成**: + - 在 `tools/synth_pipeline.py` 增加 `--use_diffusion` 或 `--diffusion_dir`，将扩散生成的 PNG 目录并入训练数据目录。 + - 配置建议新增： + ```yaml + synthetic: + diffusion: + enabled: false + png_dir: data/synthetic_diff/png + ratio: 0.1 # 与真实/程序合成的混采比例 + ``` + - 预览与质检：重用 `tools/preview_dataset.py`，并用 `tools/validate_h_consistency.py` 跳过 H 检查（扩散输出无严格几何配对），改用结构统计工具（后续补充）。 + - **🗓️ 里程碑**: + 1. 第 1 周：数据准备与统计、PoC（预训练 SD + ControlNet-Edge 的小规模微调，256 尺寸）。 + 2. 第 2–3 周：扩大训练（≥50k patch），加入骨架/距离变换条件，完善后处理。 + 3. 第 4 周：与训练管线集成（混采/可视化），对比“仅程序合成 vs 程序合成+扩散”的增益。 + 4. 第 5 周：文档、示例权重与一键脚本（可选导出 ONNX/TensorRT 推理）。 + - **⚠️ 风险与缓解**: + - 结构失真/非曼哈顿：增强条件约束（ControlNet），提高形态学后处理强度；两阶段（草图→细化）。 + - 模式崩塌/多样性不足：分层采样、数据重采样、EMA、风格/密度条件编码。 + - 训练数据不足：先用程序合成预训练，再混入少量真实数据微调。 + - **📚 参考与依赖**: + - 依赖：`diffusers`, `transformers`, `accelerate`, `albumentations`, `opencv-python`, `gdstk` + - 参考：Latent Diffusion、Stable Diffusion、ControlNet、T2I-Adapter 等论文与开源实现 + ### 二、 模型架构 (Model Architecture) > *目标：提升模型的特征提取效率和精度，降低计算资源消耗。* @@ -40,11 +128,19 @@ - 代码：`models/rord.py` - 基准：`tests/benchmark_backbones.py` - 文档：`docs/description/Backbone_FPN_Test_Change_Notes.md`, `docs/description/Performance_Benchmark.md` -- [ ] **集成注意力机制 (Attention Mechanism)** - - **✔️ 价值**: 引导模型自动关注版图中的关键几何结构（如边角、交点），忽略大面积的空白或重复区域，提升特征质量。 - - **📝 执行方案**: - 1. 寻找一个可靠的注意力模块实现，如 CBAM 或 SE-Net。 - 2. 在 `models/rord.py` 中，将该模块插入到 `self.backbone` 和两个 `head` 之间。 +- [x] **集成注意力机制 (Attention Mechanism)** + - **✔️ 价值**: 引导模型关注关键几何结构、弱化冗余区域，提升特征质量与匹配稳定性。 + - **✅ 当前进展（2025-10-20）**: + - 已集成可切换的注意力模块：`SE` 与 `CBAM`；支持通过 `model.attention.enabled/type/places` 配置开启与插入位置（`backbone_high`/`det_head`/`desc_head`）。 + - 已完成 CPU A/B 基准（none/se/cbam，resnet34，places=backbone_high+desc_head），详见 `docs/description/Performance_Benchmark.md`；脚本：`tests/benchmark_attention.py`。 + - **📝 后续动作**: + 1. 扩展更多模块：ECA、SimAM、CoordAttention、SKNet，并保持统一接口与配置。 + 2. 进行插入位置消融（仅 backbone_high / det_head / desc_head / 组合），在 GPU 上复测速度与显存峰值。 + 3. 在真实数据上评估注意力开/关的 IoU/mAP 与收敛差异。 + - **参考**: + - 代码：`models/rord.py` + - 基准：`tests/benchmark_attention.py`, `tests/benchmark_grid.py` + - 文档：`docs/description/Performance_Benchmark.md` ### 三、 训练与损失函数 (Training & Loss Function) @@ -181,6 +277,24 @@ --output-file export.md ``` +### ✅ 三维基准对比（Backbone × Attention × Single/FPN） + +- **文件**: `tests/benchmark_grid.py` ✅，JSON 输出：`benchmark_grid.json` +- **功能**: + - 遍历 `backbone × attention` 组合（当前：vgg16/resnet34/efficientnet_b0 × none/se/cbam） + - 统计单尺度与 FPN 前向的平均耗时与标准差 + - 控制台摘要 + JSON 结果落盘 +- **使用**: + ```bash + PYTHONPATH=. uv run python tests/benchmark_grid.py \ + --device cpu --image-size 512 --runs 3 \ + --backbones vgg16 resnet34 efficientnet_b0 \ + --attentions none se cbam \ + --places backbone_high desc_head + ``` +- **结果**: + - 已将 CPU（512×512，runs=3）结果写入 `docs/description/Performance_Benchmark.md` 的“三维基准”表格，原始数据位于仓库根目录 `benchmark_grid.json`。 + ### 📚 新增文档 | 文档 | 大小 | 说明 | @@ -263,6 +377,8 @@ | | 全面评估指标 | ✅ | 2025-10-19 | | **新增工作** | 性能基准测试 | ✅ | 2025-10-20 | | | TensorBoard 导出工具 | ✅ | 2025-10-20 | +| **二. 模型架构** | 注意力机制（SE/CBAM 基线） | ✅ | 2025-10-20 | +| **新增工作** | 三维基准对比（Backbone×Attention×Single/FPN） | ✅ | 2025-10-20 | ### 未完成的工作项（可选优化） @@ -270,8 +386,52 @@ |------|--------|--------|------| | **一. 数据策略与增强** | 弹性变形增强 | 🟡 低 | 便利性增强 | | | 合成版图生成器 | 🟡 低 | 数据增强 | -| **二. 模型架构** | 现代骨干网络 | 🟠 中 | 性能优化 | -| | 注意力机制 | 🟠 中 | 性能优化 | +| | 基于扩散的版图生成器 | 🟠 中 | 研究型：引入结构条件与形态学后处理，作为数据多样性来源 | + +--- + +## 扩散生成集成的实现说明（新增） + +- 配置新增节点（已添加到 `configs/base_config.yaml`）: + ```yaml + synthetic: + enabled: false + png_dir: data/synthetic/png + ratio: 0.0 + diffusion: + enabled: false + png_dir: data/synthetic_diff/png + ratio: 0.0 + ``` + +- 训练混采（已实现于 `train.py`）: + - 支持三源混采：真实数据 + 程序合成 (`synthetic`) + 扩散合成 (`synthetic.diffusion`)。 + - 目标比例：`real = 1 - (syn_ratio + diff_ratio)`；使用 `WeightedRandomSampler` 近似。 + - 验证集仅使用真实数据，避免评估偏移。 + +- 一键管线扩展（已实现于 `tools/synth_pipeline.py`）: + - 新增 `--diffusion_dir` 参数：将指定目录的 PNG 并入配置文件的 `synthetic.diffusion.png_dir` 并开启 `enabled=true`。 + - 不自动采样扩散图片（避免引入新依赖），仅做目录集成；后续可在该脚本中串联 `tools/diffusion/sample_layouts.py`。 + +- 新增脚本骨架（`tools/diffusion/`）: + - `prepare_patch_dataset.py`: 从现有 PNG 构建 patch 数据集与条件图（CLI 骨架 + TODO）。 + - `train_layout_diffusion.py`: 微调扩散模型的训练脚本（CLI 骨架 + TODO）。 + - `sample_layouts.py`: 使用已训练权重进行采样输出 PNG（CLI 骨架 + TODO）。 + +- 使用建议： + 1) 将扩散采样得到的 PNG 放入某目录，例如 `data/synthetic_diff/png`。 + 2) 运行： + ```bash + uv run python tools/synth_pipeline.py \ + --out_root data/synthetic \ + --num 200 --dpi 600 \ + --config configs/base_config.yaml \ + --ratio 0.3 \ + --diffusion_dir data/synthetic_diff/png + ``` + 3) 在 YAML 中按需设置 `synthetic.diffusion.ratio`（例如 0.1），训练时即自动按比例混采。 + +| **二. 模型架构** | 更多注意力模块（ECA/SimAM/CoordAttention/SKNet） | 🟠 中 | 扩展与消融 | | **三. 训练与损失** | 损失加权自适应 | 🟠 中 | 训练优化 | | | 困难样本采样 | 🟡 低 | 训练优化 | diff --git a/docs/todos/03_Stage3_Integration_Optimization.md b/docs/todos/03_Stage3_Integration_Optimization.md index 9845ae9..9629304 100644 --- a/docs/todos/03_Stage3_Integration_Optimization.md +++ b/docs/todos/03_Stage3_Integration_Optimization.md @@ -82,6 +82,25 @@ - [ ] 特征维度一致性检查 - [ ] GPU/CPU 切换测试 + #### 2.3 基准与评估补充（来自 NextStep 2.1 未完项） + + - [ ] GPU 环境 A/B 基准（速度/显存） + - [ ] 使用 `tests/benchmark_backbones.py` 在 GPU 上复现（20 次，512×512），记录 ms 与 VRAM + - [ ] 追加结果到 `docs/description/Performance_Benchmark.md` + + - [ ] GPU 环境 Attention A/B 基准（速度/显存） + - [ ] 使用 `tests/benchmark_attention.py` 在 GPU 上复现（10 次，512×512），覆盖 `places` 组合（`backbone_high`/`det_head`/`desc_head`） + - [ ] 记录平均耗时与 VRAM 峰值，追加摘要到 `docs/description/Performance_Benchmark.md` + + - [ ] 三维网格基准（Backbone × Attention × Single/FPN） + - [ ] 使用 `tests/benchmark_grid.py` 在 GPU 上跑最小矩阵（例如 3×3，runs=5） + - [ ] 将 JSON 存入 `results/benchmark_grid_YYYYMMDD.json`，在性能文档中追加表格摘要并链接 JSON + + - [ ] 真实数据集精度评估（IoU/mAP 与收敛曲线） + - [ ] 固定数据与超参，训练 5 个 epoch，记录 loss 曲线 + - [ ] 在验证集上评估 IoU/mAP，并与 vgg16 基线对比 + - [ ] 形成对照表与初步结论 + **验收标准**: - [ ] 所有测试用例通过 - [ ] 推理结果符合预期维度和范围 @@ -143,6 +162,12 @@ - [ ] 日志查看方法 - [ ] GPU 内存不足处理 +#### 3.4 预训练权重加载摘要（来自 NextStep 2.1 未完项） + +- [x] 在 `models/rord.py` 加载 `pretrained=true` 时，打印未命中层摘要 + - [x] 记录：加载成功/跳过的层名数量 + - [x] 提供简要输出（missing/unexpected keys，参数量统计）；实现：`models/rord.py::_summarize_pretrained_load` + #### 3.2 编写配置参数文档 - [ ] 创建 `docs/CONFIG.md` diff --git a/docs/todos/04_Stage4_Advanced_Features.md b/docs/todos/04_Stage4_Advanced_Features.md index 9f9b101..6f2f47e 100644 --- a/docs/todos/04_Stage4_Advanced_Features.md +++ b/docs/todos/04_Stage4_Advanced_Features.md @@ -218,6 +218,41 @@ --- +### 4. 注意力机制集成（来自 NextStep 2.2） + +**目标**: 在骨干高层与头部前集成 CBAM / SE，并量化收益 + +#### 4.1 模块实现与插桩 +- [ ] 实现 `CBAM` 与 `SEBlock`（或迁移可靠实现） +- [ ] 在 `models/rord.py` 通过配置插拔：`attention.enabled/type/places` +- [ ] 确保 forward 尺寸不变，默认关闭可回退 + +#### 4.2 训练与评估 +- [ ] 选择入选骨干为基线，分别开启 `cbam` 与 `se` +- [ ] 记录训练损失、验证 IoU/mAP、推理时延/显存 +- [ ] 可选：导出可视化注意力图 + +**验收标准**: +- [ ] 训练稳定，无数值异常 +- [ ] 指标不低于无注意力基线；若提升则量化收益 +- [ ] 配置可一键关闭以回退 + +#### 4.3 扩展模块与插入位置消融 +- [ ] 扩展更多注意力模块：ECA、SimAM、CoordAttention、SKNet + - [ ] 在 `models/rord.py` 实现统一接口与注册表 + - [ ] 在 `configs/base_config.yaml` 增加可选项说明 +- [ ] 插入位置消融 + - [ ] 仅 `backbone_high` / 仅 `det_head` / 仅 `desc_head` / 组合 + - [ ] 使用 `tests/benchmark_attention.py` 统一基准，记录 Single/FPN 时延与 VRAM + - [ ] 在 `docs/description/Performance_Benchmark.md` 增加“注意力插入位置”小节 + +**验收标准**: +- [ ] 所有新增模块 forward 通过，尺寸/类型与现有路径一致 +- [ ] 基准结果可复现并写入文档 +- [ ] 给出速度-精度权衡建议 + +--- + ## 🔄 实施流程 ### 第 1 周: 实验管理集成 diff --git a/models/rord.py b/models/rord.py index 11103c2..81818eb 100644 --- a/models/rord.py +++ b/models/rord.py @@ -91,7 +91,12 @@ class RoRD(nn.Module): # 默认各层通道（VGG 对齐） c2_ch, c3_ch, c4_ch = 128, 256, 512 if backbone_name == "resnet34": - res = models.resnet34(weights=models.ResNet34_Weights.DEFAULT if pretrained else None) + # 构建骨干并按需手动加载权重，便于打印加载摘要 + if pretrained: + res = models.resnet34(weights=None) + self._summarize_pretrained_load(res, models.ResNet34_Weights.DEFAULT, "resnet34") + else: + res = models.resnet34(weights=None) self.backbone = nn.Sequential( res.conv1, res.bn1, res.relu, res.maxpool, res.layer1, res.layer2, res.layer3, res.layer4, @@ -102,14 +107,23 @@ class RoRD(nn.Module): # 选择 layer2/layer3/layer4 作为 C2/C3/C4 c2_ch, c3_ch, c4_ch = 128, 256, 512 elif backbone_name == "efficientnet_b0": - eff = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT if pretrained else None) + if pretrained: + eff = models.efficientnet_b0(weights=None) + self._summarize_pretrained_load(eff, models.EfficientNet_B0_Weights.DEFAULT, "efficientnet_b0") + else: + eff = models.efficientnet_b0(weights=None) self.backbone = eff.features self._backbone_raw = eff out_channels_backbone = 1280 # 选择 features[2]/[3]/[6] 作为 C2/C3/C4（约 24/40/192） c2_ch, c3_ch, c4_ch = 24, 40, 192 else: - vgg16_features = models.vgg16(weights=models.VGG16_Weights.DEFAULT if pretrained else None).features + if pretrained: + vgg = models.vgg16(weights=None) + self._summarize_pretrained_load(vgg, models.VGG16_Weights.DEFAULT, "vgg16") + else: + vgg = models.vgg16(weights=None) + vgg16_features = vgg.features # VGG16 特征各阶段索引（conv & relu 层序列） # relu2_2 索引 8，relu3_3 索引 15，relu4_3 索引 22 self.features = vgg16_features @@ -263,4 +277,33 @@ class RoRD(nn.Module): x = feats[6](x); c4 = x return c2, c3, c4 - raise RuntimeError(f"Unsupported backbone for FPN: {self.backbone_name}") \ No newline at end of file + raise RuntimeError(f"Unsupported backbone for FPN: {self.backbone_name}") + + # --- Utils --- + def _summarize_pretrained_load(self, torch_model: nn.Module, weights_enum, arch_name: str) -> None: + """手动加载 torchvision 预训练权重并打印加载摘要。 + - 使用 strict=False 以兼容可能的键差异，打印 missing/unexpected keys。 + - 输出参数量统计，便于快速核对加载情况。 + """ + try: + state_dict = weights_enum.get_state_dict(progress=False) + except Exception: + # 回退：若权重枚举不支持 get_state_dict，则跳过摘要（通常已在构造器中加载） + print(f"[Pretrained] {arch_name}: skip summary (weights enum lacks get_state_dict)") + return + incompatible = torch_model.load_state_dict(state_dict, strict=False) + total_params = sum(p.numel() for p in torch_model.parameters()) + trainable_params = sum(p.numel() for p in torch_model.parameters() if p.requires_grad) + missing = list(getattr(incompatible, 'missing_keys', [])) + unexpected = list(getattr(incompatible, 'unexpected_keys', [])) + try: + matched = len(state_dict) - len(unexpected) + except Exception: + matched = 0 + print(f"[Pretrained] {arch_name}: ImageNet weights loaded (strict=False)") + print(f" params: total={total_params/1e6:.2f}M, trainable={trainable_params/1e6:.2f}M") + print(f" keys: matched≈{matched} | missing={len(missing)} | unexpected={len(unexpected)}") + if missing and len(missing) <= 10: + print(f" missing: {missing}") + if unexpected and len(unexpected) <= 10: + print(f" unexpected: {unexpected}") \ No newline at end of file diff --git a/tests/benchmark_attention.py b/tests/benchmark_attention.py new file mode 100644 index 0000000..bca441b --- /dev/null +++ b/tests/benchmark_attention.py @@ -0,0 +1,91 @@ +""" +注意力模块 A/B 基准测试 + +目的：在相同骨干与输入下，对比注意力开/关（none/se/cbam）在单尺度与 FPN 前向的耗时差异；可选指定插入位置。 + +示例： + PYTHONPATH=. uv run python tests/benchmark_attention.py --device cpu --image-size 512 --runs 10 --backbone resnet34 --places backbone_high desc_head +""" +from __future__ import annotations + +import argparse +import time +from typing import Dict, List + +import numpy as np +import torch + +from models.rord import RoRD + + +def bench_once(model: torch.nn.Module, x: torch.Tensor, fpn: bool = False) -> float: + if torch.cuda.is_available() and x.is_cuda: + torch.cuda.synchronize() + t0 = time.time() + with torch.inference_mode(): + _ = model(x, return_pyramid=fpn) + if torch.cuda.is_available() and x.is_cuda: + torch.cuda.synchronize() + return (time.time() - t0) * 1000.0 + + +def build_model(backbone: str, attention_type: str, places: List[str], device: torch.device) -> RoRD: + cfg = type("cfg", (), { + "model": type("m", (), { + "backbone": type("b", (), {"name": backbone, "pretrained": False})(), + "attention": type("a", (), {"enabled": attention_type != "none", "type": attention_type, "places": places})(), + })() + })() + model = RoRD(cfg=cfg).to(device) + model.eval() + return model + + +def run_suite(backbone: str, places: List[str], device: torch.device, image_size: int, runs: int) -> List[Dict[str, float]]: + x = torch.randn(1, 3, image_size, image_size, device=device) + results: List[Dict[str, float]] = [] + for attn in ["none", "se", "cbam"]: + model = build_model(backbone, attn, places, device) + # warmup + for _ in range(3): + _ = model(x, return_pyramid=False) + _ = model(x, return_pyramid=True) + # single + t_list_single = [bench_once(model, x, fpn=False) for _ in range(runs)] + # fpn + t_list_fpn = [bench_once(model, x, fpn=True) for _ in range(runs)] + results.append({ + "backbone": backbone, + "attention": attn, + "places": ",".join(places) if places else "-", + "single_ms_mean": float(np.mean(t_list_single)), + "single_ms_std": float(np.std(t_list_single)), + "fpn_ms_mean": float(np.mean(t_list_fpn)), + "fpn_ms_std": float(np.std(t_list_fpn)), + "runs": int(runs), + }) + return results + + +def main(): + parser = argparse.ArgumentParser(description="RoRD 注意力模块 A/B 基准") + parser.add_argument("--backbone", type=str, default="resnet34", choices=["vgg16","resnet34","efficientnet_b0"], help="骨干") + parser.add_argument("--places", nargs="*", default=["backbone_high"], help="插入位置：backbone_high det_head desc_head") + parser.add_argument("--image-size", type=int, default=512, help="输入尺寸") + parser.add_argument("--runs", type=int, default=10, help="重复次数") + parser.add_argument("--device", type=str, default="cpu", help="cuda 或 cpu") + args = parser.parse_args() + + device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu") + results = run_suite(args.backbone, args.places, device, args.image_size, args.runs) + + # 简要打印 + print("\n===== Attention A/B Summary =====") + for r in results: + print(f"{r['backbone']:<14} attn={r['attention']:<5} places={r['places']:<24} " + f"single {r['single_ms_mean']:.2f}±{r['single_ms_std']:.2f} | " + f"fpn {r['fpn_ms_mean']:.2f}±{r['fpn_ms_std']:.2f} ms") + + +if __name__ == "__main__": + main() diff --git a/tests/benchmark_grid.py b/tests/benchmark_grid.py new file mode 100644 index 0000000..441308a --- /dev/null +++ b/tests/benchmark_grid.py @@ -0,0 +1,96 @@ +""" +三维基准对比：Backbone × Attention × (SingleMean / FPNMean) + +示例： + PYTHONPATH=. uv run python tests/benchmark_grid.py --device cpu --image-size 512 --runs 5 \ + --backbones vgg16 resnet34 efficientnet_b0 --attentions none se cbam --places backbone_high +""" +from __future__ import annotations + +import argparse +import json +import time +from typing import Dict, List + +import numpy as np +import torch + +from models.rord import RoRD + + +def bench_once(model: torch.nn.Module, x: torch.Tensor, fpn: bool = False) -> float: + if torch.cuda.is_available() and x.is_cuda: + torch.cuda.synchronize() + t0 = time.time() + with torch.inference_mode(): + _ = model(x, return_pyramid=fpn) + if torch.cuda.is_available() and x.is_cuda: + torch.cuda.synchronize() + return (time.time() - t0) * 1000.0 + + +def build_model(backbone: str, attention: str, places: List[str], device: torch.device) -> RoRD: + cfg = type("cfg", (), { + "model": type("m", (), { + "backbone": type("b", (), {"name": backbone, "pretrained": False})(), + "attention": type("a", (), {"enabled": attention != "none", "type": attention, "places": places})(), + })() + })() + model = RoRD(cfg=cfg).to(device) + model.eval() + return model + + +def run_grid(backbones: List[str], attentions: List[str], places: List[str], device: torch.device, image_size: int, runs: int) -> List[Dict[str, float]]: + x = torch.randn(1, 3, image_size, image_size, device=device) + rows: List[Dict[str, float]] = [] + for bk in backbones: + for attn in attentions: + model = build_model(bk, attn, places, device) + # warmup + for _ in range(3): + _ = model(x, return_pyramid=False) + _ = model(x, return_pyramid=True) + # bench + t_single = [bench_once(model, x, fpn=False) for _ in range(runs)] + t_fpn = [bench_once(model, x, fpn=True) for _ in range(runs)] + rows.append({ + "backbone": bk, + "attention": attn, + "places": ",".join(places) if places else "-", + "single_ms_mean": float(np.mean(t_single)), + "single_ms_std": float(np.std(t_single)), + "fpn_ms_mean": float(np.mean(t_fpn)), + "fpn_ms_std": float(np.std(t_fpn)), + "runs": int(runs), + }) + return rows + + +def main(): + parser = argparse.ArgumentParser(description="三维基准：Backbone × Attention × (Single/FPN)") + parser.add_argument("--backbones", nargs="*", default=["vgg16","resnet34","efficientnet_b0"], help="骨干列表") + parser.add_argument("--attentions", nargs="*", default=["none","se","cbam"], help="注意力列表") + parser.add_argument("--places", nargs="*", default=["backbone_high"], help="插入位置") + parser.add_argument("--image-size", type=int, default=512) + parser.add_argument("--runs", type=int, default=5) + parser.add_argument("--device", type=str, default="cpu") + parser.add_argument("--json-out", type=str, default="benchmark_grid.json") + args = parser.parse_args() + + device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu") + rows = run_grid(args.backbones, args.attentions, args.places, device, args.image_size, args.runs) + + # 打印简表 + print("\n===== Grid Summary (Backbone × Attention) =====") + for r in rows: + print(f"{r['backbone']:<14} attn={r['attention']:<5} places={r['places']:<16} single {r['single_ms_mean']:.2f} | fpn {r['fpn_ms_mean']:.2f} ms") + + # 保存 JSON + with open(args.json_out, 'w') as f: + json.dump(rows, f, indent=2) + print(f"Saved: {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/tools/diffusion/prepare_patch_dataset.py b/tools/diffusion/prepare_patch_dataset.py new file mode 100644 index 0000000..ecedb01 --- /dev/null +++ b/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() diff --git a/tools/diffusion/sample_layouts.py b/tools/diffusion/sample_layouts.py new file mode 100644 index 0000000..0621d82 --- /dev/null +++ b/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() diff --git a/tools/diffusion/train_layout_diffusion.py b/tools/diffusion/train_layout_diffusion.py new file mode 100644 index 0000000..282b7d8 --- /dev/null +++ b/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() diff --git a/tools/generate_synthetic_layouts.py b/tools/generate_synthetic_layouts.py new file mode 100644 index 0000000..8741b24 --- /dev/null +++ b/tools/generate_synthetic_layouts.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +""" +Programmatic synthetic IC layout generator using gdstk. +Generates GDS files with simple standard-cell-like patterns, wires, and vias. +""" +from __future__ import annotations + +import argparse +from pathlib import Path +import random + +import gdstk + + +def build_standard_cell(cell_name: str, rng: random.Random, layer: int = 1, datatype: int = 0) -> gdstk.Cell: + cell = gdstk.Cell(cell_name) + # Basic cell body + w = rng.uniform(0.8, 2.0) + h = rng.uniform(1.6, 4.0) + rect = gdstk.rectangle((0, 0), (w, h), layer=layer, datatype=datatype) + cell.add(rect) + # Poly fingers + nf = rng.randint(1, 4) + pitch = w / (nf + 1) + for i in range(1, nf + 1): + x = i * pitch + poly = gdstk.rectangle((x - 0.05, 0), (x + 0.05, h), layer=layer + 1, datatype=datatype) + cell.add(poly) + # Contact/vias + for i in range(rng.randint(2, 6)): + vx = rng.uniform(0.1, w - 0.1) + vy = rng.uniform(0.1, h - 0.1) + via = gdstk.rectangle((vx - 0.05, vy - 0.05), (vx + 0.05, vy + 0.05), layer=layer + 2, datatype=datatype) + cell.add(via) + return cell + + +def generate_layout(out_path: Path, width: float, height: float, seed: int, rows: int, cols: int, density: float): + rng = random.Random(seed) + lib = gdstk.Library() + top = gdstk.Cell("TOP") + + # Create a few standard cell variants + variants = [build_standard_cell(f"SC_{i}", rng, layer=1) for i in range(4)] + + # Place instances in a grid with random skips based on density + x_pitch = width / cols + y_pitch = height / rows + for r in range(rows): + for c in range(cols): + if rng.random() > density: + continue + cell = rng.choice(variants) + dx = c * x_pitch + rng.uniform(0.0, 0.1 * x_pitch) + dy = r * y_pitch + rng.uniform(0.0, 0.1 * y_pitch) + ref = gdstk.Reference(cell, (dx, dy)) + top.add(ref) + + lib.add(*variants) + lib.add(top) + lib.write_gds(str(out_path)) + + +def main(): + parser = argparse.ArgumentParser(description="Generate synthetic IC layouts (GDS)") + parser.add_argument("--out-dir", type=str, default="data/synthetic/gds") + parser.add_argument("--out_dir", dest="out_dir", type=str, help="Alias of --out-dir") + parser.add_argument("--num-samples", type=int, default=10) + parser.add_argument("--num", dest="num_samples", type=int, help="Alias of --num-samples") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--width", type=float, default=200.0) + parser.add_argument("--height", type=float, default=200.0) + parser.add_argument("--rows", type=int, default=10) + parser.add_argument("--cols", type=int, default=10) + parser.add_argument("--density", type=float, default=0.5) + + args = parser.parse_args() + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + rng = random.Random(args.seed) + for i in range(args.num_samples): + sample_seed = rng.randint(0, 2**31 - 1) + out_path = out_dir / f"chip_{i:06d}.gds" + generate_layout(out_path, args.width, args.height, sample_seed, args.rows, args.cols, args.density) + print(f"[OK] Generated {out_path}") + + +if __name__ == "__main__": + main() diff --git a/tools/layout2png.py b/tools/layout2png.py new file mode 100644 index 0000000..ad21b3b --- /dev/null +++ b/tools/layout2png.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +""" +Batch convert GDS to PNG. + +Priority: +1) Use KLayout in headless batch mode (most accurate view fidelity for IC layouts). +2) Fallback to gdstk(read) -> write SVG -> cairosvg to PNG (no KLayout dependency at runtime). +""" +from __future__ import annotations + +import argparse +from pathlib import Path +import subprocess +import sys +import tempfile + +import cairosvg + + +def klayout_convert(gds_path: Path, png_path: Path, dpi: int, layermap: str | None = None, line_width: int | None = None, bgcolor: str | None = None) -> bool: + """Render using KLayout by invoking a temporary Python macro with paths embedded.""" + # Prepare optional display config code + layer_cfg_code = "" + if layermap: + # layermap format: "LAYER/DATATYPE:#RRGGBB,..." + layer_cfg_code += "lprops = pya.LayerPropertiesNode()\n" + for spec in layermap.split(","): + spec = spec.strip() + if not spec: + continue + try: + ld, color = spec.split(":") + layer_s, datatype_s = ld.split("/") + color = color.strip() + layer_cfg_code += ( + "lp = pya.LayerPropertiesNode()\n" + f"lp.layer = int({int(layer_s)})\n" + f"lp.datatype = int({int(datatype_s)})\n" + f"lp.fill_color = pya.Color.from_string('{color}')\n" + f"lp.frame_color = pya.Color.from_string('{color}')\n" + "lprops.insert(lp)\n" + ) + except Exception: + # Ignore malformed entries + continue + layer_cfg_code += "cv.set_layer_properties(lprops)\n" + + line_width_code = "" + if line_width is not None: + line_width_code = f"cv.set_config('default-draw-line-width', '{int(line_width)}')\n" + + bg_code = "" + if bgcolor: + bg_code = f"cv.set_config('background-color', '{bgcolor}')\n" + + script = f""" +import pya +ly = pya.Layout() +ly.read(r"{gds_path}") +cv = pya.LayoutView() +cv.load_layout(ly, 0) +cv.max_hier_levels = 20 +{bg_code} +{line_width_code} +{layer_cfg_code} +cv.zoom_fit() +cv.save_image(r"{png_path}", {dpi}, 0) +""" + try: + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tf: + tf.write(script) + tf.flush() + macro_path = Path(tf.name) + # Run klayout in batch mode + res = subprocess.run(["klayout", "-zz", "-b", "-r", str(macro_path)], check=False, capture_output=True, text=True) + ok = res.returncode == 0 and png_path.exists() + if not ok: + # Print stderr for visibility when running manually + if res.stderr: + sys.stderr.write(res.stderr) + try: + macro_path.unlink(missing_ok=True) # type: ignore[arg-type] + except Exception: + pass + return ok + except FileNotFoundError: + # klayout command not found + return False + except Exception: + return False + + +def gdstk_fallback(gds_path: Path, png_path: Path, dpi: int) -> bool: + """Fallback path: use gdstk to read GDS and write SVG, then cairosvg to PNG. + Note: This may differ visually from KLayout depending on layers/styles. + """ + try: + import gdstk # local import to avoid import cost when not needed + svg_path = png_path.with_suffix(".svg") + lib = gdstk.read_gds(str(gds_path)) + tops = lib.top_level() + if not tops: + return False + # Combine tops into a single temporary cell for rendering + cell = tops[0] + # gdstk Cell has write_svg in recent versions + try: + cell.write_svg(str(svg_path)) # type: ignore[attr-defined] + except Exception: + # Older gdstk: write_svg available on Library + try: + lib.write_svg(str(svg_path)) # type: ignore[attr-defined] + except Exception: + return False + # Convert SVG to PNG + cairosvg.svg2png(url=str(svg_path), write_to=str(png_path), dpi=dpi) + try: + svg_path.unlink() + except Exception: + pass + return True + except Exception: + return False + + +def main(): + parser = argparse.ArgumentParser(description="Convert GDS files to PNG") + parser.add_argument("--in", dest="in_dir", type=str, required=True, help="Input directory containing .gds files") + parser.add_argument("--out", dest="out_dir", type=str, required=True, help="Output directory to place .png files") + parser.add_argument("--dpi", type=int, default=600, help="Output resolution in DPI for rasterization") + parser.add_argument("--layermap", type=str, default=None, help="Layer color map, e.g. '1/0:#00FF00,2/0:#FF0000'") + parser.add_argument("--line_width", type=int, default=None, help="Default draw line width in pixels for KLayout display") + parser.add_argument("--bgcolor", type=str, default=None, help="Background color, e.g. '#000000' or 'black'") + + args = parser.parse_args() + in_dir = Path(args.in_dir) + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + gds_files = sorted(in_dir.glob("*.gds")) + if not gds_files: + print(f"[WARN] No GDS files found in {in_dir}") + return + + ok_cnt = 0 + for gds in gds_files: + png_path = out_dir / (gds.stem + ".png") + ok = klayout_convert(gds, png_path, args.dpi, layermap=args.layermap, line_width=args.line_width, bgcolor=args.bgcolor) + if not ok: + ok = gdstk_fallback(gds, png_path, args.dpi) + if ok: + ok_cnt += 1 + print(f"[OK] {gds.name} -> {png_path}") + else: + print(f"[FAIL] {gds.name}") + print(f"Done. {ok_cnt}/{len(gds_files)} converted.") + + +if __name__ == "__main__": + main() diff --git a/tools/preview_dataset.py b/tools/preview_dataset.py new file mode 100644 index 0000000..713acfa --- /dev/null +++ b/tools/preview_dataset.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +""" +Quickly preview training pairs (original, transformed, H) from ICLayoutTrainingDataset. +Saves a grid image for visual inspection. +""" +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import torch +from PIL import Image +from torchvision.utils import make_grid, save_image + +from data.ic_dataset import ICLayoutTrainingDataset +from utils.data_utils import get_transform + + +def to_pil(t: torch.Tensor) -> Image.Image: + # input normalized to [-1,1] for 3-channels; invert normalization + x = t.clone() + if x.dim() == 3 and x.size(0) == 3: + x = (x * 0.5) + 0.5 # unnormalize + x = (x * 255.0).clamp(0, 255).byte() + if x.dim() == 3 and x.size(0) == 3: + x = x + elif x.dim() == 3 and x.size(0) == 1: + x = x.repeat(3, 1, 1) + else: + raise ValueError("Unexpected tensor shape") + np_img = x.permute(1, 2, 0).cpu().numpy() + return Image.fromarray(np_img) + + +def main(): + parser = argparse.ArgumentParser(description="Preview dataset samples") + parser.add_argument("--dir", dest="image_dir", type=str, required=True, help="PNG images directory") + parser.add_argument("--out", dest="out_path", type=str, default="preview.png") + parser.add_argument("--n", dest="num", type=int, default=8) + 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() + + transform = get_transform() + ds = ICLayoutTrainingDataset( + args.image_dir, + patch_size=args.patch_size, + transform=transform, + scale_range=(1.0, 1.0), + use_albu=args.use_elastic, + albu_params={"prob": 0.5}, + ) + + images = [] + for i in range(min(args.num, len(ds))): + orig, rot, H = ds[i] + # Stack orig and rot side-by-side for each sample + images.append(orig) + images.append(rot) + + grid = make_grid(torch.stack(images, dim=0), nrow=2, padding=2) + save_image(grid, args.out_path) + print(f"Saved preview to {args.out_path}") + + +if __name__ == "__main__": + main() diff --git a/tools/smoke_test.py b/tools/smoke_test.py new file mode 100644 index 0000000..0e2ae8b --- /dev/null +++ b/tools/smoke_test.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +""" +Minimal smoke test: +1) Generate a tiny synthetic set (num=8) and rasterize to PNG +2) Validate H consistency (n=4, with/without elastic) +3) Run a short training loop (epochs=1-2) to verify end-to-end pipeline +Prints PASS/FAIL with basic stats. +""" +from __future__ import annotations + +import argparse +import subprocess +import os +import sys +from pathlib import Path + + +def run(cmd: list[str]) -> int: + print("[RUN]", " ".join(cmd)) + env = os.environ.copy() + # Ensure project root on PYTHONPATH for child processes + root = Path(__file__).resolve().parents[1] + env["PYTHONPATH"] = f"{root}:{env.get('PYTHONPATH','')}" if env.get("PYTHONPATH") else str(root) + return subprocess.call(cmd, env=env) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Minimal smoke test for E2E pipeline") + parser.add_argument("--root", type=str, default="data/smoke", help="Root dir for smoke test outputs") + parser.add_argument("--config", type=str, default="configs/base_config.yaml") + args = parser.parse_args() + + root = Path(args.root) + gds_dir = root / "gds" + png_dir = root / "png" + gds_dir.mkdir(parents=True, exist_ok=True) + png_dir.mkdir(parents=True, exist_ok=True) + + rc = 0 + + # 1) Generate a tiny set + rc |= run([sys.executable, "tools/generate_synthetic_layouts.py", "--out_dir", gds_dir.as_posix(), "--num", "8", "--seed", "123"]) + if rc != 0: + print("[FAIL] generate synthetic") + sys.exit(2) + + # 2) Rasterize + rc |= run([sys.executable, "tools/layout2png.py", "--in", gds_dir.as_posix(), "--out", png_dir.as_posix(), "--dpi", "600"]) + if rc != 0: + print("[FAIL] layout2png") + sys.exit(3) + + # 3) Validate H (n=4, both no-elastic and elastic) + rc |= run([sys.executable, "tools/validate_h_consistency.py", "--dir", png_dir.as_posix(), "--out", (root/"validate_no_elastic").as_posix(), "--n", "4"]) + rc |= run([sys.executable, "tools/validate_h_consistency.py", "--dir", png_dir.as_posix(), "--out", (root/"validate_elastic").as_posix(), "--n", "4", "--elastic"]) + if rc != 0: + print("[FAIL] validate H") + sys.exit(4) + + # 4) Write back config via synth_pipeline and run short training (1 epoch) + rc |= run([sys.executable, "tools/synth_pipeline.py", "--out_root", root.as_posix(), "--num", "0", "--dpi", "600", "--config", args.config, "--ratio", "0.3", "--enable_elastic", "--no_preview"]) + if rc != 0: + print("[FAIL] synth_pipeline config update") + sys.exit(5) + + # Train 1 epoch to smoke the loop + rc |= run([sys.executable, "train.py", "--config", args.config, "--epochs", "1" ]) + if rc != 0: + print("[FAIL] train 1 epoch") + sys.exit(6) + + print("[PASS] Smoke test completed successfully.") + + +if __name__ == "__main__": + main() diff --git a/tools/synth_pipeline.py b/tools/synth_pipeline.py new file mode 100644 index 0000000..50cf9e8 --- /dev/null +++ b/tools/synth_pipeline.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +""" +One-click synthetic data pipeline: +1) Generate synthetic GDS using tools/generate_synthetic_layouts.py +2) Rasterize GDS to PNG using tools/layout2png.py (KLayout preferred, fallback gdstk+SVG) +3) Preview random training pairs using tools/preview_dataset.py (optional) +4) Validate homography consistency using tools/validate_h_consistency.py (optional) +5) Optionally update a YAML config to enable synthetic mixing and elastic augmentation +""" +from __future__ import annotations + +import argparse +import subprocess +import sys +from pathlib import Path + +from omegaconf import OmegaConf + + +def run_cmd(cmd: list[str]) -> None: + print("[RUN]", " ".join(str(c) for c in cmd)) + res = subprocess.run(cmd) + if res.returncode != 0: + raise SystemExit(f"Command failed with code {res.returncode}: {' '.join(map(str, cmd))}") + + +essential_scripts = { + "gen": Path("tools/generate_synthetic_layouts.py"), + "gds2png": Path("tools/layout2png.py"), + "preview": Path("tools/preview_dataset.py"), + "validate": Path("tools/validate_h_consistency.py"), +} + + +def ensure_scripts_exist() -> None: + missing = [str(p) for p in essential_scripts.values() if not p.exists()] + if missing: + raise SystemExit(f"Missing required scripts: {missing}") + + +def update_config(config_path: Path, png_dir: Path, ratio: float, enable_elastic: bool) -> None: + cfg = OmegaConf.load(config_path) + # Ensure nodes exist + if "synthetic" not in cfg: + cfg.synthetic = {} + cfg.synthetic.enabled = True + cfg.synthetic.png_dir = png_dir.as_posix() + cfg.synthetic.ratio = float(ratio) + + if enable_elastic: + if "augment" not in cfg: + cfg.augment = {} + if "elastic" not in cfg.augment: + cfg.augment.elastic = {} + cfg.augment.elastic.enabled = True + # Don't override numeric params if already present + if "alpha" not in cfg.augment.elastic: + cfg.augment.elastic.alpha = 40 + if "sigma" not in cfg.augment.elastic: + cfg.augment.elastic.sigma = 6 + if "alpha_affine" not in cfg.augment.elastic: + cfg.augment.elastic.alpha_affine = 6 + if "prob" not in cfg.augment.elastic: + cfg.augment.elastic.prob = 0.3 + # Photometric defaults + if "photometric" not in cfg.augment: + cfg.augment.photometric = {"brightness_contrast": True, "gauss_noise": True} + + OmegaConf.save(config=cfg, f=config_path) + print(f"[OK] Config updated: {config_path}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="One-click synthetic data pipeline") + parser.add_argument("--out_root", type=str, default="data/synthetic", help="Root output dir for gds/png/preview") + parser.add_argument("--num", type=int, default=200, help="Number of GDS samples to generate") + parser.add_argument("--dpi", type=int, default=600, help="Rasterization DPI for PNG rendering") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--ratio", type=float, default=0.3, help="Mixing ratio for synthetic data in training") + parser.add_argument("--config", type=str, default="configs/base_config.yaml", help="YAML config to update") + parser.add_argument("--enable_elastic", action="store_true", help="Also enable elastic augmentation in config") + parser.add_argument("--no_preview", action="store_true", help="Skip preview generation") + parser.add_argument("--validate_h", action="store_true", help="Run homography consistency validation on rendered PNGs") + parser.add_argument("--validate_n", type=int, default=6, help="Number of samples for H validation") + parser.add_argument("--diffusion_dir", type=str, default=None, help="Directory of diffusion-generated PNGs to include") + # Rendering style passthrough + parser.add_argument("--layermap", type=str, default=None, help="Layer color map for KLayout, e.g. '1/0:#00FF00,2/0:#FF0000'") + parser.add_argument("--line_width", type=int, default=None, help="Default draw line width for KLayout display") + parser.add_argument("--bgcolor", type=str, default=None, help="Background color for KLayout display") + + args = parser.parse_args() + ensure_scripts_exist() + + out_root = Path(args.out_root) + gds_dir = out_root / "gds" + png_dir = out_root / "png" + gds_dir.mkdir(parents=True, exist_ok=True) + png_dir.mkdir(parents=True, exist_ok=True) + + # 1) Generate GDS + run_cmd([sys.executable, str(essential_scripts["gen"]), "--out_dir", gds_dir.as_posix(), "--num", str(args.num), "--seed", str(args.seed)]) + + # 2) GDS -> PNG + gds2png_cmd = [ + sys.executable, str(essential_scripts["gds2png"]), + "--in", gds_dir.as_posix(), + "--out", png_dir.as_posix(), + "--dpi", str(args.dpi), + ] + if args.layermap: + gds2png_cmd += ["--layermap", args.layermap] + if args.line_width is not None: + gds2png_cmd += ["--line_width", str(args.line_width)] + if args.bgcolor: + gds2png_cmd += ["--bgcolor", args.bgcolor] + run_cmd(gds2png_cmd) + + # 3) Preview (optional) + if not args.no_preview: + preview_path = out_root / "preview.png" + preview_cmd = [sys.executable, str(essential_scripts["preview"]), "--dir", png_dir.as_posix(), "--out", preview_path.as_posix(), "--n", "8"] + if args.enable_elastic: + preview_cmd.append("--elastic") + run_cmd(preview_cmd) + + # 4) Validate homography consistency (optional) + if args.validate_h: + validate_dir = out_root / "validate_h" + validate_cmd = [ + sys.executable, str(essential_scripts["validate"]), + "--dir", png_dir.as_posix(), + "--out", validate_dir.as_posix(), + "--n", str(args.validate_n), + ] + if args.enable_elastic: + validate_cmd.append("--elastic") + run_cmd(validate_cmd) + + # 5) Update YAML config + update_config(Path(args.config), png_dir, args.ratio, args.enable_elastic) + # Include diffusion dir if provided (no automatic sampling here; integration only) + if args.diffusion_dir: + cfg = OmegaConf.load(args.config) + if "synthetic" not in cfg: + cfg.synthetic = {} + if "diffusion" not in cfg.synthetic: + cfg.synthetic.diffusion = {} + cfg.synthetic.diffusion.enabled = True + cfg.synthetic.diffusion.png_dir = Path(args.diffusion_dir).as_posix() + # Keep ratio default at 0 unless user updates later; or reuse a small default like 0.1? Keep 0.0 for safety. + if "ratio" not in cfg.synthetic.diffusion: + cfg.synthetic.diffusion.ratio = 0.0 + OmegaConf.save(config=cfg, f=args.config) + print(f"[OK] Config updated with diffusion_dir: {args.diffusion_dir}") + + print("\n[Done] Synthetic pipeline completed.") + print(f"- GDS: {gds_dir}") + print(f"- PNG: {png_dir}") + if args.diffusion_dir: + print(f"- Diffusion PNGs: {Path(args.diffusion_dir)}") + if not args.no_preview: + print(f"- Preview: {out_root / 'preview.png'}") + if args.validate_h: + print(f"- H validation: {out_root / 'validate_h'}") + print(f"- Updated config: {args.config}") + + +if __name__ == "__main__": + main() diff --git a/tools/validate_h_consistency.py b/tools/validate_h_consistency.py new file mode 100644 index 0000000..edb288c --- /dev/null +++ b/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() diff --git a/train.py b/train.py index 47858b5..5e9b715 100644 --- a/train.py +++ b/train.py @@ -7,7 +7,7 @@ from datetime import datetime from pathlib import Path import torch -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, ConcatDataset, WeightedRandomSampler from torch.utils.tensorboard import SummaryWriter from data.ic_dataset import ICLayoutTrainingDataset @@ -82,25 +82,152 @@ def main(args): transform = get_transform() - dataset = ICLayoutTrainingDataset( + # 读取增强与合成配置 + augment_cfg = cfg.get("augment", {}) + elastic_cfg = augment_cfg.get("elastic", {}) if augment_cfg else {} + use_albu = bool(elastic_cfg.get("enabled", False)) + albu_params = { + "prob": elastic_cfg.get("prob", 0.3), + "alpha": elastic_cfg.get("alpha", 40), + "sigma": elastic_cfg.get("sigma", 6), + "alpha_affine": elastic_cfg.get("alpha_affine", 6), + "brightness_contrast": bool(augment_cfg.get("photometric", {}).get("brightness_contrast", True)) if augment_cfg else True, + "gauss_noise": bool(augment_cfg.get("photometric", {}).get("gauss_noise", True)) if augment_cfg else True, + } + + # 构建真实数据集 + real_dataset = ICLayoutTrainingDataset( data_dir, patch_size=patch_size, transform=transform, scale_range=scale_range, + use_albu=use_albu, + albu_params=albu_params, ) - - logger.info(f"数据集大小: {len(dataset)}") - - # 分割训练集和验证集 - train_size = int(0.8 * len(dataset)) - val_size = len(dataset) - train_size - train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size]) - - logger.info(f"训练集大小: {len(train_dataset)}, 验证集大小: {len(val_dataset)}") + + # 读取合成数据配置（程序化 + 扩散） + syn_cfg = cfg.get("synthetic", {}) + syn_enabled = bool(syn_cfg.get("enabled", False)) + syn_ratio = float(syn_cfg.get("ratio", 0.0)) + syn_dir = syn_cfg.get("png_dir", None) + + syn_dataset = None + if syn_enabled and syn_dir: + syn_dir_path = Path(to_absolute_path(syn_dir, config_dir)) + if syn_dir_path.exists(): + syn_dataset = ICLayoutTrainingDataset( + syn_dir_path.as_posix(), + patch_size=patch_size, + transform=transform, + scale_range=scale_range, + use_albu=use_albu, + albu_params=albu_params, + ) + if len(syn_dataset) == 0: + syn_dataset = None + else: + logger.warning(f"合成数据目录不存在，忽略: {syn_dir_path}") + syn_enabled = False + + # 扩散生成数据配置 + diff_cfg = syn_cfg.get("diffusion", {}) if syn_cfg else {} + diff_enabled = bool(diff_cfg.get("enabled", False)) + diff_ratio = float(diff_cfg.get("ratio", 0.0)) + diff_dir = diff_cfg.get("png_dir", None) + diff_dataset = None + if diff_enabled and diff_dir: + diff_dir_path = Path(to_absolute_path(diff_dir, config_dir)) + if diff_dir_path.exists(): + diff_dataset = ICLayoutTrainingDataset( + diff_dir_path.as_posix(), + patch_size=patch_size, + transform=transform, + scale_range=scale_range, + use_albu=use_albu, + albu_params=albu_params, + ) + if len(diff_dataset) == 0: + diff_dataset = None + else: + logger.warning(f"扩散数据目录不存在，忽略: {diff_dir_path}") + diff_enabled = False + + logger.info( + "真实数据集大小: %d%s%s" % ( + len(real_dataset), + f", 合成(程序)数据集: {len(syn_dataset)}" if syn_dataset else "", + f", 合成(扩散)数据集: {len(diff_dataset)}" if diff_dataset else "", + ) + ) + + # 验证集仅使用真实数据，避免评价受合成样本干扰 + train_size = int(0.8 * len(real_dataset)) + val_size = max(len(real_dataset) - train_size, 1) + real_train_dataset, val_dataset = torch.utils.data.random_split(real_dataset, [train_size, val_size]) + + # 训练集：可与合成数据集合并（程序合成 + 扩散） + datasets = [real_train_dataset] + weights = [] + names = [] + # 收集各源与期望比例 + n_real = len(real_train_dataset) + n_real = max(n_real, 1) + names.append("real") + # 程序合成 + if syn_dataset is not None and syn_enabled and syn_ratio > 0.0: + datasets.append(syn_dataset) + names.append("synthetic") + # 扩散合成 + if diff_dataset is not None and diff_enabled and diff_ratio > 0.0: + datasets.append(diff_dataset) + names.append("diffusion") + + if len(datasets) > 1: + mixed_train_dataset = ConcatDataset(datasets) + # 计算各源样本数 + counts = [len(real_train_dataset)] + if syn_dataset is not None and syn_enabled and syn_ratio > 0.0: + counts.append(len(syn_dataset)) + if diff_dataset is not None and diff_enabled and diff_ratio > 0.0: + counts.append(len(diff_dataset)) + # 期望比例：real = 1 - (syn_ratio + diff_ratio) + target_real = max(0.0, 1.0 - (syn_ratio + diff_ratio)) + target_ratios = [target_real] + if syn_dataset is not None and syn_enabled and syn_ratio > 0.0: + target_ratios.append(syn_ratio) + if diff_dataset is not None and diff_enabled and diff_ratio > 0.0: + target_ratios.append(diff_ratio) + # 构建每个样本的权重 + per_source_weights = [] + for count, ratio in zip(counts, target_ratios): + count = max(count, 1) + per_source_weights.append(ratio / count) + # 展开到每个样本 + weights = [] + idx = 0 + for count, w in zip(counts, per_source_weights): + weights += [w] * count + idx += count + sampler = WeightedRandomSampler(weights, num_samples=len(mixed_train_dataset), replacement=True) + train_dataloader = DataLoader(mixed_train_dataset, batch_size=batch_size, sampler=sampler, num_workers=4) + logger.info( + f"启用混采: real={target_real:.2f}, syn={syn_ratio:.2f}, diff={diff_ratio:.2f}; 总样本={len(mixed_train_dataset)}" + ) + if writer: + writer.add_text( + "dataset/mix", + f"enabled=true, ratios: real={target_real:.2f}, syn={syn_ratio:.2f}, diff={diff_ratio:.2f}; " + f"counts: real_train={len(real_train_dataset)}, syn={len(syn_dataset) if syn_dataset else 0}, diff={len(diff_dataset) if diff_dataset else 0}" + ) + else: + train_dataloader = DataLoader(real_train_dataset, batch_size=batch_size, shuffle=True, num_workers=4) + if writer: + writer.add_text("dataset/mix", f"enabled=false, real_train={len(real_train_dataset)}") + + logger.info(f"训练集大小: {len(train_dataloader.dataset)}, 验证集大小: {len(val_dataset)}") if writer: - writer.add_text("dataset/info", f"train={len(train_dataset)}, val={len(val_dataset)}") - - train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4) + writer.add_text("dataset/info", f"train={len(train_dataloader.dataset)}, val={len(val_dataset)}") + val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4) model = RoRD().cuda() -- 2.49.1