192 lines
7.4 KiB
Python
192 lines
7.4 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torchvision.transforms import GaussianBlur
|
||
|
||
def warp(image, flow): #warp操作
|
||
"""
|
||
基于光流的图像变形函数
|
||
Args:
|
||
image: [B, C, H, W] 输入图像
|
||
flow: [B, 2, H, W] 光流场（x/y方向位移）
|
||
Returns:
|
||
warped: [B, C, H, W] 变形后的图像
|
||
"""
|
||
B, C, H, W = image.shape
|
||
# 生成网格坐标
|
||
grid_x, grid_y = torch.meshgrid(torch.arange(W), torch.arange(H))
|
||
grid = torch.stack((grid_x, grid_y), dim=0).float().to(image.device) # [2,H,W]
|
||
grid = grid.unsqueeze(0).repeat(B,1,1,1) # [B,2,H,W]
|
||
|
||
# 应用光流位移（归一化到[-1,1]）
|
||
new_grid = grid + flow
|
||
new_grid[:,0,:,:] = 2.0 * new_grid[:,0,:,:] / (W-1) - 1.0 # x方向
|
||
new_grid[:,1,:,:] = 2.0 * new_grid[:,1,:,:] / (H-1) - 1.0 # y方向
|
||
new_grid = new_grid.permute(0,2,3,1) # [B,H,W,2]
|
||
|
||
# 双线性插值
|
||
return F.grid_sample(image, new_grid, align_corners=True)
|
||
|
||
# 时序归一化损失计算
|
||
def compute_ctn_loss(G, x, F_content): #公式10
|
||
"""
|
||
计算内容感知时序归一化损失
|
||
Args:
|
||
G: 生成器
|
||
x: 输入红外图像 [B,C,H,W]
|
||
F_content: 生成的光流场 [B,2,H,W]
|
||
"""
|
||
|
||
# 生成可见光图像
|
||
y_fake = G(x) # [B,3,H,W]
|
||
|
||
# 对生成结果应用光流变形
|
||
warped_fake = warp(y_fake, F_content) # [B,3,H,W]
|
||
|
||
# 对输入应用相同光流后生成图像
|
||
warped_x = warp(x, F_content) # [B,C,H,W]
|
||
y_fake_warped = G(warped_x) # [B,3,H,W]
|
||
|
||
# 计算L2损失
|
||
loss = F.mse_loss(warped_fake, y_fake_warped)
|
||
return loss
|
||
|
||
class ContentAwareOptimization(nn.Module):
|
||
def __init__(self, lambda_inc=2.0, eta_ratio=0.4):
|
||
super().__init__()
|
||
self.lambda_inc = lambda_inc # 权重增强系数
|
||
self.eta_ratio = eta_ratio # 选择内容区域的比例
|
||
|
||
def compute_cosine_similarity(self, gradients):
|
||
"""
|
||
计算每个patch梯度与平均梯度的余弦相似度
|
||
Args:
|
||
gradients: [B, N, D] 判别器输出的每个patch的梯度（N=w*h）
|
||
Returns:
|
||
cosine_sim: [B, N] 每个patch的余弦相似度
|
||
"""
|
||
mean_grad = torch.mean(gradients, dim=1, keepdim=True) # [B, 1, D]
|
||
# 计算余弦相似度
|
||
cosine_sim = F.cosine_similarity(gradients, mean_grad, dim=2) # [B, N]
|
||
return cosine_sim
|
||
|
||
def generate_weight_map(self, gradients_real, gradients_fake):
|
||
"""
|
||
生成内容感知权重图
|
||
Args:
|
||
gradients_real: [B, N, D] 真实图像判别器梯度
|
||
gradients_fake: [B, N, D] 生成图像判别器梯度
|
||
Returns:
|
||
weight_real: [B, N] 真实图像权重图
|
||
weight_fake: [B, N] 生成图像权重图
|
||
"""
|
||
# 计算真实图像块的余弦相似度
|
||
cosine_real = self.compute_cosine_similarity(gradients_real) # [B, N] 公式5
|
||
# 计算生成图像块的余弦相似度
|
||
cosine_fake = self.compute_cosine_similarity(gradients_fake) # [B, N]
|
||
|
||
# 选择内容丰富的区域（余弦相似度最低的eta_ratio比例）
|
||
k = int(self.eta_ratio * cosine_real.shape[1])
|
||
|
||
# 对真实图像生成权重图
|
||
_, real_indices = torch.topk(-cosine_real, k, dim=1) # 选择最不相似的区域
|
||
weight_real = torch.ones_like(cosine_real)
|
||
for b in range(cosine_real.shape[0]):
|
||
weight_real[b, real_indices[b]] = self.lambda_inc / (1e-6 + torch.abs(cosine_real[b, real_indices[b]])) #公式6
|
||
|
||
# 对生成图像生成权重图（同理）
|
||
_, fake_indices = torch.topk(-cosine_fake, k, dim=1)
|
||
weight_fake = torch.ones_like(cosine_fake)
|
||
for b in range(cosine_fake.shape[0]):
|
||
weight_fake[b, fake_indices[b]] = self.lambda_inc / (1e-6 + torch.abs(cosine_fake[b, fake_indices[b]]))
|
||
|
||
return weight_real, weight_fake
|
||
|
||
def forward(self, D_real, D_fake, real_scores, fake_scores):
|
||
"""
|
||
计算内容感知对抗损失
|
||
Args:
|
||
D_real: 判别器对真实图像的特征输出 [B, C, H, W]
|
||
D_fake: 判别器对生成图像的特征输出 [B, C, H, W]
|
||
real_scores: 真实图像的判别器预测 [B, N] (N=H*W)
|
||
fake_scores: 生成图像的判别器预测 [B, N]
|
||
Returns:
|
||
loss_co_adv: 内容感知对抗损失
|
||
"""
|
||
B, C, H, W = D_real.shape
|
||
N = H * W
|
||
|
||
# 注册钩子获取梯度
|
||
gradients_real = []
|
||
gradients_fake = []
|
||
|
||
def hook_real(grad):
|
||
gradients_real.append(grad.detach().view(B, N, -1))
|
||
|
||
def hook_fake(grad):
|
||
gradients_fake.append(grad.detach().view(B, N, -1))
|
||
|
||
D_real.register_hook(hook_real)
|
||
D_fake.register_hook(hook_fake)
|
||
|
||
# 计算原始对抗损失以触发梯度计算
|
||
loss_real = torch.mean(torch.log(real_scores + 1e-8))
|
||
loss_fake = torch.mean(torch.log(1 - fake_scores + 1e-8))
|
||
# 添加与 D_real、D_fake 相关的 dummy 项，确保梯度传递
|
||
loss_dummy = 1e-8 * (D_real.sum() + D_fake.sum())
|
||
total_loss = loss_real + loss_fake + loss_dummy
|
||
total_loss.backward(retain_graph=True)
|
||
|
||
# 获取梯度数据
|
||
gradients_real = gradients_real[0] # [B, N, D]
|
||
gradients_fake = gradients_fake[0] # [B, N, D]
|
||
|
||
# 生成权重图
|
||
self.weight_real, self.weight_fake = self.generate_weight_map(gradients_real, gradients_fake)
|
||
|
||
# 应用权重到对抗损失
|
||
loss_co_real = torch.mean(self.weight_real * torch.log(real_scores + 1e-8))
|
||
loss_co_fake = torch.mean(self.weight_fake * torch.log(1 - fake_scores + 1e-8))
|
||
|
||
# 计算并返回最终内容感知对抗损失
|
||
loss_co_adv = -(loss_co_real + loss_co_fake)
|
||
|
||
return loss_co_adv
|
||
|
||
class ContentAwareTemporalNorm(nn.Module):
|
||
def __init__(self, gamma_stride=0.1, kernel_size=21, sigma=5.0):
|
||
super().__init__()
|
||
self.gamma_stride = gamma_stride # 控制整体运动幅度
|
||
self.smoother = GaussianBlur(kernel_size, sigma=sigma) # 高斯平滑层
|
||
|
||
def forward(self, weight_map):
|
||
"""
|
||
生成内容感知光流
|
||
Args:
|
||
weight_map: [B, 1, H, W] 权重图（来自内容感知优化模块）
|
||
Returns:
|
||
F_content: [B, 2, H, W] 生成的光流场（x/y方向位移）
|
||
"""
|
||
B, _, H, W = weight_map.shape
|
||
|
||
# 1. 归一化权重图
|
||
# 保持区域相对强度，同时限制数值范围
|
||
weight_norm = F.normalize(weight_map, p=1, dim=(2,3)) # L1归一化 [B,1,H,W]
|
||
|
||
# 2. 生成高斯噪声（与光流场同尺寸）
|
||
z = torch.randn(B, 2, H, W, device=weight_map.device) # [B,2,H,W]
|
||
|
||
# 3. 合成基础光流
|
||
# 将权重图扩展为2通道（x/y方向共享权重）
|
||
weight_expanded = weight_norm.expand(-1, 2, -1, -1) # [B,2,H,W]
|
||
F_raw = self.gamma_stride * weight_expanded * z # [B,2,H,W] #公式9
|
||
|
||
# 4. 平滑处理（保持结构连续性）
|
||
# 对每个通道独立进行高斯模糊
|
||
F_smooth = self.smoother(F_raw) # [B,2,H,W]
|
||
|
||
# 5. 动态范围调整（可选）
|
||
# 限制光流幅值，避免极端位移
|
||
F_content = torch.tanh(F_smooth) # 缩放到[-1,1]范围
|
||
|
||
return F_content |