import torch
import torch.nn as nn
from skimage.metrics import peak_signal_noise_ratio


dtype = torch.cuda.FloatTensor



def psnr3d(x,y): 

    #y = y.cpu().detach().numpy()
    x = x.cpu().detach().numpy()
    # y = cv2.GaussianBlur(y, (3,3), 0,0)
    ps = 0
    for i in range(x.shape[2]):
        ps += peak_signal_noise_ratio(x[:,:,i], y[:,:,i])
    return ps/x.shape[2]


class permute_in(nn.Module):
    def __init__(self, dim):
        super(permute_in, self).__init__()
        self.dim = dim
    
    def forward(self, x):
        x = x.permute(2,0,1)
        x_in = x.view(1,self.dim,x.shape[1],x.shape[2])
        return x_in
    
class permute_out(nn.Module):
    def __init__(self, dim):
        super(permute_out, self).__init__()
        self.dim = dim
    
    def forward(self, x):
        x = x.view(self.dim,x.shape[2],x.shape[3])
        x_out = x.permute(1,2,0)
        return x_out