initial commit

data_units.py

@@ -0,0 +1,31 @@
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+
+def layout_transforms():
+    """定义数据增强和预处理"""
+    return transforms.Compose([
+        transforms.Resize((256, 256)),       # 调整尺寸到固定大小
+        transforms.RandomRotation(30),       # 随机旋转（增强方向不变性）
+        transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 彩色抖动（如果使用图像数据）
+        transforms.ToTensor(),               # 转换为张量
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化（如ImageNet均值和方差）
+    ])
+
+def layout_to_tensor(layout_path, target_size=(256, 256)):
+    """将版图转换为标准化张量"""
+    # 实际应用中可能需要解析GDSII/LEF格式，此处简化处理
+    img = Image.open(layout_path).convert('L')  # 灰度化
+    img = img.resize(target_size, resample=Image.BILINEAR)
+    return np.array(img) / 255.0                 # 归一化到[0,1]
+
+def tile_layout(large_layout, block_size=64):
+    """将大版图分割为小块（滑动窗口方式）"""
+    height, width = large_layout.shape
+    stride = block_size // 2                     # 步长设置重叠区域
+    tiles = []
+    for y in range(0, height - block_size +1, stride):
+        for x in range(0, width - block_size +1, stride):
+            tile = large_layout[y:y+block_size, x:x+block_size]
+            tiles.append((x, y, tile))
+    return tiles

inference.py

@@ -0,0 +1,48 @@
+import faiss
+import numpy as np
+import torch
+# 导入 models.rotation_cnn 模块中的 RotationInvariantNet 类
+from models.rotation_cnn import RotationInvariantNet
+
+from models.rotation_cnn import get_rotational_features
+
+# 导入 data_utils 中的 layout_to_tensor 函数（假设该函数存在）
+from data_units import layout_to_tensor  # 如果 data_utils.py 存在此函数
+
+from data_units import tile_layout
+
+def main():
+    # 配置参数（需根据实际调整）
+    block_size = 64           # 分块尺寸
+    target_module_path = "target.png"
+    large_layout_path = "layout_large.png"
+
+    # 加载模型
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = RotationInvariantNet().to(device)
+    model.load_state_dict(torch.load("rotation_cnn.pth"))
+    model.eval()
+
+    # 预处理目标模块与大版图
+    target_tensor = layout_to_tensor(target_module_path, (block_size, block_size))
+    target_feat = get_rotational_features(model, torch.tensor(target_tensor).to(device))
+
+    large_layout = layout_to_tensor(large_layout_path)
+    tiles = tile_layout(large_layout)
+
+    # 构建特征索引（使用Faiss加速）
+    index = faiss.IndexFlatL2(64)  # 特征维度由模型决定
+    features_db = []
+    for (x,y,tile) in tiles:
+        feat = get_rotational_features(model, torch.tensor(tile).to(device))
+        features_db.append(feat)
+    index.add(np.stack(features_db))
+
+    # 检索相似区域
+    D, I = index.search(target_feat[np.newaxis,:], k=10)
+    for idx in I[0]:
+        x,y,_ = tiles[idx]
+        print(f"匹配区域坐标: ({x}, {y}), 相似度: {D[0][idx]}")
+
+if __name__ == "__main__":
+    main()

models/rotation_cnn.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+class RotationInvariantNet(nn.Module):
+    """轻量级旋转不变特征提取网络"""
+    def __init__(self, input_channels=1, num_features=64):
+        super().__init__()
+        self.cnn = nn.Sequential(
+            # 基础卷积层
+            nn.Conv2d(input_channels, 32, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(2),                     # 下采样
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1,1))          # 全局池化获取全局特征
+        )
+        
+    def forward(self, x):
+        features = self.cnn(x)
+        return torch.flatten(features, 1)         # 展平为特征向量
+    
+def get_rotational_features(model, input_image):
+    """计算输入图像所有旋转角度的特征平均值"""
+    rotations = [0, 90, 180, 270]
+    features_list = []
+    for angle in rotations:
+        rotated_img = torch.rot90(input_image, k=angle//90, dims=[2,3])
+        feat = model(rotated_img.unsqueeze(0))
+        features_list.append(feat)
+    return torch.mean(torch.stack(features_list), dim=0).detach().numpy()

train.py

@@ -0,0 +1,99 @@
+import os
+import torch
+from torch import nn, optim
+from torch.utils.data import DataLoader, random_split
+import numpy as np
+from datetime import datetime
+import argparse
+
+# 导入项目模块（根据你的路径调整）
+from models.rotation_cnn import RotationInvariantCNN  # 模型实现
+from data_units import LayoutDataset, layout_transforms  # 数据集和预处理函数
+
+# 设置随机种子（可选）
+torch.manual_seed(42)
+np.random.seed(42)
+
+def main():
+    """训练流程"""
+    # 解析命令行参数
+    parser = argparse.ArgumentParser(description="Train Rotation-Invariant Layout Matcher")
+    parser.add_argument("--data_dir", type=str, default="./data/train/", help="训练数据目录")
+    parser.add_argument("--val_split", type=float, default=0.2, help="验证集比例")
+    parser.add_argument("--batch_size", type=int, default=16, help="批量大小")
+    parser.add_argument("--epochs", type=int, default=50, help="训练轮次")
+    parser.add_argument("--lr", type=float, default=1e-3, help="学习率")
+    parser.add_argument("--model_save_dir", type=str, default="./models/", help="模型保存路径")
+    args = parser.parse_args()
+
+    # 创建输出目录
+    os.makedirs(args.model_save_dir, exist_ok=True)
+
+    # 数据加载
+    dataset = LayoutDataset(root_dir=args.data_dir, transform=layout_transforms())
+    total_samples = len(dataset)
+    val_size = int(total_samples * args.val_split)
+    train_size = total_samples - val_size
+
+    # 划分训练集和验证集
+    train_dataset, val_dataset = random_split(
+        dataset,
+        [train_size, val_size],
+        generator=torch.Generator().manual_seed(42)
+    )
+    
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
+
+    # 初始化模型、损失函数和优化器
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = RotationInvariantCNN().to(device)  # 根据你的模型结构调整参数
+    criterion = nn.CrossEntropyLoss()          # 分类任务示例，根据任务类型选择损失函数
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+    # 训练循环
+    best_val_loss = float("inf")
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        train_loss = 0.0
+        for batch_idx, (data, targets) in enumerate(train_loader):
+            data, targets = data.to(device), targets.to(device)
+            
+            # 前向传播
+            outputs = model(data)
+            loss = criterion(outputs, targets)
+
+            # 反向传播和优化
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+            
+            if (batch_idx + 1) % 10 == 0:
+                print(f"Epoch [{epoch}/{args.epochs}] Batch {batch_idx+1}/{len(train_loader)} Loss: {loss.item():.4f}")
+
+        # 验证
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for data, targets in val_loader:
+                data, targets = data.to(device), targets.to(device)
+                outputs = model(data)
+                loss = criterion(outputs, targets)
+                val_loss += loss.item()
+
+        avg_train_loss = train_loss / len(train_loader)
+        avg_val_loss = val_loss / len(val_loader)
+
+        print(f"Epoch {epoch} - Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
+
+        # 保存最佳模型
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            torch.save(model.state_dict(), os.path.join(args.model_save_dir, f"best_model_{datetime.now().strftime('%Y%m%d%H%M')}.pth"))
+    
+    print("训练完成！")
+
+if __name__ == "__main__":
+    main()