236 lines
8.5 KiB
Python
236 lines
8.5 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.utils.data import Dataset, DataLoader
|
||
from torchvision import transforms
|
||
from PIL import Image
|
||
import numpy as np
|
||
import cv2
|
||
import os
|
||
from models.rord import RoRD
|
||
|
||
# 数据集类：生成随机旋转的训练对
|
||
class ICLayoutTrainingDataset(Dataset):
|
||
def __init__(self, image_dir, patch_size=256, transform=None):
|
||
"""
|
||
初始化 IC 版图训练数据集。
|
||
|
||
参数：
|
||
image_dir (str): 存储 PNG 格式 IC 版图图像的目录路径。
|
||
patch_size (int): 裁剪的 patch 大小（默认 256x256）。
|
||
transform (callable, optional): 应用于图像的变换。
|
||
"""
|
||
self.image_dir = image_dir
|
||
self.image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith('.png')]
|
||
self.patch_size = patch_size
|
||
self.transform = transform
|
||
|
||
def __len__(self):
|
||
"""
|
||
返回数据集中的图像数量。
|
||
|
||
返回：
|
||
int: 数据集大小。
|
||
"""
|
||
return len(self.image_paths)
|
||
|
||
def __getitem__(self, index):
|
||
"""
|
||
获取指定索引的训练对（原始 patch、旋转 patch、Homography 矩阵）。
|
||
|
||
参数：
|
||
index (int): 图像索引。
|
||
|
||
返回：
|
||
tuple: (patch, rotated_patch, H_tensor)
|
||
- patch: 原始 patch 张量。
|
||
- rotated_patch: 旋转后的 patch 张量。
|
||
- H_tensor: Homography 矩阵张量。
|
||
"""
|
||
img_path = self.image_paths[index]
|
||
image = Image.open(img_path).convert('L') # 灰度图像
|
||
|
||
# 获取图像大小
|
||
W, H = image.size
|
||
|
||
# 随机选择裁剪的左上角坐标
|
||
x = np.random.randint(0, W - self.patch_size + 1)
|
||
y = np.random.randint(0, H - self.patch_size + 1)
|
||
patch = image.crop((x, y, x + self.patch_size, y + self.patch_size))
|
||
|
||
# 转换为 NumPy 数组
|
||
patch_np = np.array(patch)
|
||
|
||
# 随机旋转角度（0°~360°）
|
||
theta = np.random.uniform(0, 360)
|
||
theta_rad = np.deg2rad(theta)
|
||
cos_theta = np.cos(theta_rad)
|
||
sin_theta = np.sin(theta_rad)
|
||
|
||
# 计算旋转中心（patch 的中心）
|
||
cx = self.patch_size / 2.0
|
||
cy = self.patch_size / 2.0
|
||
|
||
# 计算旋转的齐次矩阵（Homography）
|
||
H = np.array([
|
||
[cos_theta, -sin_theta, cx * (1 - cos_theta) + cy * sin_theta],
|
||
[sin_theta, cos_theta, cy * (1 - cos_theta) - cx * sin_theta],
|
||
[0, 0, 1]
|
||
], dtype=np.float32)
|
||
|
||
# 应用旋转到 patch
|
||
rotated_patch_np = cv2.warpPerspective(patch_np, H, (self.patch_size, self.patch_size))
|
||
|
||
# 转换回 PIL Image
|
||
rotated_patch = Image.fromarray(rotated_patch_np)
|
||
|
||
# 应用变换
|
||
if self.transform:
|
||
patch = self.transform(patch)
|
||
rotated_patch = self.transform(rotated_patch)
|
||
|
||
# 转换 H 为张量
|
||
H_tensor = torch.from_numpy(H).float()
|
||
|
||
return patch, rotated_patch, H_tensor
|
||
|
||
# 特征图变换函数
|
||
def warp_feature_map(feature_map, H_inv):
|
||
"""
|
||
使用逆 Homography 矩阵变换特征图。
|
||
|
||
参数：
|
||
feature_map (torch.Tensor): 输入特征图，形状为 [B, C, H, W]。
|
||
H_inv (torch.Tensor): 逆 Homography 矩阵，形状为 [B, 3, 3]。
|
||
|
||
返回：
|
||
torch.Tensor: 变换后的特征图，形状为 [B, C, H, W]。
|
||
"""
|
||
B, C, H, W = feature_map.size()
|
||
# 生成网格
|
||
grid_y, grid_x = torch.meshgrid(
|
||
torch.linspace(-1, 1, H, device=feature_map.device),
|
||
torch.linspace(-1, 1, W, device=feature_map.device),
|
||
indexing='ij'
|
||
)
|
||
grid = torch.stack((grid_x, grid_y, torch.ones_like(grid_x)), dim=-1) # [H, W, 3]
|
||
grid = grid.unsqueeze(0).expand(B, H, W, 3) # [B, H, W, 3]
|
||
|
||
# 将网格转换为齐次坐标并应用 H_inv
|
||
grid_flat = grid.view(B, -1, 3) # [B, H*W, 3]
|
||
grid_transformed = torch.bmm(grid_flat, H_inv.transpose(1, 2)) # [B, H*W, 3]
|
||
grid_transformed = grid_transformed.view(B, H, W, 3) # [B, H, W, 3]
|
||
grid_transformed = grid_transformed[..., :2] / (grid_transformed[..., 2:3] + 1e-8) # [B, H, W, 2]
|
||
|
||
# 使用 grid_sample 进行变换
|
||
warped_feature = F.grid_sample(feature_map, grid_transformed, align_corners=True)
|
||
return warped_feature
|
||
|
||
# 检测损失函数
|
||
def compute_detection_loss(det_original, det_rotated, H):
|
||
"""
|
||
计算检测损失（MSE），比较原始检测图与旋转检测图（逆变换后）。
|
||
|
||
参数：
|
||
det_original (torch.Tensor): 原始图像的检测图，形状为 [B, 1, H, W]。
|
||
det_rotated (torch.Tensor): 旋转图像的检测图，形状为 [B, 1, H, W]。
|
||
H (torch.Tensor): Homography 矩阵，形状为 [B, 3, 3]。
|
||
|
||
返回：
|
||
torch.Tensor: 检测损失。
|
||
"""
|
||
H_inv = torch.inverse(H) # 计算逆 Homography
|
||
warped_det_rotated = warp_feature_map(det_rotated, H_inv)
|
||
return F.mse_loss(det_original, warped_det_rotated)
|
||
|
||
# 描述子损失函数
|
||
def compute_description_loss(desc_original, desc_rotated, H, margin=1.0):
|
||
"""
|
||
计算描述子损失（三元组损失），基于对应点的描述子。
|
||
|
||
参数：
|
||
desc_original (torch.Tensor): 原始图像的描述子图，形状为 [B, 128, H, W]。
|
||
desc_rotated (torch.Tensor): 旋转图像的描述子图，形状为 [B, 128, H, W]。
|
||
H (torch.Tensor): Homography 矩阵，形状为 [B, 3, 3]。
|
||
margin (float): 三元组损失的边距。
|
||
|
||
返回：
|
||
torch.Tensor: 描述子损失。
|
||
"""
|
||
B, C, H, W = desc_original.size()
|
||
# 随机选择锚点（anchor）
|
||
num_samples = min(100, H * W) # 每张图像采样 100 个点
|
||
idx = torch.randint(0, H * W, (B, num_samples), device=desc_original.device)
|
||
idx_y = idx // W
|
||
idx_x = idx % W
|
||
coords = torch.stack((idx_x.float(), idx_y.float()), dim=-1) # [B, num_samples, 2]
|
||
|
||
# 转换为齐次坐标
|
||
coords_hom = torch.cat((coords, torch.ones(B, num_samples, 1, device=coords.device)), dim=-1) # [B, num_samples, 3]
|
||
coords_transformed = torch.bmm(coords_hom, H.transpose(1, 2)) # [B, num_samples, 3]
|
||
coords_transformed = coords_transformed[..., :2] / (coords_transformed[..., 2:3] + 1e-8) # [B, num_samples, 2]
|
||
|
||
# 归一化到 [-1, 1] 用于 grid_sample
|
||
coords_transformed = coords_transformed / torch.tensor([W/2, H/2], device=coords.device) - 1
|
||
|
||
# 提取锚点和正样本描述子
|
||
anchor = desc_original.view(B, C, -1)[:, :, idx.view(-1)] # [B, 128, num_samples]
|
||
positive = F.grid_sample(desc_rotated, coords_transformed.unsqueeze(2), align_corners=True).squeeze(3) # [B, 128, num_samples]
|
||
|
||
# 随机选择负样本
|
||
neg_idx = torch.randint(0, H * W, (B, num_samples), device=desc_original.device)
|
||
negative = desc_rotated.view(B, C, -1)[:, :, neg_idx.view(-1)] # [B, 128, num_samples]
|
||
|
||
# 三元组损失
|
||
triplet_loss = nn.TripletMarginLoss(margin=margin, p=2)
|
||
loss = triplet_loss(anchor.transpose(1, 2), positive.transpose(1, 2), negative.transpose(1, 2))
|
||
return loss
|
||
|
||
# 定义变换
|
||
transform = transforms.Compose([
|
||
transforms.ToTensor(), # (1, 256, 256)
|
||
transforms.Lambda(lambda x: x.repeat(3, 1, 1)), # (3, 256, 256)
|
||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
|
||
])
|
||
|
||
# 创建数据集和 DataLoader
|
||
dataset = ICLayoutTrainingDataset('path/to/layouts', patch_size=256, transform=transform)
|
||
dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4)
|
||
|
||
# 定义模型
|
||
model = RoRD().cuda()
|
||
|
||
# 定义优化器
|
||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
|
||
|
||
# 训练循环
|
||
num_epochs = 10
|
||
for epoch in range(num_epochs):
|
||
model.train()
|
||
total_loss = 0
|
||
for batch in dataloader:
|
||
original, rotated, H = batch
|
||
original = original.cuda()
|
||
rotated = rotated.cuda()
|
||
H = H.cuda()
|
||
|
||
# 前向传播
|
||
det_original, _, desc_rord_original = model(original)
|
||
det_rotated, _, desc_rord_rotated = model(rotated)
|
||
|
||
# 计算损失
|
||
detection_loss = compute_detection_loss(det_original, det_rotated, H)
|
||
description_loss = compute_description_loss(desc_rord_original, desc_rord_rotated, H)
|
||
total_loss_batch = detection_loss + description_loss
|
||
|
||
# 反向传播
|
||
optimizer.zero_grad()
|
||
total_loss_batch.backward()
|
||
optimizer.step()
|
||
|
||
total_loss += total_loss_batch.item()
|
||
|
||
print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(dataloader):.4f}")
|
||
|
||
# 保存模型
|
||
torch.save(model.state_dict(), 'path/to/save/model.pth') |