Spaces:

Tennineee
/

PDFNet

Running

Upload 84 files

2581217 verified over 1 year ago

22 kB

# Copyright (c) Meta Platforms, Inc. and affiliates.

# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import numpy.random as random
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from math import exp

class SiLogLoss(nn.Module):
    def __init__(self, lambd=0.5, eps=1e-6):
        super().__init__()
        self.lambd = lambd
        self.eps = eps  # 防止log(0)的微小常数

    def forward(self, pred, target):
        # 将输入转换为float32计算关键部分
        pred = pred.float()
        target = target.float()
        
        # 添加eps防止log(0)并提升数值稳定性
        diff_log = torch.log(target + self.eps) - torch.log(pred + self.eps)
        loss = torch.sqrt(
            (diff_log ** 2).mean() - self.lambd * (diff_log.mean() ** 2) + self.eps
        )
        return loss

class IntegrityPriorLoss(nn.Module):
    def __init__(self, epsilon=1e-8):
        super().__init__()

        self.epsilon = epsilon
        self.max_variance = 0.05
        self.max_grad = 0.05

        self.sobel_x = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)
        self.sobel_y = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)
        
        sobel_kernel_x = torch.tensor([[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]], dtype=torch.float32)
        sobel_kernel_y = torch.tensor([[[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]]], dtype=torch.float32)
        
        self.sobel_x.weight.data = sobel_kernel_x
        self.sobel_y.weight.data = sobel_kernel_y
        
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, mask, depth_map, gt):
        #对FP计算与均值的差距，越远loss越高
        #对FN计算与均值的差距，越近loss越高
        py = gt*mask + (1-gt)*(1-mask)
        FP = (1-py)*mask
        FN = (1-py)*gt
        logP = -torch.log(py + self.epsilon)
        diff = (depth_map-((depth_map*gt).sum()/gt.sum()))**2
        FPdiff = (diff)*FP
        FNdiff = (1-diff)*FN
        vareight = (FPdiff+FNdiff)*py    
        variance = logP * vareight  # [B,1]
        variance_loss = torch.mean(variance)

        grad_x = abs(self.sobel_x(depth_map))  # [B,1,H,W]
        grad_y = abs(self.sobel_y(depth_map))  # [B,1,H,W]
        
        masked_grad_x = grad_x * logP
        masked_grad_y = grad_y * logP
        
        grad = (masked_grad_x + masked_grad_y)
        grad_loss = torch.mean(grad)

        total_loss = variance_loss + grad_loss
        return total_loss
def gaussian(window_size, sigma):
    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
    return gauss/gauss.sum()

def create_window(window_size, channel):
    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
    return window

def _ssim(img1, img2, window, window_size, channel, size_average=True):
    mu1 = F.conv2d(img1, window, padding = window_size//2, groups=channel)
    mu2 = F.conv2d(img2, window, padding = window_size//2, groups=channel)

    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1*mu2

    sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=channel) - mu1_sq
    sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=channel) - mu2_sq
    sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=channel) - mu1_mu2

    C1 = 0.01**2
    C2 = 0.03**2

    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))

    if size_average:
        return ssim_map.mean()
    else:
        return ssim_map.mean(1).mean(1).mean(1)

class SSIMLoss(torch.nn.Module):
    def __init__(self, window_size=11, size_average=True):
        super(SSIMLoss, self).__init__()
        self.window_size = window_size
        self.size_average = size_average
        self.channel = 1
        self.window = create_window(window_size, self.channel)

    def forward(self, img1, img2):
        (_, channel, _, _) = img1.size()
        if channel == self.channel and self.window.data.type() == img1.data.type():
            window = self.window
        else:
            window = create_window(self.window_size, channel)
            if img1.is_cuda:
                window = window.cuda(img1.get_device())
            window = window.type_as(img1)
            self.window = window
            self.channel = channel
        return 1 - (1 + _ssim(img1, img2, window, self.window_size, channel, self.size_average)) / 2

def circular_highPassFiltering(img, ratio):
    device = img.device
    batch_size,_,height,width = img.shape
    sigma = (height * (ratio[...,None,None])) / 4
    center_h = height // 2
    center_w = width // 2
    grid_y, grid_x = torch.meshgrid(torch.arange(-center_h, height - center_h),
                                    torch.arange(-center_w, width - center_w))
    grid_y = grid_y[None,None,...].repeat(batch_size, 1, 1, 1).to(device)
    grid_x = grid_x[None,None,...].repeat(batch_size, 1, 1, 1).to(device)
    # 按照二维高斯分布公式计算每个位置的值
    gaussian_values = (1 / (2 * torch.pi * sigma ** 2)) * torch.exp(-(grid_x ** 2 + grid_y ** 2) / (2 * sigma ** 2))
    gmin = gaussian_values.flatten(-2).min(dim=-1)[0][...,None,None]
    gmax = gaussian_values.flatten(-2).max(dim=-1)[0][...,None,None]
    decreasing_matrix = (gaussian_values-gmin) / (gmax-gmin)  # 根据归一化距离计算灰度值
    mask = ((0.5-decreasing_matrix)*100).sigmoid()
    fft = torch.fft.fft2(img)
    fft_shift = torch.fft.fftshift(fft,dim=(2,3))
    fft_shift = torch.mul(fft_shift, mask)
    idft_shift = torch.fft.ifftshift(fft_shift,dim=(2,3))
    ifimg = torch.fft.ifft2(idft_shift)
    ifimg = torch.abs(ifimg)
    ifmin = ifimg.flatten(-2).min(dim=-1)[0][...,None,None]
    ifmax = ifimg.flatten(-2).max(dim=-1)[0][...,None,None]
    ifimg = (ifimg-ifmin) / (ifmax-ifmin)  # 根据归一化距离计算灰度值
    return mask,ifimg

def _upsample_like(src,tar,mode='bilinear'):
    if mode == 'bilinear':
        src = F.upsample(src,size=tar.shape[2:],mode=mode,align_corners=True)
    else:
        src = F.upsample(src,size=tar.shape[2:],mode=mode)
    return src

def _upsample_(src,size,mode='bilinear'):
    if mode == 'bilinear':
        src = F.upsample(src,size=size,mode=mode,align_corners=True)
    else:
        src = F.upsample(src,size=size,mode=mode)
    return src

def patchfy(x,p=4,c=4):
    h = w = x.shape[2] // p
    x = x.reshape(shape=(x.shape[0], c, h, p, w, p))
    x = torch.einsum('nchpwq->nhwpqc', x)
    x = x.reshape(shape=(x.shape[0], h * w, p**2 * c))
    return x
    
def unpatchfy(x,p=4,c=4):
    h = w = round(x.shape[1]**0.5)
    x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
    x = torch.einsum('nhwpqc->nchpwq', x)
    x = x.reshape(shape=(x.shape[0], c, h * p, h * p))
    return x


# def structure_loss(pred, mask):
#     size = 15
#     pad = size//2
#     M_edge = mask
#     N_edge = 1 - M_edge
#     for i in range(2):
#         M_edge = abs(torch.nn.functional.max_pool2d(M_edge, kernel_size=size, stride=1, padding=pad))
#         N_edge = abs(torch.nn.functional.max_pool2d(N_edge, kernel_size=size, stride=1, padding=pad))
#     edge = M_edge + N_edge - 1
#     edge = abs(torch.nn.functional.avg_pool2d(edge, kernel_size=size, stride=1, padding=pad))
#     weit  = 1+2.5*edge

#     wbce  = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
#     wbce  = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))
#     pred  = torch.sigmoid(pred)
#     inter = ((pred * mask) * weit).sum(dim=(2, 3))
#     union = ((pred + mask) * weit).sum(dim=(2, 3))
#     wiou  = 1-(inter+1)/(union-inter+1)
#     return (wbce+wiou).mean()

def structure_loss(pred, mask):
    weit  = 1+5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15)-mask)
    wbce  = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
    wbce  = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))


    pred  = torch.sigmoid(pred)
    inter = ((pred * mask) * weit).sum(dim=(2, 3))

    union = ((pred + mask) * weit).sum(dim=(2, 3))
    wiou  = 1-(inter+1)/(union-inter+1)

    return (wbce+wiou).mean()

def iou_loss(pred, mask):
    eps = 1e-6
    inter = (pred * mask).sum(dim=(2, 3)) #交集
    union = (pred + mask).sum(dim=(2, 3)) - inter #并集-交集
    iou = 1 - (inter + eps) / (union + eps)
    return iou.mean()

def dice_loss(pred, mask):
    eps = 1e-6
    N = pred.size()[0]
    pred_flat = pred.view(N,-1)
    mask_flat = mask.view(N,-1)

    intersection = (pred_flat * mask_flat).sum(1)
    dice_coefficient = (2. * intersection + eps) / (pred_flat.sum(1) + mask_flat.sum(1) + eps)
    dice_loss_value = 1 - dice_coefficient.sum()/N
    return dice_loss_value

class LargeK(nn.Module):
    """ LargeK Block.
    
    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
    """
    def __init__(self, dim):
        super().__init__()
        self.channel_split = nn.Conv2d(dim,dim*3,kernel_size=1)
        self.dwconv1 = nn.Conv2d(dim, dim, kernel_size=7, dilation=2, padding=6, groups=dim) # depthwise conv
        self.dwconv2 = nn.Conv2d(dim, dim, kernel_size=7, dilation=4, padding=12, groups=dim) # depthwise conv
        self.dwconv3 = nn.Conv2d(dim, dim, kernel_size=7, dilation=8, padding=24, groups=dim) # depthwise conv
        self.channel_mix = nn.Conv2d(dim*3,dim,kernel_size=1)

    def forward(self, x):
        x = self.channel_split(x)
        x1,x2,x3 = torch.chunk(x,3,dim=1)
        x1 = self.dwconv1(x1)
        x2 = self.dwconv2(x2)
        x3 = self.dwconv3(x3)
        x = torch.cat([x1,x2,x3],dim=1)
        x = self.channel_mix(x)
        return x

class GANLoss(nn.Module):
    """Define different GAN objectives.
    The GANLoss class abstracts away the need to create the target label tensor
    that has the same size as the input.
    """

    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
        """ Initialize the GANLoss class.
        Parameters:
            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
            target_real_label (bool) - - label for a real image
            target_fake_label (bool) - - label of a fake image
        Note: Do not use sigmoid as the last layer of Discriminator.
        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
        """
        super(GANLoss, self).__init__()
        self.register_buffer('real_label', torch.tensor(target_real_label))
        self.register_buffer('fake_label', torch.tensor(target_fake_label))
        self.gan_mode = gan_mode
        if gan_mode == 'lsgan':
            self.loss = nn.MSELoss()
        elif gan_mode == 'vanilla':
            self.loss = nn.BCEWithLogitsLoss()
        elif gan_mode in ['wgangp']:
            self.loss = None
        else:
            raise NotImplementedError('gan mode %s not implemented' % gan_mode)

    def get_target_tensor(self, prediction, target_is_real):
        """Create label tensors with the same size as the input.
        Parameters:
            prediction (tensor) - - tpyically the prediction from a discriminator
            target_is_real (bool) - - if the ground truth label is for real images or fake images
        Returns:
            A label tensor filled with ground truth label, and with the size of the input
        """

        if target_is_real:
            target_tensor = self.real_label
        else:
            target_tensor = self.fake_label
        return target_tensor.expand_as(prediction)

    def __call__(self, prediction, target_is_real):
        """Calculate loss given Discriminator's output and grount truth labels.
        Parameters:
            prediction (tensor) - - tpyically the prediction output from a discriminator
            target_is_real (bool) - - if the ground truth label is for real images or fake images
        Returns:
            the calculated loss.
        """
        if self.gan_mode in ['lsgan', 'vanilla']:
            target_tensor = self.get_target_tensor(prediction, target_is_real)
            loss = self.loss(prediction, target_tensor)
        elif self.gan_mode == 'wgangp':
            if target_is_real:
                loss = -prediction.mean()
            else:
                loss = prediction.mean()
        return loss

def sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, device):
    # (max_len, 1)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(-1)
    # (output_dim//2)
    ids = torch.arange(0, output_dim // 2, dtype=torch.float)  # 即公式里的i, i的范围是 [0,d/2]
    theta = torch.pow(10000, -2 * ids / output_dim)

    # (max_len, output_dim//2)
    embeddings = position * theta  # 即公式里的：pos / (10000^(2i/d))

    # (max_len, output_dim//2, 2)
    embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)

    # (bs, head, max_len, output_dim//2, 2)
    embeddings = embeddings.repeat((batch_size, nums_head, *([1] * len(embeddings.shape))))  # 在bs维度重复，其他维度都是1不重复

    # (bs, head, max_len, output_dim)
    # reshape后就是：偶数sin, 奇数cos了
    embeddings = torch.reshape(embeddings, (batch_size, nums_head, max_len, output_dim))
    embeddings = embeddings.to(device)
    return embeddings

def RoPE(q, k):
    # q,k: (bs, head, max_len, output_dim)
    use_multi_head = True
    if q.size() == 3 and k.size() == 3:
        use_multi_head = False
        q, k = q[:,None,...], k[:,None,...]
    batch_size = q.shape[0]
    nums_head = q.shape[1]
    max_len = q.shape[2]
    output_dim = q.shape[-1]

    # (bs, head, max_len, output_dim)
    pos_emb = sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, q.device)


    # cos_pos,sin_pos: (bs, head, max_len, output_dim)
    # 看rope公式可知，相邻cos，sin之间是相同的，所以复制一遍。如(1,2,3)变成(1,1,2,2,3,3)
    cos_pos = pos_emb[...,  1::2].repeat_interleave(2, dim=-1)  # 将奇数列信息抽取出来也就是cos 拿出来并复制
    sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1)  # 将偶数列信息抽取出来也就是sin 拿出来并复制

    # q,k: (bs, head, max_len, output_dim)
    q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1)
    q2 = q2.reshape(q.shape)  # reshape后就是正负交替了

    # 更新qw, *对应位置相乘
    q = q * cos_pos + q2 * sin_pos

    k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1)
    k2 = k2.reshape(k.shape)
    # 更新kw, *对应位置相乘
    k = k * cos_pos + k2 * sin_pos
    if not use_multi_head:
        q, k = q[:,0], k[:,0]
    return q, k

class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
        super().__init__()
        self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.w3 = nn.Linear(hidden_size, intermediate_size, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch_size, seq_len, hidden_size)
        # w1(x) -> (batch_size, seq_len, intermediate_size)
        # w3(x) -> (batch_size, seq_len, intermediate_size)
        # w2(*) -> (batch_size, seq_len, hidden_size)
    	return self.w2(F.silu(self.w1(x)) * self.w3(x))

class LayerNorm(nn.Module):
    """ LayerNorm that supports two data formats: channels_last (default) or channels_first. 
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
    with shape (batch_size, channels, height, width).
    """
    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError 
        self.normalized_shape = (normalized_shape, )
    
    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x

class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6, data_format="channels_first") -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.data_format = data_format
  
    def _norm(self, hidden_states):
        if self.data_format == "channels_first":
            variance = hidden_states.pow(2).mean(dim=(1), keepdim=True)  # 在高和宽维度上计算均值
        elif self.data_format == "channels_last":
            variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
        return hidden_states * torch.rsqrt(variance + self.eps)
  
    def forward(self, hidden_states):
        if self.data_format == "channels_first":
            return self.weight[..., None, None] * self._norm(hidden_states.float()).type_as(hidden_states)
        elif self.data_format == "channels_last":
            return self.weight * self._norm(hidden_states.float()).type_as(hidden_states)

class GRN(nn.Module):
    """ GRN (Global Response Normalization) layer
    """
    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(1,2), keepdim=True)
        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x
    
class DUpsampling(nn.Module):
    def __init__(self, inplanes, scale, pad=0):
        super(DUpsampling, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, inplanes* scale * scale, kernel_size=1, padding = pad)
        self.scale = scale
    
    def forward(self, x):
        x = self.conv1(x)
        N, C, H, W = x.size()
        # N, H, W, C
        x_permuted = x.permute(0, 2, 3, 1) 

        # N, H, W*scale, C/scale
        x_permuted = x_permuted.contiguous().view((N, H, W * self.scale, int(C / (self.scale))))

        # N, W*scale,H, C/scale
        x_permuted = x_permuted.permute(0, 2, 1, 3)
        # N, W*scale,H*scale, C/(scale**2)
        x_permuted = x_permuted.contiguous().view((N, W * self.scale, H * self.scale, int(C / (self.scale * self.scale))))

        # N,C/(scale**2),W*scale,H*scale
        x = x_permuted.permute(0, 3, 2, 1)
        
        return x
    
class REsampling(nn.Module):
    def __init__(self, scale):
        super(REsampling, self).__init__()
        self.scale = scale
    
    def forward(self, x):
        N, C, H, W = x.size()
        # N, H, W, C
        x_permuted = x.permute(0, 2, 3, 1) 

        # N, H, W*scale, C/scale
        x_permuted = x_permuted.contiguous().view((N, H, W * self.scale, int(C / (self.scale))))

        # N, W*scale,H, C/scale
        x_permuted = x_permuted.permute(0, 2, 1, 3)
        # N, W*scale,H*scale, C/(scale**2)
        x_permuted = x_permuted.contiguous().view((N, W * self.scale, H * self.scale, int(C / (self.scale * self.scale))))

        # N,C/(scale**2),W*scale,H*scale
        x = x_permuted.permute(0, 3, 2, 1)
        
        return x
    
class Dcrop(nn.Module):
    def __init__(self,inplanes,cropscale=2):
        super(Dcrop, self).__init__()
        self.conv = nn.Conv2d(inplanes*cropscale*cropscale, inplanes*cropscale*cropscale, kernel_size=3, padding = 1)
        self.cropscale = cropscale
    
    def forward(self, x):
        B,C,H,W = x.size()
        x_permuted = x.permute(0, 2, 3, 1) 
        x_permuted = x_permuted.contiguous().view((B, H, W//self.cropscale, C*self.cropscale))
        x_permuted = x_permuted.permute(0, 2, 1, 3)
        x_permuted = x_permuted.contiguous().view((B, W//self.cropscale, H//self.cropscale, C*self.cropscale*self.cropscale))
        x = x_permuted.permute(0, 3, 2, 1)
        x = self.conv(x)+x
        return x
    
def show_gray_images(images, m=8, alpha=3, cmap='coolwarm',save_path=None):
    if len(images.size()) == 2:
        plt.imshow(images, cmap=cmap)
        plt.axis('off')
    else:
        n, h, w = images.shape
        if n == 1:
            plt.imshow(images[0], cmap=cmap)
            plt.axis('off')
        else:
            if m > n: m = n
            num_rows = (n + m - 1) // m
            fig, axes = plt.subplots(num_rows, m, figsize=(m * 2*alpha, num_rows * 2*alpha))
            plt.subplots_adjust(wspace=0.05, hspace=0.05)
            for i in range(num_rows):
                for j in range(m):
                    idx = i*m + j
                    if m == 1 or num_rows == 1:
                        axes[idx].imshow(images[idx], cmap=cmap)
                        axes[idx].axis('off')
                    elif idx < n:
                        axes[i, j].imshow(images[idx], cmap=cmap)
                        axes[i, j].axis('off')
    if save_path is not None:
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()
    

	# Copyright (c) Meta Platforms, Inc. and affiliates.

	# All rights reserved.

	# This source code is licensed under the license found in the
	# LICENSE file in the root directory of this source tree.


	import numpy.random as random
	import matplotlib.pyplot as plt
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.autograd import Variable
	from math import exp

	class SiLogLoss(nn.Module):
	def __init__(self, lambd=0.5, eps=1e-6):
	super().__init__()
	self.lambd = lambd
	self.eps = eps # 防止log(0)的微小常数

	def forward(self, pred, target):
	# 将输入转换为float32计算关键部分
	pred = pred.float()
	target = target.float()

	# 添加eps防止log(0)并提升数值稳定性
	diff_log = torch.log(target + self.eps) - torch.log(pred + self.eps)
	loss = torch.sqrt(
	(diff_log ** 2).mean() - self.lambd * (diff_log.mean() ** 2) + self.eps
	)
	return loss

	class IntegrityPriorLoss(nn.Module):
	def __init__(self, epsilon=1e-8):
	super().__init__()

	self.epsilon = epsilon
	self.max_variance = 0.05
	self.max_grad = 0.05

	self.sobel_x = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)
	self.sobel_y = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)

	sobel_kernel_x = torch.tensor([[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]], dtype=torch.float32)
	sobel_kernel_y = torch.tensor([[[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]]], dtype=torch.float32)

	self.sobel_x.weight.data = sobel_kernel_x
	self.sobel_y.weight.data = sobel_kernel_y

	for param in self.parameters():
	param.requires_grad = False

	def forward(self, mask, depth_map, gt):
	#对FP计算与均值的差距，越远loss越高
	#对FN计算与均值的差距，越近loss越高
	py = gtmask + (1-gt)(1-mask)
	FP = (1-py)*mask
	FN = (1-py)*gt
	logP = -torch.log(py + self.epsilon)
	diff = (depth_map-((depth_mapgt).sum()/gt.sum()))*2
	FPdiff = (diff)*FP
	FNdiff = (1-diff)*FN
	vareight = (FPdiff+FNdiff)*py
	variance = logP * vareight # [B,1]
	variance_loss = torch.mean(variance)

	grad_x = abs(self.sobel_x(depth_map)) # [B,1,H,W]
	grad_y = abs(self.sobel_y(depth_map)) # [B,1,H,W]

	masked_grad_x = grad_x * logP
	masked_grad_y = grad_y * logP

	grad = (masked_grad_x + masked_grad_y)
	grad_loss = torch.mean(grad)

	total_loss = variance_loss + grad_loss
	return total_loss
	def gaussian(window_size, sigma):
	gauss = torch.Tensor([exp(-(x - window_size//2)*2/float(2sigma**2)) for x in range(window_size)])
	return gauss/gauss.sum()

	def create_window(window_size, channel):
	_1D_window = gaussian(window_size, 1.5).unsqueeze(1)
	_2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
	window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
	return window

	def _ssim(img1, img2, window, window_size, channel, size_average=True):
	mu1 = F.conv2d(img1, window, padding = window_size//2, groups=channel)
	mu2 = F.conv2d(img2, window, padding = window_size//2, groups=channel)

	mu1_sq = mu1.pow(2)
	mu2_sq = mu2.pow(2)
	mu1_mu2 = mu1*mu2

	sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=channel) - mu1_sq
	sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=channel) - mu2_sq
	sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=channel) - mu1_mu2

	C1 = 0.01**2
	C2 = 0.03**2

	ssim_map = ((2mu1_mu2 + C1)(2sigma12 + C2))/((mu1_sq + mu2_sq + C1)(sigma1_sq + sigma2_sq + C2))

	if size_average:
	return ssim_map.mean()
	else:
	return ssim_map.mean(1).mean(1).mean(1)

	class SSIMLoss(torch.nn.Module):
	def __init__(self, window_size=11, size_average=True):
	super(SSIMLoss, self).__init__()
	self.window_size = window_size
	self.size_average = size_average
	self.channel = 1
	self.window = create_window(window_size, self.channel)

	def forward(self, img1, img2):
	(_, channel, _, _) = img1.size()
	if channel == self.channel and self.window.data.type() == img1.data.type():
	window = self.window
	else:
	window = create_window(self.window_size, channel)
	if img1.is_cuda:
	window = window.cuda(img1.get_device())
	window = window.type_as(img1)
	self.window = window
	self.channel = channel
	return 1 - (1 + _ssim(img1, img2, window, self.window_size, channel, self.size_average)) / 2

	def circular_highPassFiltering(img, ratio):
	device = img.device
	batch_size,_,height,width = img.shape
	sigma = (height * (ratio[...,None,None])) / 4
	center_h = height // 2
	center_w = width // 2
	grid_y, grid_x = torch.meshgrid(torch.arange(-center_h, height - center_h),
	torch.arange(-center_w, width - center_w))
	grid_y = grid_y[None,None,...].repeat(batch_size, 1, 1, 1).to(device)
	grid_x = grid_x[None,None,...].repeat(batch_size, 1, 1, 1).to(device)
	# 按照二维高斯分布公式计算每个位置的值
	gaussian_values = (1 / (2 * torch.pi * sigma ** 2)) * torch.exp(-(grid_x 2 + grid_y 2) / (2 * sigma ** 2))
	gmin = gaussian_values.flatten(-2).min(dim=-1)[0][...,None,None]
	gmax = gaussian_values.flatten(-2).max(dim=-1)[0][...,None,None]
	decreasing_matrix = (gaussian_values-gmin) / (gmax-gmin) # 根据归一化距离计算灰度值
	mask = ((0.5-decreasing_matrix)*100).sigmoid()
	fft = torch.fft.fft2(img)
	fft_shift = torch.fft.fftshift(fft,dim=(2,3))
	fft_shift = torch.mul(fft_shift, mask)
	idft_shift = torch.fft.ifftshift(fft_shift,dim=(2,3))
	ifimg = torch.fft.ifft2(idft_shift)
	ifimg = torch.abs(ifimg)
	ifmin = ifimg.flatten(-2).min(dim=-1)[0][...,None,None]
	ifmax = ifimg.flatten(-2).max(dim=-1)[0][...,None,None]
	ifimg = (ifimg-ifmin) / (ifmax-ifmin) # 根据归一化距离计算灰度值
	return mask,ifimg

	def _upsample_like(src,tar,mode='bilinear'):
	if mode == 'bilinear':
	src = F.upsample(src,size=tar.shape[2:],mode=mode,align_corners=True)
	else:
	src = F.upsample(src,size=tar.shape[2:],mode=mode)
	return src

	def _upsample_(src,size,mode='bilinear'):
	if mode == 'bilinear':
	src = F.upsample(src,size=size,mode=mode,align_corners=True)
	else:
	src = F.upsample(src,size=size,mode=mode)
	return src

	def patchfy(x,p=4,c=4):
	h = w = x.shape[2] // p
	x = x.reshape(shape=(x.shape[0], c, h, p, w, p))
	x = torch.einsum('nchpwq->nhwpqc', x)
	x = x.reshape(shape=(x.shape[0], h * w, p*2 c))
	return x

	def unpatchfy(x,p=4,c=4):
	h = w = round(x.shape[1]**0.5)
	x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
	x = torch.einsum('nhwpqc->nchpwq', x)
	x = x.reshape(shape=(x.shape[0], c, h * p, h * p))
	return x


	# def structure_loss(pred, mask):
	# size = 15
	# pad = size//2
	# M_edge = mask
	# N_edge = 1 - M_edge
	# for i in range(2):
	# M_edge = abs(torch.nn.functional.max_pool2d(M_edge, kernel_size=size, stride=1, padding=pad))
	# N_edge = abs(torch.nn.functional.max_pool2d(N_edge, kernel_size=size, stride=1, padding=pad))
	# edge = M_edge + N_edge - 1
	# edge = abs(torch.nn.functional.avg_pool2d(edge, kernel_size=size, stride=1, padding=pad))
	# weit = 1+2.5*edge

	# wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
	# wbce = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))
	# pred = torch.sigmoid(pred)
	# inter = ((pred * mask) * weit).sum(dim=(2, 3))
	# union = ((pred + mask) * weit).sum(dim=(2, 3))
	# wiou = 1-(inter+1)/(union-inter+1)
	# return (wbce+wiou).mean()

	def structure_loss(pred, mask):
	weit = 1+5*torch.abs(F.avg_pool2d(mask, kernel_size=31, stride=1, padding=15)-mask)
	wbce = F.binary_cross_entropy_with_logits(pred, mask, reduction='none')
	wbce = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))


	pred = torch.sigmoid(pred)
	inter = ((pred * mask) * weit).sum(dim=(2, 3))

	union = ((pred + mask) * weit).sum(dim=(2, 3))
	wiou = 1-(inter+1)/(union-inter+1)

	return (wbce+wiou).mean()

	def iou_loss(pred, mask):
	eps = 1e-6
	inter = (pred * mask).sum(dim=(2, 3)) #交集
	union = (pred + mask).sum(dim=(2, 3)) - inter #并集-交集
	iou = 1 - (inter + eps) / (union + eps)
	return iou.mean()

	def dice_loss(pred, mask):
	eps = 1e-6
	N = pred.size()[0]
	pred_flat = pred.view(N,-1)
	mask_flat = mask.view(N,-1)

	intersection = (pred_flat * mask_flat).sum(1)
	dice_coefficient = (2. * intersection + eps) / (pred_flat.sum(1) + mask_flat.sum(1) + eps)
	dice_loss_value = 1 - dice_coefficient.sum()/N
	return dice_loss_value

	class LargeK(nn.Module):
	""" LargeK Block.

	Args:
	dim (int): Number of input channels.
	drop_path (float): Stochastic depth rate. Default: 0.0
	"""
	def __init__(self, dim):
	super().__init__()
	self.channel_split = nn.Conv2d(dim,dim*3,kernel_size=1)
	self.dwconv1 = nn.Conv2d(dim, dim, kernel_size=7, dilation=2, padding=6, groups=dim) # depthwise conv
	self.dwconv2 = nn.Conv2d(dim, dim, kernel_size=7, dilation=4, padding=12, groups=dim) # depthwise conv
	self.dwconv3 = nn.Conv2d(dim, dim, kernel_size=7, dilation=8, padding=24, groups=dim) # depthwise conv
	self.channel_mix = nn.Conv2d(dim*3,dim,kernel_size=1)

	def forward(self, x):
	x = self.channel_split(x)
	x1,x2,x3 = torch.chunk(x,3,dim=1)
	x1 = self.dwconv1(x1)
	x2 = self.dwconv2(x2)
	x3 = self.dwconv3(x3)
	x = torch.cat([x1,x2,x3],dim=1)
	x = self.channel_mix(x)
	return x

	class GANLoss(nn.Module):
	"""Define different GAN objectives.

	The GANLoss class abstracts away the need to create the target label tensor
	that has the same size as the input.
	"""

	def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
	""" Initialize the GANLoss class.

	Parameters:
	gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
	target_real_label (bool) - - label for a real image
	target_fake_label (bool) - - label of a fake image

	Note: Do not use sigmoid as the last layer of Discriminator.
	LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
	"""
	super(GANLoss, self).__init__()
	self.register_buffer('real_label', torch.tensor(target_real_label))
	self.register_buffer('fake_label', torch.tensor(target_fake_label))
	self.gan_mode = gan_mode
	if gan_mode == 'lsgan':
	self.loss = nn.MSELoss()
	elif gan_mode == 'vanilla':
	self.loss = nn.BCEWithLogitsLoss()
	elif gan_mode in ['wgangp']:
	self.loss = None
	else:
	raise NotImplementedError('gan mode %s not implemented' % gan_mode)

	def get_target_tensor(self, prediction, target_is_real):
	"""Create label tensors with the same size as the input.

	Parameters:
	prediction (tensor) - - tpyically the prediction from a discriminator
	target_is_real (bool) - - if the ground truth label is for real images or fake images

	Returns:
	A label tensor filled with ground truth label, and with the size of the input
	"""

	if target_is_real:
	target_tensor = self.real_label
	else:
	target_tensor = self.fake_label
	return target_tensor.expand_as(prediction)

	def __call__(self, prediction, target_is_real):
	"""Calculate loss given Discriminator's output and grount truth labels.

	Parameters:
	prediction (tensor) - - tpyically the prediction output from a discriminator
	target_is_real (bool) - - if the ground truth label is for real images or fake images

	Returns:
	the calculated loss.
	"""
	if self.gan_mode in ['lsgan', 'vanilla']:
	target_tensor = self.get_target_tensor(prediction, target_is_real)
	loss = self.loss(prediction, target_tensor)
	elif self.gan_mode == 'wgangp':
	if target_is_real:
	loss = -prediction.mean()
	else:
	loss = prediction.mean()
	return loss

	def sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, device):
	# (max_len, 1)
	position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(-1)
	# (output_dim//2)
	ids = torch.arange(0, output_dim // 2, dtype=torch.float) # 即公式里的i, i的范围是 [0,d/2]
	theta = torch.pow(10000, -2 * ids / output_dim)

	# (max_len, output_dim//2)
	embeddings = position * theta # 即公式里的：pos / (10000^(2i/d))

	# (max_len, output_dim//2, 2)
	embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)

	# (bs, head, max_len, output_dim//2, 2)
	embeddings = embeddings.repeat((batch_size, nums_head, ([1] len(embeddings.shape)))) # 在bs维度重复，其他维度都是1不重复

	# (bs, head, max_len, output_dim)
	# reshape后就是：偶数sin, 奇数cos了
	embeddings = torch.reshape(embeddings, (batch_size, nums_head, max_len, output_dim))
	embeddings = embeddings.to(device)
	return embeddings

	def RoPE(q, k):
	# q,k: (bs, head, max_len, output_dim)
	use_multi_head = True
	if q.size() == 3 and k.size() == 3:
	use_multi_head = False
	q, k = q[:,None,...], k[:,None,...]
	batch_size = q.shape[0]
	nums_head = q.shape[1]
	max_len = q.shape[2]
	output_dim = q.shape[-1]

	# (bs, head, max_len, output_dim)
	pos_emb = sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, q.device)


	# cos_pos,sin_pos: (bs, head, max_len, output_dim)
	# 看rope公式可知，相邻cos，sin之间是相同的，所以复制一遍。如(1,2,3)变成(1,1,2,2,3,3)
	cos_pos = pos_emb[..., 1::2].repeat_interleave(2, dim=-1) # 将奇数列信息抽取出来也就是cos 拿出来并复制
	sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1) # 将偶数列信息抽取出来也就是sin 拿出来并复制

	# q,k: (bs, head, max_len, output_dim)
	q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1)
	q2 = q2.reshape(q.shape) # reshape后就是正负交替了

	# 更新qw, *对应位置相乘
	q = q * cos_pos + q2 * sin_pos

	k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1)
	k2 = k2.reshape(k.shape)
	# 更新kw, *对应位置相乘
	k = k * cos_pos + k2 * sin_pos
	if not use_multi_head:
	q, k = q[:,0], k[:,0]
	return q, k

	class SwiGLU(nn.Module):
	def __init__(self, hidden_size: int, intermediate_size: int) -> None:
	super().__init__()
	self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False)
	self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False)
	self.w3 = nn.Linear(hidden_size, intermediate_size, bias=False)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	# x: (batch_size, seq_len, hidden_size)
	# w1(x) -> (batch_size, seq_len, intermediate_size)
	# w3(x) -> (batch_size, seq_len, intermediate_size)
	# w2(*) -> (batch_size, seq_len, hidden_size)
	return self.w2(F.silu(self.w1(x)) * self.w3(x))

	class LayerNorm(nn.Module):
	""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
	The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
	shape (batch_size, height, width, channels) while channels_first corresponds to inputs
	with shape (batch_size, channels, height, width).
	"""
	def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
	super().__init__()
	self.weight = nn.Parameter(torch.ones(normalized_shape))
	self.bias = nn.Parameter(torch.zeros(normalized_shape))
	self.eps = eps
	self.data_format = data_format
	if self.data_format not in ["channels_last", "channels_first"]:
	raise NotImplementedError
	self.normalized_shape = (normalized_shape, )

	def forward(self, x):
	if self.data_format == "channels_last":
	return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
	elif self.data_format == "channels_first":
	u = x.mean(1, keepdim=True)
	s = (x - u).pow(2).mean(1, keepdim=True)
	x = (x - u) / torch.sqrt(s + self.eps)
	x = self.weight[:, None, None] * x + self.bias[:, None, None]
	return x

	class RMSNorm(nn.Module):
	def __init__(self, hidden_size: int, eps: float = 1e-6, data_format="channels_first") -> None:
	super().__init__()
	self.eps = eps
	self.weight = nn.Parameter(torch.ones(hidden_size))
	self.data_format = data_format

	def _norm(self, hidden_states):
	if self.data_format == "channels_first":
	variance = hidden_states.pow(2).mean(dim=(1), keepdim=True) # 在高和宽维度上计算均值
	elif self.data_format == "channels_last":
	variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
	return hidden_states * torch.rsqrt(variance + self.eps)

	def forward(self, hidden_states):
	if self.data_format == "channels_first":
	return self.weight[..., None, None] * self._norm(hidden_states.float()).type_as(hidden_states)
	elif self.data_format == "channels_last":
	return self.weight * self._norm(hidden_states.float()).type_as(hidden_states)

	class GRN(nn.Module):
	""" GRN (Global Response Normalization) layer
	"""
	def __init__(self, dim):
	super().__init__()
	self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
	self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))

	def forward(self, x):
	Gx = torch.norm(x, p=2, dim=(1,2), keepdim=True)
	Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
	return self.gamma * (x * Nx) + self.beta + x

	class DUpsampling(nn.Module):
	def __init__(self, inplanes, scale, pad=0):
	super(DUpsampling, self).__init__()
	self.conv1 = nn.Conv2d(inplanes, inplanes* scale * scale, kernel_size=1, padding = pad)
	self.scale = scale

	def forward(self, x):
	x = self.conv1(x)
	N, C, H, W = x.size()
	# N, H, W, C
	x_permuted = x.permute(0, 2, 3, 1)

	# N, H, W*scale, C/scale
	x_permuted = x_permuted.contiguous().view((N, H, W * self.scale, int(C / (self.scale))))

	# N, W*scale,H, C/scale
	x_permuted = x_permuted.permute(0, 2, 1, 3)
	# N, Wscale,Hscale, C/(scale**2)
	x_permuted = x_permuted.contiguous().view((N, W * self.scale, H * self.scale, int(C / (self.scale * self.scale))))

	# N,C/(scale*2),Wscale,H*scale
	x = x_permuted.permute(0, 3, 2, 1)

	return x

	class REsampling(nn.Module):
	def __init__(self, scale):
	super(REsampling, self).__init__()
	self.scale = scale

	def forward(self, x):
	N, C, H, W = x.size()
	# N, H, W, C
	x_permuted = x.permute(0, 2, 3, 1)

	# N, H, W*scale, C/scale
	x_permuted = x_permuted.contiguous().view((N, H, W * self.scale, int(C / (self.scale))))

	# N, W*scale,H, C/scale
	x_permuted = x_permuted.permute(0, 2, 1, 3)
	# N, Wscale,Hscale, C/(scale**2)
	x_permuted = x_permuted.contiguous().view((N, W * self.scale, H * self.scale, int(C / (self.scale * self.scale))))

	# N,C/(scale*2),Wscale,H*scale
	x = x_permuted.permute(0, 3, 2, 1)

	return x

	class Dcrop(nn.Module):
	def __init__(self,inplanes,cropscale=2):
	super(Dcrop, self).__init__()
	self.conv = nn.Conv2d(inplanescropscalecropscale, inplanescropscalecropscale, kernel_size=3, padding = 1)
	self.cropscale = cropscale

	def forward(self, x):
	B,C,H,W = x.size()
	x_permuted = x.permute(0, 2, 3, 1)
	x_permuted = x_permuted.contiguous().view((B, H, W//self.cropscale, C*self.cropscale))
	x_permuted = x_permuted.permute(0, 2, 1, 3)
	x_permuted = x_permuted.contiguous().view((B, W//self.cropscale, H//self.cropscale, Cself.cropscaleself.cropscale))
	x = x_permuted.permute(0, 3, 2, 1)
	x = self.conv(x)+x
	return x

	def show_gray_images(images, m=8, alpha=3, cmap='coolwarm',save_path=None):
	if len(images.size()) == 2:
	plt.imshow(images, cmap=cmap)
	plt.axis('off')
	else:
	n, h, w = images.shape
	if n == 1:
	plt.imshow(images[0], cmap=cmap)
	plt.axis('off')
	else:
	if m > n: m = n
	num_rows = (n + m - 1) // m
	fig, axes = plt.subplots(num_rows, m, figsize=(m * 2alpha, num_rows 2*alpha))
	plt.subplots_adjust(wspace=0.05, hspace=0.05)
	for i in range(num_rows):
	for j in range(m):
	idx = i*m + j
	if m == 1 or num_rows == 1:
	axes[idx].imshow(images[idx], cmap=cmap)
	axes[idx].axis('off')
	elif idx < n:
	axes[i, j].imshow(images[idx], cmap=cmap)
	axes[i, j].axis('off')
	if save_path is not None:
	plt.savefig(save_path)
	plt.close()
	else:
	plt.show()