File size: 2,787 Bytes
28e6f98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import torch
from torch import nn
import numpy as np
def cc(img1, img2):
eps = torch.finfo(torch.float32).eps
"""Correlation coefficient for (N, C, H, W) image; torch.float32 [0.,1.]."""
N, C, _, _ = img1.shape
img1 = img1.reshape(N, C, -1)
img2 = img2.reshape(N, C, -1)
img1 = img1 - img1.mean(dim=-1, keepdim=True)
img2 = img2 - img2.mean(dim=-1, keepdim=True)
cc = torch.sum(img1 * img2, dim=-1) / (eps + torch.sqrt(torch.sum(
img1 **2, dim=-1)) * torch.sqrt(torch.sum(img2**2, dim=-1)))
cc = torch.clamp(cc, -1., 1.)
return cc.mean()
def gradient_calllback(network):
for name, param in network.named_parameters():
if param.grad is not None:
# print("Gradient of {}: {}".format(name, param.grad.abs().mean()))
if param.grad.abs().mean() == 0:
print("Gradient of {} is 0".format(name))
else:
print("Gradient of {} is None".format(name))
def bright(x, a,b):
# input datatype np.uint8
x = np.array(x, dtype='float')
x = x/(b-a) - 255*a/(b-a)
x[x>255.0] = 255.0
x[x<0.0] = 0.0
x = x.astype(np.uint8)
return x
def trunc(x):
# input datatype float
x[x>255.0] = 255.0
x[x<0.0] = 0.0
return x
def gradient_calllback(network):
for name, param in network.named_parameters():
if param.grad is not None:
if param.grad.abs().mean() == 0:
print("Gradient of {} is 0".format(name))
else:
print("Gradient of {} is None".format(name))
class Frequency_Loss(nn.Module):
def __init__(self):
super(Frequency_Loss, self).__init__()
self.cri = nn.L1Loss()
self.cri_sum = nn.L1Loss(reduction="sum")
def forward(self, x, y, mask=None):
x = torch.fft.fftshift(torch.fft.fft2(x)) # rfft2
y = torch.fft.fftshift(torch.fft.fft2(y))
# def apply_tofre(x_start, mask):
# # B, C, H, W = x_start.shape
# kspace = fftshift(fft2(x_start, norm=None, dim=(-2, -1)), dim=(-2, -1)) # Default: all dimensions
# mask = mask.to(kspace.device)
# return kspace, mask
x_mag = torch.abs(x)
y_mag = torch.abs(y)
x_ang = torch.angle(x)
y_ang = torch.angle(y)
if isinstance(mask, type(None)):
return self.cri(x_mag,y_mag) + self.cri(x_ang, y_ang)
k = (1 - mask.to(x.device)).detach()
# W = x.shape[-1]
# k = k[..., :W // 2 + 1]
k_total = torch.sum(k)
x_mag = x_mag * k
y_mag = y_mag * k
x_ang = x_ang * k
y_ang = y_ang * k
# Compute L1 loss between magnitudes
return self.cri_sum(x_mag, y_mag) / k_total + self.cri_sum(x_ang, y_ang) / k_total
|