初始化web demo.

app.py

@@ -0,0 +1,119 @@
+import gradio as gr
+from PIL import Image
+import torch
+from torchvision import transforms
+from models.models_x import *
+import torchvision_x_functional as TF_x
+import torchvision.transforms.functional as TF
+from torchvision import transforms
+import cv2
+from timm.models.hub import download_cached_file
+
+
+cuda = True if torch.cuda.is_available() else False
+Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
+trans = transforms.ToTensor()
+
+
+LUT0 = Generator3DLUT_identity()
+LUT1 = Generator3DLUT_zero()
+LUT2 = Generator3DLUT_zero()
+classifier = Classifier()
+trilinear_ = Tritri()
+if cuda:
+    LUT0 = LUT0.cuda()
+    LUT1 = LUT1.cuda()
+    LUT2 = LUT2.cuda()
+    classifier = classifier.cuda()
+
+# Load pretrained models
+cache = download_cached_file('https://drive.google.com/uc?export=download&id=1tzeECo1m4MBqvfLv4H4SQ7by4YMEP17H',
+                             check_hash=False, progress=True)
+LUTs = torch.load(cache, map_location=torch.device('cpu'))
+LUT0.load_state_dict(LUTs["0"])
+LUT1.load_state_dict(LUTs["1"])
+LUT2.load_state_dict(LUTs["2"])
+LUT0.eval()
+LUT1.eval()
+LUT2.eval()
+
+cache = download_cached_file('https://drive.google.com/uc?export=download&id=1rQ_p3NMRFxZ52MOYj0jPewYtD3JQTJGi',
+                             check_hash=False, progress=True)
+classifier.load_state_dict(torch.load(cache, map_location=torch.device('cpu')))
+classifier.eval()
+
+
+XLUT0 = Generator3DLUT_identity()
+XLUT1 = Generator3DLUT_zero()
+XLUT2 = Generator3DLUT_zero()
+Xclassifier = Classifier()
+Xtrilinear_ = Tritri()
+if cuda:
+    XLUT0 = XLUT0.cuda()
+    XLUT1 = XLUT1.cuda()
+    XLUT2 = XLUT2.cuda()
+    Xclassifier = Xclassifier.cuda()
+
+# Load pretrained models
+cache = download_cached_file('https://drive.google.com/uc?export=download&id=1ossTzgbgpZL4Jy5uhiRJDGfCWw9vOv0c',
+                             check_hash=False, progress=True)
+XLUTs = torch.load(cache, map_location=torch.device('cpu'))
+XLUT0.load_state_dict(XLUTs["0"])
+XLUT1.load_state_dict(XLUTs["1"])
+XLUT2.load_state_dict(XLUTs["2"])
+XLUT0.eval()
+XLUT1.eval()
+XLUT2.eval()
+
+cache = download_cached_file('https://drive.google.com/uc?export=download&id=1279CoaqQZK-eK83283MERoRxtRbIgRew',
+                             check_hash=False, progress=True)
+Xclassifier.load_state_dict(torch.load(cache, map_location=torch.device('cpu')))
+Xclassifier.eval()
+
+
+def generate_LUT(img):
+    pred = classifier(img).squeeze()
+
+    LUT = pred[0] * LUT0.LUT + pred[1] * LUT1.LUT + pred[2] * LUT2.LUT  # + pred[3] * LUT3.LUT + pred[4] * LUT4.LUT
+
+    return LUT
+
+def generate_XLUT(img):
+    pred = Xclassifier(img).squeeze()
+
+    XLUT = pred[0] * XLUT0.LUT + pred[1] * XLUT1.LUT + pred[2] * XLUT2.LUT  # + pred[3] * LUT3.LUT + pred[4] * LUT4.LUT
+
+    return XLUT
+
+
+def inference(ori_image, models_n):
+    with torch.no_grad():
+        if models_n == 'sRGB':
+            # img = Image.open(ori_image)
+            # img = TF.to_tensor(img).type(Tensor)
+            img = trans(ori_image)
+            img = img.unsqueeze(0)
+            LUT = generate_LUT(img)
+            result = trilinear_(LUT, img)
+            result = result.permute(0, 3, 1, 2)
+            ndarr = result.squeeze().mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(ndarr)
+        elif models_n == 'XYZ':
+            img = trans(ori_image)
+            img = img.unsqueeze(0)
+            XLUT = generate_XLUT(img)
+            result = Xtrilinear_(XLUT, img)
+            result = result.permute(0, 3, 1, 2)
+            ndarr = result.squeeze().mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
+            im = Image.fromarray(ndarr)
+    return im
+
+
+inputs = [gr.inputs.Image(type='pil', label='待增强图片'),
+          gr.inputs.Radio(choices=['sRGB', 'XYZ'], type="value", default="sRGB", label="图片色彩空间")]
+outputs = [gr.outputs.Image(type='pil', label='增强后图片')]
+
+title = '基于LUT的图像增强演示'
+
+gr.Interface(inference, inputs, outputs, title=title, allow_flagging= 'never',
+             examples=[['./examples/example.jpg', 'sRGB']]).launch(enable_queue=True)

models/models_x.py

@@ -0,0 +1,329 @@
+from doctest import OutputChecker
+from turtle import forward
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+import torch
+import numpy as np
+import math
+
+from models.trilinear_test import bing_lut_trilinearInterplt,Tritri
+
+from re import I
+import time
+from PIL import Image
+###########################################
+# use this module for  pytorch 1.x,together with trilinear_cpp
+###########################################
+
+
+def weights_init_normal_classifier(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        torch.nn.init.xavier_normal_(m.weight.data)
+
+    elif classname.find("BatchNorm2d") != -1 or classname.find("InstanceNorm2d") != -1:
+        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
+        torch.nn.init.constant_(m.bias.data, 0.0)
+
+class resnet18_224(nn.Module):
+
+    def __init__(self, out_dim=5, aug_test=False):
+        super(resnet18_224, self).__init__()
+
+        self.aug_test = aug_test
+        net = models.resnet18(pretrained=True)
+        # self.mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda()
+        # self.std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda()
+
+        self.upsample = nn.Upsample(size=(224,224),mode='bilinear')
+        net.fc = nn.Linear(512, out_dim)
+        self.model = net
+
+
+    def forward(self, x):
+
+        x = self.upsample(x)
+        if self.aug_test:
+            # x = torch.cat((x, torch.rot90(x, 1, [2, 3]), torch.rot90(x, 3, [2, 3])), 0)
+            x = torch.cat((x, torch.flip(x, [3])), 0)
+        f = self.model(x)
+
+        return f
+
+##############################
+#        Discriminator
+##############################
+
+
+def discriminator_block(in_filters, out_filters, normalization=False):
+    """Returns downsampling layers of each discriminator block"""
+    layers = [nn.Conv2d(in_filters, out_filters, 3, stride=2, padding=1)]
+    layers.append(nn.LeakyReLU(0.2))
+    if normalization:
+        layers.append(nn.InstanceNorm2d(out_filters, affine=True))
+        #layers.append(nn.BatchNorm2d(out_filters))
+
+    return layers
+
+class Discriminator(nn.Module):
+    def __init__(self, in_channels=3):
+        super(Discriminator, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Upsample(size=(256,256),mode='bilinear'),
+            nn.Conv2d(3, 16, 3, stride=2, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.InstanceNorm2d(16, affine=True),
+            *discriminator_block(16, 32),
+            *discriminator_block(32, 64),
+            *discriminator_block(64, 128),
+            *discriminator_block(128, 128),
+            #*discriminator_block(128, 128),
+            nn.Conv2d(128, 1, 8, padding=0)
+        )
+
+    def forward(self, img_input):
+        return self.model(img_input)
+
+class Classifier(nn.Module):
+    def __init__(self, in_channels=3):
+        super(Classifier, self).__init__()
+
+        self.model = nn.Sequential(
+            # nn.Downsample(size=(256,256),mode='bilinear'), 
+            nn.Upsample(size=(256,256),mode='bilinear'),            #original
+
+            nn.Conv2d(3, 16, 3, stride=2, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.InstanceNorm2d(16, affine=True),
+            *discriminator_block(16, 32, normalization=True),
+            *discriminator_block(32, 64, normalization=True),
+            *discriminator_block(64, 128, normalization=True),
+            *discriminator_block(128, 128),
+            #*discriminator_block(128, 128, normalization=True),
+            nn.Dropout(p=0.5),
+            nn.Conv2d(128, 3, 8, padding=0),
+        )
+        
+
+    def forward(self, img_input):
+        return self.model(img_input)
+    
+
+class Classifier_unpaired(nn.Module):
+    def __init__(self, in_channels=3):
+        super(Classifier_unpaired, self).__init__()
+
+        self.model = nn.Sequential(
+            nn.Upsample(size=(256,256),mode='bilinear'),
+            nn.Conv2d(3, 16, 3, stride=2, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.InstanceNorm2d(16, affine=True),
+            *discriminator_block(16, 32),
+            *discriminator_block(32, 64),
+            *discriminator_block(64, 128),
+            *discriminator_block(128, 128),
+            #*discriminator_block(128, 128),
+            nn.Conv2d(128, 3, 8, padding=0),
+        )
+
+    def forward(self, img_input):
+        return self.model(img_input)
+
+
+class Generator3DLUT_identity(nn.Module):
+    def __init__(self, dim=33):
+        super(Generator3DLUT_identity, self).__init__()
+        if dim == 33:
+            file = open("IdentityLUT33.txt", 'r')
+        elif dim == 64:
+            file = open("IdentityLUT64.txt", 'r')
+        lines = file.readlines()
+        buffer = np.zeros((3,dim,dim,dim), dtype=np.float32)
+
+        for i in range(0,dim):
+            for j in range(0,dim):
+                for k in range(0,dim):
+                    n = i * dim*dim + j * dim + k
+                    x = lines[n].split()
+                    buffer[0,i,j,k] = float(x[0])
+                    buffer[1,i,j,k] = float(x[1])
+                    buffer[2,i,j,k] = float(x[2])
+        self.LUT = nn.Parameter(torch.from_numpy(buffer).requires_grad_(True))
+        self.TrilinearInterpolation = Tritri()
+        # self.trilinearItp = bing_lut_trilinearInterplt()
+
+
+    def forward(self, x):
+        _, output = self.TrilinearInterpolation(self.LUT, x)
+        # output = self.trilinearItp(self.LUT,x)
+
+        #self.LUT, output = self.TrilinearInterpolation(self.LUT, x)
+        return output
+
+class Generator3DLUT_zero(nn.Module):
+    def __init__(self, dim=33):
+        super(Generator3DLUT_zero, self).__init__()
+
+        self.LUT = torch.zeros(3,dim,dim,dim, dtype=torch.float)
+        self.LUT = nn.Parameter(torch.tensor(self.LUT))
+        self.TrilinearInterpolation = Tritri()
+        # self.trilinearItp = bing_lut_trilinearInterplt()
+
+    def forward(self, x):
+        _, output = self.TrilinearInterpolation(self.LUT, x)
+        # output = self.trilinearItp(self.LUT,x)
+
+        return output
+
+class LUT_all(nn.Module):
+    def __init__(self,
+        path_LUT="saved_models/LUTs/paired/fiveK_480p_3LUT_sm_1e-4_mn_10_sRGB/LUTs_399.pth",
+        path_classifier="saved_models/LUTs/paired/fiveK_480p_3LUT_sm_1e-4_mn_10_sRGB/classifier_399.pth") -> None:        
+        super(LUT_all,self).__init__()
+        self.classifier=Classifier()
+        self.classifier.load_state_dict(torch.load(path_classifier))
+        
+        self.LUT0 = Generator3DLUT_identity()
+        self.LUT1 = Generator3DLUT_zero()
+        self.LUT2 = Generator3DLUT_zero()
+        LUTs = torch.load(path_LUT)
+        self.LUT0.load_state_dict(LUTs["0"])
+        self.LUT1.load_state_dict(LUTs["1"])
+        self.LUT2.load_state_dict(LUTs["2"])
+        # self.trilinear_ = TrilinearInterpolation() 
+        # self.trilinear_ = bing_lut_trilinearInterplt()
+        self.trilinear_=Tritri()
+
+    def forward(self,img):
+        pred = self.classifier(img).squeeze()
+
+        # #numpy squeeze方法去掉矩阵中维度为1的维度，返回np.ndarray
+        # LUT = pred[0] * self.LUT0.LUT 
+        LUT = pred[0] * self.LUT0.LUT + pred[1] * self.LUT1.LUT + pred[2] * self.LUT2.LUT 
+        output = self.trilinear_(LUT, img)
+        # _,output = self.trilinear_(LUT, img)
+        return output
+        # return LUT
+
+
+
+# class TrilinearInterpolationFunction(torch.autograd.Function):
+#     @staticmethod
+#     def forward(ctx, lut, x):
+        
+#         x = x.contiguous()
+
+#         output = x.new(x.size())
+#         dim = lut.size()[-1]
+#         shift = dim ** 3
+#         binsize = 1.000001 / (dim-1)
+#         W = x.size(2)
+#         H = x.size(3)
+#         batch = x.size(0)
+#         #trilinear这个包是作者自己实现的
+#         assert 1 == trilinear.forward(lut, 
+#                                       x, 
+#                                       output,
+#                                       dim, 
+#                                       shift, 
+#                                       binsize, 
+#                                       W, 
+#                                       H, 
+#                                       batch)
+
+#         int_package = torch.IntTensor([dim, shift, W, H, batch])
+#         float_package = torch.FloatTensor([binsize])
+#         variables = [lut, x, int_package, float_package]
+        
+#         ctx.save_for_backward(*variables)
+        
+#         return lut, output
+    
+#     @staticmethod
+#     def backward(ctx, lut_grad, x_grad):
+        
+#         lut, x, int_package, float_package = ctx.saved_variables
+#         dim, shift, W, H, batch = int_package
+#         dim, shift, W, H, batch = int(dim), int(shift), int(W), int(H), int(batch)
+#         binsize = float(float_package[0])
+            
+#         assert 1 == trilinear.backward(x, 
+#                                        x_grad, 
+#                                        lut_grad,
+#                                        dim, 
+#                                        shift, 
+#                                        binsize, 
+#                                        W, 
+#                                        H, 
+#                                        batch)
+#         return lut_grad, x_grad
+
+
+# class TrilinearInterpolation(torch.nn.Module):
+#     def __init__(self):
+#         super(TrilinearInterpolation, self).__init__()
+
+#     def forward(self, lut, x):
+#         return TrilinearInterpolationFunction.apply(lut, x)
+
+
+class TV_3D(nn.Module):
+    def __init__(self, dim=33):
+        super(TV_3D,self).__init__()
+
+        self.weight_r = torch.ones(3,dim,dim,dim-1, dtype=torch.float)
+        self.weight_r[:,:,:,(0,dim-2)] *= 2.0
+        self.weight_g = torch.ones(3,dim,dim-1,dim, dtype=torch.float)
+        self.weight_g[:,:,(0,dim-2),:] *= 2.0
+        self.weight_b = torch.ones(3,dim-1,dim,dim, dtype=torch.float)
+        self.weight_b[:,(0,dim-2),:,:] *= 2.0
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, LUT):
+
+        dif_r = LUT.LUT[:,:,:,:-1] - LUT.LUT[:,:,:,1:]
+        dif_g = LUT.LUT[:,:,:-1,:] - LUT.LUT[:,:,1:,:]
+        dif_b = LUT.LUT[:,:-1,:,:] - LUT.LUT[:,1:,:,:]
+        tv = torch.mean(torch.mul((dif_r ** 2),self.weight_r)) + torch.mean(torch.mul((dif_g ** 2),self.weight_g)) + torch.mean(torch.mul((dif_b ** 2),self.weight_b))
+
+        mn = torch.mean(self.relu(dif_r)) + torch.mean(self.relu(dif_g)) + torch.mean(self.relu(dif_b))
+
+        return tv, mn
+
+
+##new by bing##
+if __name__=='__main__':
+    def img_process_256(img):   
+        # 将PIL类型的图片文件（mode=RGB size=3840x2160，三通道）转换为tensor，tensor维度是[N,C,H,W](即[1,3,256,256])
+        img=img.resize((256,256))
+        trans=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
+        img = trans(img)    
+        img = torch.unsqueeze(img,0) # 填充一维
+        print("img",img.size())
+        # # 将其由HWC格式改成NCHW格式，N=1    
+        # img=np.array(img)
+        return img
+
+    def img_process_4k(img):   
+        # 将PIL类型的图片文件（mode=RGB size=3840x2160，三通道）转换为tensor，tensor维度是[N,C,H,W](即[1,3,256,256])
+        trans=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
+        img = trans(img)    
+        img = torch.unsqueeze(img,0) # 填充一维
+        print("img",img.size())
+        # # 将其由HWC格式改成NCHW格式，N=1    
+        # img=np.array(img)
+        return img
+
+
+    img_ori=Image.open("/home/elle/bing/proj/code/download-4k-img/picture/%s" % ("X4_Animal2_BIC_g_03.png"))
+    img=img_process_256(img_ori)
+    img_4k=img_process_4k(img_ori)
+    model=LUT_all()    
+
+    out=model(img_4k)
+    print(out)
+

models/trilinear_test.py

@@ -0,0 +1,608 @@
+from re import A
+import time
+from turtle import width
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+##new####
+# https://github.com/tedyhabtegebrial/PyTorch-Trilinear-Interpolation
+class TrilinearIntepolation(nn.Module):
+    """TrilinearIntepolation in PyTorch."""
+
+    def __init__(self):
+        super(TrilinearIntepolation, self).__init__()
+
+    def sample_at_integer_locs(self, input_feats, index_tensor):
+        assert input_feats.ndimension()==5, 'input_feats should be of shape [Batch,F,D,Height,Width]'
+        assert index_tensor.ndimension()==4, 'index_tensor should be of shape [Batch,Height,Width,3]'
+        # first sample pixel locations using nearest neighbour interpolation
+        batch_size, num_chans, num_d, height, width = input_feats.shape
+        grid_height, grid_width = index_tensor.shape[1],index_tensor.shape[2]
+
+        xy_grid = index_tensor[..., 0:2]
+        # 0：2是包括0但是不包括2的，因此取出来的是最后一个维度的0维和1维
+        xy_grid[..., 0] = xy_grid[..., 0] - ((width-1.0)/2.0)
+        xy_grid[..., 0] = xy_grid[..., 0] / ((width-1.0)/2.0)
+        xy_grid[..., 1] = xy_grid[..., 1] - ((height-1.0)/2.0)
+        xy_grid[..., 1] = xy_grid[..., 1] / ((height-1.0)/2.0)
+        xy_grid = torch.clamp(xy_grid, min=-1.0, max=1.0)
+        #clamp限制每个元素的最大值和最小值
+        sampled_in_2d = F.grid_sample(input=input_feats.view(batch_size, num_chans*num_d, height, width),
+                                        grid=xy_grid, mode='nearest').view(batch_size, num_chans, num_d, grid_height, grid_width)
+        # grid_sample双线性插值https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html?highlight=grid_sample#torch.nn.functional.grid_sample
+        # view函数https://blog.csdn.net/york1996/article/details/81949843
+        z_grid = index_tensor[..., 2].view(batch_size, 1, 1, grid_height, grid_width)
+        z_grid = z_grid.long().clamp(min=0, max=num_d-1)
+        # .long()将张量转换为int64类型
+        z_grid = z_grid.expand(batch_size,num_chans, 1, grid_height, grid_width)
+        # expand对原张量中维度为1的维度进行扩展 https://blog.csdn.net/weixin_42782150/article/details/108615706
+        # 本例中是使用expand对dim=1的维度进行扩展，扩展成num_chans
+        sampled_in_3d = sampled_in_2d.gather(2, z_grid).squeeze(2)
+        return sampled_in_3d
+
+
+    def forward(self, input_feats, sampling_grid):
+        assert input_feats.ndimension()==5, 'input_feats should be of shape [B,F,D,H,W]'
+        assert sampling_grid.ndimension()==4, 'sampling_grid should be of shape [B,H,W,3]'
+        batch_size, num_chans, num_d, height, width = input_feats.shape
+        grid_height, grid_width = sampling_grid.shape[1],sampling_grid.shape[2]
+        # make sure sampling grid lies between -1, 1
+        sampling_grid = torch.clamp(sampling_grid, min=-1.0, max=1.0)
+        # map to 0,1
+        sampling_grid = (sampling_grid+1)/2.0
+        # Scale grid to floating point pixel locations
+        scaling_factor = torch.FloatTensor([width-1.0, height-1.0, num_d-1.0]).to(input_feats.device).view(1, 1, 1, 3)
+        sampling_grid = scaling_factor*sampling_grid
+        # Now sampling grid is between [0, w-1; 0,h-1; 0,d-1]
+        x, y, z = torch.split(sampling_grid, split_size_or_sections=1, dim=3)
+        #这个(x,y,z)是输入的浮点数（在这篇文章中是每个像素点的rgb值）
+        #这个(x0,y0,z0)是输入的浮点数向下取整
+        #把sampling_grid维度是3的那个维度切成每份大小为1
+        x_0, y_0, z_0 = torch.split(sampling_grid.floor(), split_size_or_sections=1, dim=3)
+        x_1, y_1, z_1 = x_0+1.0, y_0+1.0, z_0+1.0
+        u, v, w = x-x_0, y-y_0, z-z_0
+        print("v:",x_0,y_0,z_0)
+        print("s:",x_0.size(),y_0.size(),z_0.size())
+        print("size,cat",torch.cat([x_0, y_0, z_0],dim=3).size())
+        u, v, w = map(lambda x:x.view(batch_size, 1, grid_height, grid_width).expand(
+                                    batch_size, num_chans, grid_height, grid_width),  [u, v, w])
+        c_000 = self.sample_at_integer_locs(input_feats, torch.cat([x_0, y_0, z_0], dim=3))
+        # torch.cat 函数目的： 在给定维度上对输入的张量序列seq 进行连接操作。
+        c_001 = self.sample_at_integer_locs(input_feats, torch.cat([x_0, y_0, z_1], dim=3))
+        c_010 = self.sample_at_integer_locs(input_feats, torch.cat([x_0, y_1, z_0], dim=3))
+        c_011 = self.sample_at_integer_locs(input_feats, torch.cat([x_0, y_1, z_1], dim=3))
+        c_100 = self.sample_at_integer_locs(input_feats, torch.cat([x_1, y_0, z_0], dim=3))
+        c_101 = self.sample_at_integer_locs(input_feats, torch.cat([x_1, y_0, z_1], dim=3))
+        c_110 = self.sample_at_integer_locs(input_feats, torch.cat([x_1, y_1, z_0], dim=3))
+        c_111 = self.sample_at_integer_locs(input_feats, torch.cat([x_1, y_1, z_1], dim=3))
+        c_xyz = (1.0-u)*(1.0-v)*(1.0-w)*c_000 + \
+                (1.0-u)*(1.0-v)*(w)*c_001 + \
+                (1.0-u)*(v)*(1.0-w)*c_010 + \
+                (1.0-u)*(v)*(w)*c_011 + \
+                (u)*(1.0-v)*(1.0-w)*c_100 + \
+                (u)*(1.0-v)*(w)*c_101 + \
+                (u)*(v)*(1.0-w)*c_110 + \
+                (u)*(v)*(w)*c_111
+        return c_xyz
+# class bing_lut_trilinearInterplt(nn.Module):
+
+#     def __init__(self):
+#         super(bing_lut_trilinearInterplt, self).__init__()
+    
+#     def test(self,LUT,img_input):
+#         # batch_size, num_chans, height, width = img_input.shape
+#         # grid_height, grid_width = LUT.shape[1],LUT.shape[2]
+#         grid_in=img_input.transpose(1,2).transpose(2,3)
+#         # 原本img_input NCHW,改成 NHWC
+#         xy_grid=grid_in[...,0:2]
+#         yz_grid=grid_in[...,1:3]
+#         #只取3通道中的第0和第1通道（0：2不含2）
+#         input_LUT=LUT[:,:,0,:]
+#         input_LUT_ori=input_LUT.squeeze(2)
+#         # LUT[33,33,33,3]->[33,33,3],把dim=2的数据丢掉了
+#         input_LUT=input_LUT_ori[...,0:2]
+#         input_LUT2=input_LUT_ori[...,1:]
+#         print("input_LUT2.size()",input_LUT2.size())
+#         # LUT[33,33,2]
+#         input_LUT=input_LUT.transpose(1,2).transpose(0,1)
+#         input_LUT2=input_LUT2.transpose(1,2).transpose(0,1)
+#         # LUT[2,33,33]
+#         input_LUT=input_LUT.unsqueeze(0)
+#         input_LUT2=input_LUT2.unsqueeze(0)
+#         print(input_LUT.size())
+#         print(input_LUT2.size())
+#         print(grid_in.size())
+#         sampled_in_2d = F.grid_sample(input=input_LUT,grid=xy_grid, mode='nearest')
+#                                         # .view(batch_size, num_chans, num_d, grid_height, grid_width)
+#         sampled_in_2d_2 = F.grid_sample(input=input_LUT2,grid=yz_grid, mode='nearest')
+#                                         # .view(batch_size, num_chans, num_d, grid_height, grid_width)
+        
+#         # print("sampled_in_2d.size()",sampled_in_2d.size())
+#         # print("sampled_in_2d.size()",sampled_in_2d_2.size())
+#         # # [1,2,2160,3840]
+#         # print("ss")
+#         # print(sampled_in_2d.size())
+#         # print(sampled_in_2d_2.size())
+#         res=torch.cat([sampled_in_2d,sampled_in_2d_2[:,1:,:,:]],dim=1)
+#         print(res.size())
+#         return res
+#         # z_grid = grid_in[..., 2]
+#         # print(z_grid.size())
+#         # # [1,2160,3840]
+#         # print("sss")
+
+
+
+#     def gen_Cout_ijk(self,LUT,x_i,y_i,z_i):
+#     # def gen_Cout_ijk(LUT,x_i,y_i,z_i,channel=3):
+#         # LUT size [3,33,33,33]
+#         # x_i,y_i,z_i size [1,1,2160,3840]
+#         # N=batch_size
+#         #img_input.size()=[1,3,2160,3840]\
+#         # LUT.size()=[3,33,33,33]        
+#         # assert LUT.ndimension()==4, 'LUT should be of shape [C,M,M,M](M=33)'
+#         channel=3
+#         batch_size,_,height,width=x_i.size()
+#         print(batch_size,height,width)
+#         output=torch.zeros([batch_size,channel,height,width])
+#         # 设置输出大小为[1,3,2160,3840]
+#         if batch_size==1:
+#             # x_i=x_i.view(height*width)
+#             # y_i=y_i.view(height*width)
+#             # z_i=z_i.view(height*width)
+#             x_i=x_i.view(height*width).long()
+#             y_i=y_i.view(height*width).long()
+#             z_i=z_i.view(height*width).long()
+#             # x_i=x_i.view(1, height*width)
+#             # y_i=y_i.view(1, height*width)
+#             # z_i=z_i.view(1, height*width)
+#             # 2维tensor，[1, 2160*3840]
+#             # xyz_i=torch.cat([x_i,y_i,z_i],dim=0)
+#             # # xyz_i 2维tensor，[3, 2160*3840]
+
+#             # print("xyz_i.size()",xyz_i.size())
+#         else:
+#             print("error:batch size must be 1")
+#         for i in range(height*width):
+#             h_index=int(i/width)
+#             w_index=int(i%width)
+#             # print(h_index)
+#             # print(w_index)
+#             # print(x_i.size())
+#             # print(batch_size)
+#             # print(output.size())
+#             # print(output[0,0,h_index,w_index])
+#             if(i%10000==0):
+#                 print(i)
+#             output[batch_size-1,0,h_index,w_index]=LUT[x_i[i],y_i[i],z_i[i],0]
+#             output[batch_size-1,1,h_index,w_index]=LUT[x_i[i],y_i[i],z_i[i],1]
+#             output[batch_size-1,2,h_index,w_index]=LUT[x_i[i],y_i[i],z_i[i],2]
+
+#         # x_i=x_i.view(batch_size,height*width)
+#         # y_i=y_i.view(batch_size,height*width)
+#         # z_i=z_i.view(batch_size,height*width)
+#         # 1,2160*3840
+
+
+#         return output
+
+
+#     def forward(self, LUT, img_input):
+#         assert img_input.ndimension()==4, 'img_input should be of shape [N,C,H,W]'
+#         # N=batch_size
+#         #img_input.size()=[1,3,2160,3840]\
+#         # LUT.size()=[3,33,33,33]
+#         assert LUT.ndimension()==4, 'LUT should be of shape [C,M,M,M](M=33)'
+#         batch_size, num_chans, height, width = img_input.shape
+#         dim = LUT.shape[1] # M
+#         img_size=img_input.size()
+#         Cmax=255.0
+#         s=Cmax/dim
+#         r,g,b=torch.split(img_input,split_size_or_sections=1,dim=1)
+#         # 将[1,3,2160,3840]以维度为1切成[1,1,2160,3840]的三部分        
+#         #r,g,b.size()=[1,1,2160,3840]
+#         # r=img_input[:,0,:,:]
+#         # g=img_input[:,1,:,:]
+#         # b=img_input[:,2,:,:]
+#         x=r/s
+#         y=g/s
+#         z=b/s
+#         # tmptmp=self.test(LUT,img_input)
+#         # x,y,z.size=[1,1,,2160,3840]
+#         # x_0,y_0,z_0.size=[1,1,,2160,3840]
+#         # x_1, y_1, z_1.size=[1,1,,2160,3840]
+#         x_0,y_0,z_0=x.floor(),y.floor(),z.floor()
+#         x_1, y_1, z_1 = x_0+1.0, y_0+1.0, z_0+1.0
+#         u, v, w = x-x_0, y-y_0, z-z_0
+#         # u,v,w.size=[1,1,2160,3840]
+#         # print("x_0.size",x_0.size())
+#         c_000 = self.test(LUT,torch.cat([x_0,y_0,z_0],dim=1))
+#         print(c_000.size())
+#         # x_i是顶点，大小为[1,1,2160,3840]
+#         # 输出c_xxx是对应顶点的LUT的值，大小为[1,3,2160,3840]
+#         c_100 = self.test(LUT,torch.cat([x_1,y_0,z_0],dim=1))
+#         c_010 = self.test(LUT,torch.cat([x_0,y_1,z_0],dim=1))
+#         c_110 = self.test(LUT,torch.cat([x_1,y_1,z_0],dim=1))
+#         c_001 = self.test(LUT,torch.cat([x_0,y_0,z_1],dim=1))
+#         c_101 = self.test(LUT,torch.cat([x_1,y_0,z_1],dim=1))
+#         c_011 = self.test(LUT,torch.cat([x_0,y_1,z_1],dim=1))
+#         c_111 = self.test(LUT,torch.cat([x_1,y_1,z_1],dim=1))
+
+#         # c_000 = self.gen_Cout_ijk(LUT,x_0,y_0,z_0)
+#         # # x_i是顶点，大小为[1,1,2160,3840]
+#         # # 输出c_xxx是对应顶点的LUT的值，大小为[1,3,2160,3840]
+#         # c_100 = self.gen_Cout_ijk(LUT,x_1,y_0,z_0)
+#         # c_010 = self.gen_Cout_ijk(LUT,x_0,y_1,z_0)
+#         # c_110 = self.gen_Cout_ijk(LUT,x_1,y_1,z_0)
+#         # c_001 = self.gen_Cout_ijk(LUT,x_0,y_0,z_1)
+#         # c_101 = self.gen_Cout_ijk(LUT,x_1,y_0,z_1)
+#         # c_011 = self.gen_Cout_ijk(LUT,x_0,y_1,z_1)
+#         # c_111 = self.gen_Cout_ijk(LUT,x_1,y_1,z_1)
+#         c_xyz = (1.0-u)*(1.0-v)*(1.0-w)*c_000 + \
+#         (1.0-u)*(1.0-v)*(w)*c_001 + \
+#         (1.0-u)*(v)*(1.0-w)*c_010 + \
+#         (1.0-u)*(v)*(w)*c_011 + \
+#         (u)*(1.0-v)*(1.0-w)*c_100 + \
+#         (u)*(1.0-v)*(w)*c_101 + \
+#         (u)*(v)*(1.0-w)*c_110 + \
+#         (u)*(v)*(w)*c_111
+#         # 广播机制，输出[1,3,2160,3840]
+#         print("c_xyz",c_xyz.size())
+#         return c_xyz
+
+#         # id100 = x_0 + 1.0 + y_0 * dim + z_0 * dim * dim
+#         # id010 = x_0 + (y_0 + 1.0) * dim + z_0 * dim * dim
+#         # id110 = x_0 + 1.0 + (y_0 + 1.0) * dim + z_0 * dim * dim
+#         # id001 = x_0 + y_0 * dim + (z_0 + 1.0) * dim * dim
+#         # id101 = x_0 + 1.0 + y_0 * dim + (z_0 + 1.0) * dim * dim
+#         # id011 = x_0 + (y_0 + 1.0) * dim + (z_0 + 1.0) * dim * dim
+#         # id111 = x_0 + 1.0 + (y_0 + 1.0) * dim + (z_0 + 1.0) * dim * dim
+
+#         # w000 = (1.0-u)*(1-v)*(1-w)
+#         # #大概也许得改成点乘
+#         # w100 = u*(1-v)*(1-w)
+#         # w010 = (1-u)*v*(1-w)
+#         # w110 = u*v*(1-w)
+#         # w001 = (1-u)*(1-v)*w
+#         # w101 = u*(1-v)*w
+#         # w011 = (1-u)*v*w
+#         # w111 = u*v*w
+#         # output=
+
+#         # print("v:",x_0,y_0,z_0)
+#         # print("s:",x_0.size(),y_0.size(),z_0.size())
+#         # u,v,w=u/s,v/s,w/s
+#         # c_000 = self.gen_Cout_ijk(x_0,y_0,z_0)
+#         # c_100 = self.gen_Cout_ijk(x_1,y_0,z_0)
+#         # c_010 = self.gen_Cout_ijk(x_0,y_1,z_0)
+#         # c_110 = self.gen_Cout_ijk(x_1,y_1,z_0)
+#         # c_001 = self.gen_Cout_ijk(x_0,y_0,z_1)
+#         # c_101 = self.gen_Cout_ijk(x_1,y_0,z_1)
+#         # c_011 = self.gen_Cout_ijk(x_0,y_1,z_1)
+#         # c_111 = self.gen_Cout_ijk(x_1,y_1,z_1)
+        
+
+#         # c_xyz = (1.0-u)*(1.0-v)*(1.0-w)*c_000 + \
+#         #         (1.0-u)*(1.0-v)*(w)*c_001 + \
+#         #         (1.0-u)*(v)*(1.0-w)*c_010 + \
+#         #         (1.0-u)*(v)*(w)*c_011 + \
+#         #         (u)*(1.0-v)*(1.0-w)*c_100 + \
+#         #         (u)*(1.0-v)*(w)*c_101 + \
+#         #         (u)*(v)*(1.0-w)*c_110 + \
+#         #         (u)*(v)*(w)*c_111
+#         # return c_xyz
+
+class Tritri(nn.Module):
+
+    def __init__(self):
+        super(Tritri, self).__init__()
+    
+    def forward(self,LUT,img):
+        img = (img - .5) * 2.
+        # grid_sample expects NxDxHxWx3 (1x1xHxWx3)
+        img = img.permute(0, 2, 3, 1)[:, None]
+        # add batch dim to LUT
+        LUT = LUT[None]
+        # grid sample
+        result = F.grid_sample(LUT, img, mode='bilinear', padding_mode='border', align_corners=True)
+        # drop added dimensions and permute back
+        result = result[:, :, 0].permute(0, 2, 3, 1)
+        return result
+
+
+
+class bing_lut_trilinearInterplt(nn.Module):
+
+    def __init__(self):
+        super(bing_lut_trilinearInterplt, self).__init__()
+    
+    def test(self,LUT,img_input):
+        # batch_size, num_chans, height, width = img_input.shape
+        # grid_height, grid_width = LUT.shape[1],LUT.shape[2]
+        grid_in=img_input.transpose(1,2).transpose(2,3)
+        # 1
+        # 原本img_input NCHW,改成 NHWC
+        xy_grid=grid_in[...,0:2]
+        yz_grid=grid_in[...,1:3]
+        # 23
+        #只取3通道中的第0和第1通道（0：2不含2）
+
+        # LUT正确版本应该是[3,33,33,33]
+        # 在这里弄错成为[33,33,33,3]
+        input_LUT=LUT[:,:,:,0:1]
+        input_LUT_ori=input_LUT.squeeze(3)
+        # 45
+        
+        # [3,33,33,33]->[3,33,33] 把dim=3的数据丢掉了
+
+        # input_LUT=LUT[:,:,0,:]
+        # input_LUT_ori=input_LUT.squeeze(2)
+        # # LUT[33,33,33,3]->[33,33,3],把dim=2的数据丢掉了
+
+        input_LUT=input_LUT_ori[0:2,...]
+        input_LUT2=input_LUT_ori[1:,...]
+        input_LUT=input_LUT.unsqueeze(0)
+        input_LUT2=input_LUT2.unsqueeze(0)
+        # 6-9
+
+        # 都是[1,2,33,33]
+        # print(input_LUT.size())
+        # print("dtype:")
+        # print(input_LUT.dtype)
+        # print(input_LUT2.dtype)
+        # print(xy_grid.dtype)
+        # print(yz_grid.dtype)
+        # input_LUT.int()
+        # input_LUT2.int()
+        # xy_grid.int()
+        # yz_grid.int()
+        
+        # # print(grid_in.size())
+        sampled_in_2d = F.grid_sample(input=input_LUT,grid=xy_grid, mode='nearest',align_corners=False)
+                                        # .view(batch_size, num_chans, num_d, grid_height, grid_width)
+        sampled_in_2d_2 = F.grid_sample(input=input_LUT2,grid=yz_grid, mode='nearest',align_corners=False)
+                                        # .view(batch_size, num_chans, num_d, grid_height, grid_width)
+        # 10
+        res=torch.cat([sampled_in_2d,sampled_in_2d_2[:,1:,:,:]],dim=1)
+        # print(res.size())
+        return res
+
+    def forward(self, LUT, img_input):
+        assert img_input.ndimension()==4, 'img_input should be of shape [N,C,H,W]'
+        # N=batch_size
+        #img_input.size()=[1,3,2160,3840]\
+        # LUT.size()=[3,33,33,33]
+        assert LUT.ndimension()==4, 'LUT should be of shape [C,M,M,M](M=33)'
+        # batch_size, num_chans, height, width = img_input.shape
+        dim = LUT.shape[1] # M
+        # img_size=img_input.size()
+        # Cmax=1.00001
+        Cmax=10
+        s=Cmax/(dim-1.0)
+        s=torch.Tensor([s])
+        #谢谢小黄鸭！！#data types int64 and int32 do not match in BroadcastRel
+
+        r,g,b=torch.split(img_input,split_size_or_sections=1,dim=1)
+        # 将[1,3,2160,3840]以维度为1切成[1,1,2160,3840]的三部分        
+        #r,g,b.size()=[1,1,2160,3840]
+        # r=img_input[:,0,:,:]
+        # g=img_input[:,1,:,:]
+        # b=img_input[:,2,:,:]
+        s=s.to(r.device)
+        x=r/s
+        y=g/s
+        z=b/s
+        # tmptmp=self.test(LUT,img_input)
+        # x,y,z.size=[1,1,,2160,3840]
+        # x_0,y_0,z_0.size=[1,1,,2160,3840]
+        # x_1, y_1, z_1.size=[1,1,,2160,3840]
+        x_0,y_0,z_0=x.floor(),y.floor(),z.floor()
+        x_1, y_1, z_1 = x_0+1.0, y_0+1.0, z_0+1.0
+        u, v, w = x-x_0, y-y_0, z-z_0
+        # u,v,w.size=[1,1,2160,3840]
+        # print("x_0.size",x_0.size())
+        
+        c_000 = self.test(LUT,torch.cat([x_0,y_0,z_0],dim=1))
+        # print(c_000.size())
+        # x_i是顶点，大小为[1,1,2160,3840]
+        # 输出c_xxx是对应顶点的LUT的值，大小为[1,3,2160,3840]
+        c_100 = self.test(LUT,torch.cat([x_1,y_0,z_0],dim=1))
+        c_010 = self.test(LUT,torch.cat([x_0,y_1,z_0],dim=1))
+        c_110 = self.test(LUT,torch.cat([x_1,y_1,z_0],dim=1))
+        c_001 = self.test(LUT,torch.cat([x_0,y_0,z_1],dim=1))
+        c_101 = self.test(LUT,torch.cat([x_1,y_0,z_1],dim=1))
+        c_011 = self.test(LUT,torch.cat([x_0,y_1,z_1],dim=1))
+        c_111 = self.test(LUT,torch.cat([x_1,y_1,z_1],dim=1))
+
+        c_xyz = (1.0-u)*(1.0-v)*(1.0-w)*c_000 + \
+        (1.0-u)*(1.0-v)*(w)*c_001 + \
+        (1.0-u)*(v)*(1.0-w)*c_010 + \
+        (1.0-u)*(v)*(w)*c_011 + \
+        (u)*(1.0-v)*(1.0-w)*c_100 + \
+        (u)*(1.0-v)*(w)*c_101 + \
+        (u)*(v)*(1.0-w)*c_110 + \
+        (u)*(v)*(w)*c_111
+        # 广播机制，输出[1,3,2160,3840]
+        print("c_xyz",c_xyz.size())
+        return c_xyz
+        
+class bing_lut_trilinearInterplt_backup(nn.Module):
+
+    def __init__(self):
+        super(bing_lut_trilinearInterplt, self).__init__()
+    
+    def test(self,LUT,img_input):
+        # batch_size, num_chans, height, width = img_input.shape
+        # grid_height, grid_width = LUT.shape[1],LUT.shape[2]
+        grid_in=img_input.transpose(1,2).transpose(2,3)
+        # 1
+        # 原本img_input NCHW,改成 NHWC
+        xy_grid=grid_in[...,0:2]
+        yz_grid=grid_in[...,1:3]
+        # 23
+        #只取3通道中的第0和第1通道（0：2不含2）
+
+        # LUT正确版本应该是[3,33,33,33]
+        # 在这里弄错成为[33,33,33,3]
+        input_LUT=LUT[:,:,:,0:1]
+        input_LUT_ori=input_LUT.squeeze(3)
+        # 45
+        
+        # [3,33,33,33]->[3,33,33] 把dim=3的数据丢掉了
+
+        # input_LUT=LUT[:,:,0,:]
+        # input_LUT_ori=input_LUT.squeeze(2)
+        # # LUT[33,33,33,3]->[33,33,3],把dim=2的数据丢掉了
+
+        input_LUT=input_LUT_ori[0:2,...]
+        input_LUT2=input_LUT_ori[1:,...]
+        input_LUT=input_LUT.unsqueeze(0)
+        input_LUT2=input_LUT2.unsqueeze(0)
+        # 6-9
+
+        # 都是[1,2,33,33]
+        # print(input_LUT.size())
+        # print("dtype:")
+        # print(input_LUT.dtype)
+        # print(input_LUT2.dtype)
+        # print(xy_grid.dtype)
+        # print(yz_grid.dtype)
+        # input_LUT.int()
+        # input_LUT2.int()
+        # xy_grid.int()
+        # yz_grid.int()
+        
+        # # print(grid_in.size())
+        sampled_in_2d = F.grid_sample(input=input_LUT,grid=xy_grid, mode='nearest')
+                                        # .view(batch_size, num_chans, num_d, grid_height, grid_width)
+        sampled_in_2d_2 = F.grid_sample(input=input_LUT2,grid=yz_grid, mode='nearest')
+                                        # .view(batch_size, num_chans, num_d, grid_height, grid_width)
+        # 10
+        res=torch.cat([sampled_in_2d,sampled_in_2d_2[:,1:,:,:]],dim=1)
+        # print(res.size())
+        return res
+
+    def forward(self, LUT, img_input):
+        assert img_input.ndimension()==4, 'img_input should be of shape [N,C,H,W]'
+        # N=batch_size
+        #img_input.size()=[1,3,2160,3840]\
+        # LUT.size()=[3,33,33,33]
+        assert LUT.ndimension()==4, 'LUT should be of shape [C,M,M,M](M=33)'
+        # batch_size, num_chans, height, width = img_input.shape
+        dim = LUT.shape[1] # M
+        # img_size=img_input.size()
+        Cmax=255.0
+        s=Cmax/dim
+        s=torch.Tensor([s])
+        #谢谢小黄鸭！！#data types int64 and int32 do not match in BroadcastRel
+
+        r,g,b=torch.split(img_input,split_size_or_sections=1,dim=1)
+        # 将[1,3,2160,3840]以维度为1切成[1,1,2160,3840]的三部分        
+        #r,g,b.size()=[1,1,2160,3840]
+        # r=img_input[:,0,:,:]
+        # g=img_input[:,1,:,:]
+        # b=img_input[:,2,:,:]
+        x=r/s
+        y=g/s
+        z=b/s
+        # tmptmp=self.test(LUT,img_input)
+        # x,y,z.size=[1,1,,2160,3840]
+        # x_0,y_0,z_0.size=[1,1,,2160,3840]
+        # x_1, y_1, z_1.size=[1,1,,2160,3840]
+        x_0,y_0,z_0=x.floor(),y.floor(),z.floor()
+        x_1, y_1, z_1 = x_0+1.0, y_0+1.0, z_0+1.0
+        u, v, w = x-x_0, y-y_0, z-z_0
+        # u,v,w.size=[1,1,2160,3840]
+        # print("x_0.size",x_0.size())
+        
+        c_000 = self.test(LUT,torch.cat([x_0,y_0,z_0],dim=1))
+        # print(c_000.size())
+        # x_i是顶点，大小为[1,1,2160,3840]
+        # 输出c_xxx是对应顶点的LUT的值，大小为[1,3,2160,3840]
+        c_100 = self.test(LUT,torch.cat([x_1,y_0,z_0],dim=1))
+        c_010 = self.test(LUT,torch.cat([x_0,y_1,z_0],dim=1))
+        c_110 = self.test(LUT,torch.cat([x_1,y_1,z_0],dim=1))
+        c_001 = self.test(LUT,torch.cat([x_0,y_0,z_1],dim=1))
+        c_101 = self.test(LUT,torch.cat([x_1,y_0,z_1],dim=1))
+        c_011 = self.test(LUT,torch.cat([x_0,y_1,z_1],dim=1))
+        c_111 = self.test(LUT,torch.cat([x_1,y_1,z_1],dim=1))
+
+        # c_000 = self.gen_Cout_ijk(LUT,x_0,y_0,z_0)
+        # # x_i是顶点，大小为[1,1,2160,3840]
+        # # 输出c_xxx是对应顶点的LUT的值，大小为[1,3,2160,3840]
+        # c_100 = self.gen_Cout_ijk(LUT,x_1,y_0,z_0)
+        # c_010 = self.gen_Cout_ijk(LUT,x_0,y_1,z_0)
+        # c_110 = self.gen_Cout_ijk(LUT,x_1,y_1,z_0)
+        # c_001 = self.gen_Cout_ijk(LUT,x_0,y_0,z_1)
+        # c_101 = self.gen_Cout_ijk(LUT,x_1,y_0,z_1)
+        # c_011 = self.gen_Cout_ijk(LUT,x_0,y_1,z_1)
+        # c_111 = self.gen_Cout_ijk(LUT,x_1,y_1,z_1)
+        c_xyz = (1.0-u)*(1.0-v)*(1.0-w)*c_000 + \
+        (1.0-u)*(1.0-v)*(w)*c_001 + \
+        (1.0-u)*(v)*(1.0-w)*c_010 + \
+        (1.0-u)*(v)*(w)*c_011 + \
+        (u)*(1.0-v)*(1.0-w)*c_100 + \
+        (u)*(1.0-v)*(w)*c_101 + \
+        (u)*(v)*(1.0-w)*c_110 + \
+        (u)*(v)*(w)*c_111
+        # 广播机制，输出[1,3,2160,3840]
+        print("c_xyz",c_xyz.size())
+        return c_xyz
+
+
+    
+    # @staticmethod
+    # def backward(ctx, lut_grad, x_grad):
+        
+    #     lut, x, int_package, float_package = ctx.saved_variables
+    #     dim, shift, W, H, batch = int_package
+    #     dim, shift, W, H, batch = int(dim), int(shift), int(W), int(H), int(batch)
+    #     binsize = float(float_package[0])
+            
+    #     assert 1 == trilinear.backward(x, 
+    #                                    x_grad, 
+    #                                    lut_grad,
+    #                                    dim, 
+    #                                    shift, 
+    #                                    binsize, 
+    #                                    W, 
+    #                                    H, 
+    #                                    batch)
+    #     return lut_grad, x_grad
+
+class Tri(nn.Module):
+    def __init__(self):
+        super(Tri,self).__init__()
+
+if __name__=='__main__':
+    # input_features: shape [B, num_channels, depth, height, width]
+    # sampling_grid: shape  [B,depth, height, 3]
+    data = torch.rand(1, 32, 16, 128, 128)
+    # data = torch.rand(1, 3, 16, 128, 128)
+    sampling_grid = (torch.rand(1, 256, 256, 3) - 0.5)*2.0
+    data = data.float().cuda(0)
+    sampling_grid = sampling_grid.float().cuda(0)
+    trilinear_interpolation = TrilinearIntepolation().cuda(0)
+    # LUT.type() torch.cuda.FloatTensor
+    # LUT.size() torch.Size([3, 33, 33, 33])
+    # img: torch.Size([1, 3, 2160, 3840])
+    data2 = torch.rand(1, 3,2160,3840)
+    # LUT2 = torch.rand(33,33,33,3)
+    LUT2 = torch.rand(3,33,33,33)
+
+    trilinear_interpolation2 = bing_lut_trilinearInterplt()
+    t_start = time.time()    
+    interp_data2=trilinear_interpolation2(LUT2,data2)
+
+    # interpolated_data = trilinear_interpolation(data, sampling_grid)
+    # print(interpolated_data.shape)
+    torch.cuda.synchronize()
+    print('time per iteration ', time.time()-t_start)
+    # for i in range(100):
+    #     t_start = time.time()
+    #     interpolated_data = trilinear_interpolation(data, sampling_grid)
+    #     print(interpolated_data.shape)
+    #     torch.cuda.synchronize()
+    #     print('time per iteration ', time.time()-t_start)

requirements.txt

@@ -0,0 +1,6 @@
+torch~=1.11.0
+torchvision~=0.12.0
+opencv-python~=4.5.5.64
+pillow~=9.1.1
+numpy~=1.22.3
+scipy~=1.8.1

torchvision_x_functional.py

@@ -0,0 +1,554 @@
+import collections
+import numbers
+from functools import wraps
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+from scipy.ndimage.filters import gaussian_filter
+
+__numpy_type_map = {
+    'float64': torch.DoubleTensor,
+    'float32': torch.FloatTensor,
+    'float16': torch.HalfTensor,
+    'int64': torch.LongTensor,
+    'int32': torch.IntTensor,
+    'int16': torch.ShortTensor,
+    'uint16': torch.ShortTensor,
+    'int8': torch.CharTensor,
+    'uint8': torch.ByteTensor,
+}
+
+'''image functional utils
+
+'''
+
+# NOTE: all the function should recive the ndarray like image, should be W x H x C or W x H
+
+# 如果将所有输出的维度够搞成height，width，channel 那么可以不用to_tensor??, 不行
+def preserve_channel_dim(func):
+    """Preserve dummy channel dim."""
+    @wraps(func)
+    def wrapped_function(img, *args, **kwargs):
+        shape = img.shape
+        result = func(img, *args, **kwargs)
+        if len(shape) == 3 and shape[-1] == 1 and len(result.shape) == 2:
+            result = np.expand_dims(result, axis=-1)
+        return result
+
+    return wrapped_function
+
+
+def _is_tensor_image(img):
+    return torch.is_tensor(img) and img.ndimension() == 3
+
+
+def _is_numpy_image(img):
+    return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
+
+
+def to_tensor(img):
+    '''convert numpy.ndarray to torch tensor. \n
+        if the image is uint8 , it will be divided by 255;\n
+        if the image is uint16 , it will be divided by 65535;\n
+        if the image is float , it will not be divided, we suppose your image range should between [0~1] ;\n
+    
+    Arguments:
+        img {numpy.ndarray} -- image to be converted to tensor.
+    '''
+    if not _is_numpy_image(img):
+        raise TypeError('data should be numpy ndarray. but got {}'.format(type(img)))
+
+    if img.ndim == 2:
+        img = img[:, :, None]
+
+    if img.dtype == np.uint8:
+        img = img.astype(np.float32)/255
+    elif img.dtype == np.uint16:
+        img = img.astype(np.float32)/65535
+    elif img.dtype in [np.float32, np.float64]:
+        img = img.astype(np.float32)/1
+    else:
+        raise TypeError('{} is not support'.format(img.dtype))
+    
+    img = torch.from_numpy(img.transpose((2, 0, 1)))
+
+    return img
+
+
+def to_pil_image(tensor):
+    # TODO
+    pass
+
+
+def to_tiff_image(tensor):
+    # TODO
+    pass
+
+
+def normalize(tensor, mean, std, inplace=False):
+    """Normalize a tensor image with mean and standard deviation.
+
+    .. note::
+        This transform acts out of place by default, i.e., it does not mutates the input tensor.
+
+    See :class:`~torchsat.transforms.Normalize` for more details.
+
+    Args:
+        tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channel.
+
+    Returns:
+        Tensor: Normalized Tensor image.
+    """
+    if not _is_tensor_image(tensor):
+        raise TypeError('tensor is not a torch image.')
+
+    if not inplace:
+        tensor = tensor.clone()
+
+    mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device)
+    std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device)
+    tensor.sub_(mean[:, None, None]).div_(std[:, None, None])
+    return tensor
+
+def noise(img, mode='gaussain', percent=0.02):
+    """
+    TODO: Not good for uint16 data
+    """
+    original_dtype = img.dtype
+    if mode == 'gaussian':
+        mean = 0
+        var = 0.1
+        sigma = var*0.5
+        
+        if img.ndim == 2:
+            h, w = img.shape
+            gauss = np.random.normal(mean, sigma, (h, w))
+        else:
+            h, w, c = img.shape
+            gauss = np.random.normal(mean, sigma, (h, w, c))
+            
+        if img.dtype not in [np.float32, np.float64]:
+            gauss = gauss * np.iinfo(img.dtype).max
+            img = np.clip(img.astype(np.float) + gauss, 0, np.iinfo(img.dtype).max)
+        else:
+            img = np.clip(img.astype(np.float) + gauss, 0, 1)
+
+    elif mode == 'salt':
+        print(img.dtype)
+        s_vs_p = 1
+        num_salt = np.ceil(percent * img.size * s_vs_p)
+        coords = tuple([np.random.randint(0, i - 1, int(num_salt)) for i in img.shape])
+        
+        if img.dtype in [np.float32, np.float64]:
+            img[coords] = 1
+        else:
+            img[coords] = np.iinfo(img.dtype).max
+            print(img.dtype)
+    elif mode == 'pepper':
+        s_vs_p = 0
+        num_pepper = np.ceil(percent * img.size * (1. - s_vs_p))
+        coords = tuple([np.random.randint(0, i - 1, int(num_pepper)) for i in img.shape])
+        img[coords] = 0
+
+    elif mode == 's&p':
+        s_vs_p = 0.5
+
+        # Salt mode
+        num_salt = np.ceil(percent * img.size * s_vs_p)
+        coords = tuple([np.random.randint(0, i - 1, int(num_salt)) for i in img.shape])
+        if img.dtype in [np.float32, np.float64]:
+            img[coords] = 1
+        else:
+            img[coords] = np.iinfo(img.dtype).max
+
+        # Pepper mode
+        num_pepper = np.ceil(percent* img.size * (1. - s_vs_p))
+        coords = tuple([np.random.randint(0, i - 1, int(num_pepper)) for i in img.shape])
+        img[coords] = 0
+    else:
+        raise ValueError('not support mode for {}'.format(mode))
+        
+    noisy = img.astype(original_dtype)
+    
+    return noisy
+
+
+def gaussian_blur(img, kernel_size):
+    # When sigma=0, it is computed as `sigma = 0.3*((ksize-1)*0.5 - 1) + 0.8`
+    return cv2.GaussianBlur(img, (kernel_size, kernel_size), sigmaX=0)
+
+
+def adjust_brightness(img, value=0):
+    if img.dtype in [np.float, np.float32, np.float64, np.float128]:
+        dtype_min, dtype_max = 0, 1
+        dtype = np.float32
+    else:
+        dtype_min = np.iinfo(img.dtype).min
+        dtype_max = np.iinfo(img.dtype).max
+        dtype = np.iinfo(img.dtype)
+    
+    result = np.clip(img.astype(np.float)+value, dtype_min, dtype_max).astype(dtype)
+    
+    return result
+    
+
+def adjust_contrast(img, factor):
+    if img.dtype in [np.float, np.float32, np.float64, np.float128]:
+        dtype_min, dtype_max = 0, 1
+        dtype = np.float32
+    else:
+        dtype_min = np.iinfo(img.dtype).min
+        dtype_max = np.iinfo(img.dtype).max
+        dtype = np.iinfo(img.dtype)
+    
+    result = np.clip(img.astype(np.float)*factor, dtype_min, dtype_max).astype(dtype)
+    
+    return result
+
+def adjust_saturation():
+    # TODO
+    pass
+
+def adjust_hue():
+    # TODO
+    pass
+
+
+
+def to_grayscale(img, output_channels=1):
+    """convert input ndarray image to gray sacle image.
+    
+    Arguments:
+        img {ndarray} -- the input ndarray image
+    
+    Keyword Arguments:
+        output_channels {int} -- output gray image channel (default: {1})
+    
+    Returns:
+        ndarray -- gray scale ndarray image
+    """
+    if img.ndim == 2:
+        gray_img = img
+    elif img.shape[2] == 3:
+        gray_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    else:
+        gray_img = np.mean(img, axis=2)
+        gray_img = gray_img.astype(img.dtype)
+
+    if output_channels != 1:
+        gray_img = np.tile(gray_img, (output_channels, 1, 1))
+        gray_img = np.transpose(gray_img, [1,2,0])
+        
+    return gray_img
+
+
+def shift(img, top, left):
+    (h, w) = img.shape[0:2]
+    matrix = np.float32([[1, 0, left], [0, 1, top]])
+    dst = cv2.warpAffine(img, matrix, (w, h))
+
+    return dst
+    
+
+def rotate(img, angle, center=None, scale=1.0):
+    (h, w) = img.shape[:2]
+ 
+    if center is None:
+        center = (w / 2, h / 2)
+ 
+    M = cv2.getRotationMatrix2D(center, angle, scale)
+    rotated = cv2.warpAffine(img, M, (w, h))
+ 
+    return rotated
+
+
+def resize(img, size, interpolation=Image.BILINEAR):
+    '''resize the image
+    TODO: opencv resize 之后图像就成了0~1了
+    Arguments:
+        img {ndarray} -- the input ndarray image
+        size {int, iterable} -- the target size, if size is intger,  width and height will be resized to same \
+                                otherwise, the size should be tuple (height, width) or list [height, width]
+                                
+    
+    Keyword Arguments:
+        interpolation {Image} -- the interpolation method (default: {Image.BILINEAR})
+    
+    Raises:
+        TypeError -- img should be ndarray
+        ValueError -- size should be intger or iterable vaiable and length should be 2.
+    
+    Returns:
+        img -- resize ndarray image
+    '''
+
+    if not _is_numpy_image(img):
+        raise TypeError('img shoud be ndarray image [w, h, c] or [w, h], but got {}'.format(type(img)))
+    if not (isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size)==2)):
+        raise ValueError('size should be intger or iterable vaiable(length is 2), but got {}'.format(type(size)))
+
+    if isinstance(size, int):
+        height, width = (size, size)
+    else:
+        height, width = (size[0], size[1])
+
+    return cv2.resize(img, (width, height), interpolation=interpolation)
+
+
+def pad(img, padding, fill=0, padding_mode='constant'):
+    if isinstance(padding, int):
+        pad_left = pad_right = pad_top = pad_bottom = padding
+    if isinstance(padding, collections.Iterable) and len(padding) == 2:
+        pad_left = pad_right = padding[0]
+        pad_bottom = pad_top = padding[1]
+    if isinstance(padding, collections.Iterable) and len(padding) == 4:
+        pad_left = padding[0]
+        pad_top = padding[1]
+        pad_right = padding[2]
+        pad_bottom = padding[3]
+
+    if img.ndim == 2:
+        if padding_mode == 'constant':
+            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), mode=padding_mode, constant_values=fill)
+        else:
+            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), mode=padding_mode)
+    if img.ndim == 3:
+        if padding_mode == 'constant':
+            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode=padding_mode, constant_values=fill)
+        else:
+            img = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode=padding_mode)
+    return img
+
+
+def crop(img, top, left, height, width):
+    '''crop image 
+    
+    Arguments:
+        img {ndarray} -- image to be croped
+        top {int} -- top size
+        left {int} -- left size 
+        height {int} -- croped height
+        width {int} -- croped width
+    '''
+    if not _is_numpy_image(img):
+        raise TypeError('the input image should be numpy ndarray with dimension 2 or 3.'
+            'but got {}'.format(type(img))
+        )
+    
+    if width<0 or height<0 or left <0 or height<0:
+        raise ValueError('the input left, top, width, height should be greater than 0'
+            'but got left={}, top={} width={} height={}'.format(left, top, width, height)
+        )
+    if img.ndim == 2:
+        img_height, img_width = img.shape
+    else:
+        img_height, img_width, _ = img.shape
+    if (left+width) > img_width or (top+height) > img_height:
+        raise ValueError('the input crop width and height should be small or \
+         equal to image width and height. ')
+
+    if img.ndim == 2:
+        return img[top:(top+height), left:(left+width)]
+    elif img.ndim == 3:
+        return img[top:(top+height), left:(left+width), :]
+
+
+def center_crop(img, output_size):
+    '''crop image
+    
+    Arguments:
+        img {ndarray} -- input image
+        output_size {number or sequence} -- the output image size. if sequence, should be [h, w]
+    
+    Raises:
+        ValueError -- the input image is large than original image.
+    
+    Returns:
+        ndarray image -- return croped ndarray image.
+    '''
+    if img.ndim == 2:
+        img_height, img_width = img.shape
+    else:
+        img_height, img_width, _ = img.shape
+
+    if isinstance(output_size, numbers.Number):
+        output_size = (int(output_size), int(output_size))
+    if output_size[0] > img_height or output_size[1] > img_width:
+        raise ValueError('the output_size should not greater than image size, but got {}'.format(output_size))
+    
+    target_height, target_width = output_size
+
+    top = int(round((img_height - target_height)/2))
+    left = int(round((img_width - target_width)/2))
+
+    return crop(img, top, left, target_height, target_width)
+    
+
+def resized_crop(img, top, left, height, width, size, interpolation=Image.BILINEAR):
+
+    img = crop(img, top, left, height, width)
+    img = resize(img, size, interpolation)
+    return img
+
+def vflip(img):
+    return cv2.flip(img, 0)
+
+def hflip(img):
+    return cv2.flip(img, 1)
+
+def flip(img, flip_code):
+    return cv2.flip(img, flip_code)
+
+
+def elastic_transform(image, alpha, sigma, alpha_affine, interpolation=cv2.INTER_LINEAR,
+                      border_mode=cv2.BORDER_REFLECT_101, random_state=None, approximate=False):
+    """Elastic deformation of images as described in [Simard2003]_ (with modifications).
+    Based on https://gist.github.com/erniejunior/601cdf56d2b424757de5
+    .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
+         Convolutional Neural Networks applied to Visual Document Analysis", in
+         Proc. of the International Conference on Document Analysis and
+         Recognition, 2003.
+    """
+    if random_state is None:
+        random_state = np.random.RandomState(1234)
+
+    height, width = image.shape[:2]
+
+    # Random affine
+    center_square = np.float32((height, width)) // 2
+    square_size = min((height, width)) // 3
+    alpha = float(alpha)
+    sigma = float(sigma)
+    alpha_affine = float(alpha_affine)
+
+    pts1 = np.float32([center_square + square_size, [center_square[0] + square_size, center_square[1] - square_size],
+                       center_square - square_size])
+    pts2 = pts1 + random_state.uniform(-alpha_affine, alpha_affine, size=pts1.shape).astype(np.float32)
+    matrix = cv2.getAffineTransform(pts1, pts2)
+
+    image = cv2.warpAffine(image, matrix, (width, height), flags=interpolation, borderMode=border_mode)
+
+    if approximate:
+        # Approximate computation smooth displacement map with a large enough kernel.
+        # On large images (512+) this is approximately 2X times faster
+        dx = (random_state.rand(height, width).astype(np.float32) * 2 - 1)
+        cv2.GaussianBlur(dx, (17, 17), sigma, dst=dx)
+        dx *= alpha
+
+        dy = (random_state.rand(height, width).astype(np.float32) * 2 - 1)
+        cv2.GaussianBlur(dy, (17, 17), sigma, dst=dy)
+        dy *= alpha
+    else:
+        dx = np.float32(gaussian_filter((random_state.rand(height, width) * 2 - 1), sigma) * alpha)
+        dy = np.float32(gaussian_filter((random_state.rand(height, width) * 2 - 1), sigma) * alpha)
+
+    x, y = np.meshgrid(np.arange(width), np.arange(height))
+
+    mapx = np.float32(x + dx)
+    mapy = np.float32(y + dy)
+
+    return cv2.remap(image, mapx, mapy, interpolation, borderMode=border_mode)
+
+
+def bbox_shift(bboxes, top, left):
+    pass
+
+
+def bbox_vflip(bboxes, img_height):
+    """vertical flip the bboxes
+    ...........
+    .         .
+    .         .
+   >...........<
+    .         .
+    .         .
+    ...........
+    Args:
+        bbox (ndarray): bbox ndarray [box_nums, 4]
+        flip_code (int, optional): [description]. Defaults to 0.
+    """
+    flipped = bboxes.copy()
+    flipped[...,1::2] = img_height - bboxes[...,1::2]
+    flipped = flipped[..., [0, 3, 2, 1]]
+    return flipped
+
+
+def bbox_hflip(bboxes, img_width):
+    """horizontal flip the bboxes
+          ^
+    .............
+    .     .     .
+    .     .     .
+    .     .     .
+    .     .     .
+    .............
+          ^
+    Args:
+        bbox (ndarray): bbox ndarray [box_nums, 4]
+        flip_code (int, optional): [description]. Defaults to 0.
+    """
+    flipped = bboxes.copy()
+    flipped[..., 0::2] = img_width - bboxes[...,0::2]
+    flipped = flipped[..., [2, 1, 0, 3]]
+    return flipped
+
+
+def bbox_resize(bboxes, img_size, target_size):
+    """resize the bbox
+    
+    Args:
+        bboxes (ndarray): bbox ndarray [box_nums, 4]
+        img_size (tuple): the image height and width
+        target_size (int, or tuple): the target bbox size. 
+                Int or Tuple, if tuple the shape should be (height, width)
+    """
+    if isinstance(target_size, numbers.Number):
+        target_size = (target_size, target_size)
+
+    ratio_height = target_size[0]/img_size[0]
+    ratio_width = target_size[1]/img_size[1]
+
+    return bboxes[...,]*[ratio_width,ratio_height,ratio_width,ratio_height]
+
+
+def bbox_crop(bboxes, top, left, height, width):
+    '''crop bbox 
+    
+    Arguments:
+        img {ndarray} -- image to be croped
+        top {int} -- top size
+        left {int} -- left size 
+        height {int} -- croped height
+        width {int} -- croped width
+    '''
+    croped_bboxes = bboxes.copy()
+
+    right = width + left
+    bottom = height + top
+    
+    croped_bboxes[..., 0::2] = bboxes[..., 0::2].clip(left, right) - left
+    croped_bboxes[..., 1::2] = bboxes[..., 1::2].clip(top, bottom) - top
+
+    return croped_bboxes
+
+def bbox_pad(bboxes, padding):
+    if isinstance(padding, int):
+        pad_left = pad_right = pad_top = pad_bottom = padding
+    if isinstance(padding, collections.Iterable) and len(padding) == 2:
+        pad_left = pad_right = padding[0]
+        pad_bottom = pad_top = padding[1]
+    if isinstance(padding, collections.Iterable) and len(padding) == 4:
+        pad_left = padding[0]
+        pad_top = padding[1]
+        pad_right = padding[2]
+        pad_bottom = padding[3]
+
+    pad_bboxes = bboxes.copy()
+    pad_bboxes[..., 0::2] = bboxes[..., 0::2] + pad_left
+    pad_bboxes[..., 1::2] = bboxes[..., 1::2] + pad_top
+
+    return pad_bboxes