models/modules/image_sampler.py

import sys
sys.path.append('../..')
import torch
import torch.nn as nn
import math
from models.modules.unet import RollOut_Conv
from einops import rearrange, reduce
MB =1024.0*1024.0
def mask_kernel(x, sigma=1):
    return torch.abs(x) < sigma #if the distance is smaller than the kernel size, return True

def mask_kernel_close_false(x, sigma=1):
    return torch.abs(x) > sigma #if the distance is smaller than the kernel size, return False

class Image_Local_Sampler(nn.Module):
    def __init__(self,reso,padding=0.1,in_channels=1280,out_channels=512):
        super().__init__()
        self.triplane_reso=reso
        self.padding=padding
        self.get_triplane_coord()
        self.img_proj=nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=1)
    def get_triplane_coord(self):
        '''xz plane firstly, z is at the '''
        x=torch.arange(self.triplane_reso)
        z=torch.arange(self.triplane_reso)
        X,Z=torch.meshgrid(x,z,indexing='xy')
        xz_coords=torch.cat([X[:,:,None],torch.ones_like(X[:,:,None])*(self.triplane_reso-1)/2,Z[:,:,None]],dim=-1) #in xyz order

        '''xy plane'''
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        X, Y = torch.meshgrid(x, y, indexing='xy')
        xy_coords = torch.cat([X[:, :, None],  Y[:, :, None],torch.ones_like(X[:, :, None])*(self.triplane_reso-1)/2], dim=-1)  # in xyz order

        '''yz plane'''
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)
        Y,Z = torch.meshgrid(y,z,indexing='xy')
        yz_coords= torch.cat([torch.ones_like(Y[:, :, None])*(self.triplane_reso-1)/2,Y[:,:,None],Z[:,:,None]], dim=-1)

        triplane_coords=torch.cat([xz_coords,xy_coords,yz_coords],dim=0)
        triplane_coords=triplane_coords/(self.triplane_reso-1)
        triplane_coords=(triplane_coords-0.5)*2*(1 + self.padding + 10e-6)
        self.triplane_coords=triplane_coords.float().cuda()

    def forward(self,image_feat,proj_mat):
        image_feat=self.img_proj(image_feat)
        batch_size=image_feat.shape[0]
        triplane_coords=self.triplane_coords.unsqueeze(0).expand(batch_size,-1,-1,-1) #B,192,64,3
        #print(torch.amin(triplane_coords),torch.amax(triplane_coords))
        coord_homo=torch.cat([triplane_coords,torch.ones((batch_size,triplane_coords.shape[1],triplane_coords.shape[2],1)).float().cuda()],dim=-1)
        coord_inimg=torch.einsum('bhwc,bck->bhwk',coord_homo,proj_mat.transpose(1,2))
        x=coord_inimg[:,:,:,0]/coord_inimg[:,:,:,2]
        y=coord_inimg[:,:,:,1]/coord_inimg[:,:,:,2]
        x=(x/(224.0-1.0)-0.5)*2 #-1~1
        y=(y/(224.0-1.0)-0.5)*2 #-1~1
        dist=coord_inimg[:,:,:,2]

        xy=torch.cat([x[:,:,:,None],y[:,:,:,None]],dim=-1)
        #print(image_feat.shape,xy.shape)
        sample_feat=torch.nn.functional.grid_sample(image_feat,xy,align_corners=True,mode='bilinear')
        return sample_feat

def position_encoding(d_model, length):
    if d_model % 2 != 0:
        raise ValueError("Cannot use sin/cos positional encoding with "
                         "odd dim (got dim={:d})".format(d_model))
    pe = torch.zeros(length, d_model)
    position = torch.arange(0, length).unsqueeze(1) #length,1
    div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) *
                         -(math.log(10000.0) / d_model))) #d_model//2, this is the frequency
    pe[:, 0::2] = torch.sin(position.float() * div_term) #length*(d_model//2)
    pe[:, 1::2] = torch.cos(position.float() * div_term)

    return pe

class Image_Vox_Local_Sampler(nn.Module):
    def __init__(self,reso,padding=0.1,in_channels=1280,inner_channel=128,out_channels=64,n_heads=8):
        super().__init__()
        self.triplane_reso=reso
        self.padding=padding
        self.get_vox_coord()
        self.out_channels=out_channels
        self.img_proj=nn.Conv2d(in_channels=in_channels,out_channels=inner_channel,kernel_size=1)

        self.vox_process=nn.Sequential(
            nn.Conv3d(in_channels=inner_channel,out_channels=inner_channel,kernel_size=3,padding=1,),
        )
        self.k=nn.Linear(in_features=inner_channel,out_features=inner_channel)
        self.q=nn.Linear(in_features=inner_channel,out_features=inner_channel)
        self.v=nn.Linear(in_features=inner_channel,out_features=inner_channel)
        self.attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)

        self.proj_out=nn.Conv2d(in_channels=inner_channel,out_channels=out_channels,kernel_size=1)
        self.condition_pe = position_encoding(inner_channel, self.triplane_reso).unsqueeze(0)
    def get_vox_coord(self):
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)

        X,Y,Z=torch.meshgrid(x,y,z,indexing='ij')
        vox_coor=torch.cat([X[:,:,:,None],Y[:,:,:,None],Z[:,:,:,None]],dim=-1)
        vox_coor=vox_coor/(self.triplane_reso-1)
        vox_coor=(vox_coor-0.5)*2*(1+self.padding+10e-6)
        self.vox_coor=vox_coor.view(-1,3).float().cuda()


    def forward(self,triplane_feat,image_feat,proj_mat):
        xz_feat,xy_feat,yz_feat=torch.split(triplane_feat,triplane_feat.shape[2]//3,dim=2) #B,C,64,64
        image_feat=self.img_proj(image_feat)
        batch_size=image_feat.shape[0]
        vox_coords=self.vox_coor.unsqueeze(0).expand(batch_size,-1,-1) #B,64*64*64,3
        vox_homo=torch.cat([vox_coords,torch.ones((batch_size,self.triplane_reso**3,1)).float().cuda()],dim=-1)
        coord_inimg=torch.einsum('bhc,bck->bhk',vox_homo,proj_mat.transpose(1,2))
        x=coord_inimg[:,:,0]/coord_inimg[:,:,2]
        y=coord_inimg[:,:,1]/coord_inimg[:,:,2]
        x=(x/(224.0-1.0)-0.5)*2 #-1~1
        y=(y/(224.0-1.0)-0.5)*2 #-1~1

        xy=torch.cat([x[:,:,None],y[:,:,None]],dim=-1).unsqueeze(1).contiguous() #B, 1,64**3,2
        #print(image_feat.shape,xy.shape)
        grid_feat=torch.nn.functional.grid_sample(image_feat,xy,align_corners=True,mode='bilinear').squeeze(2).\
            view(batch_size,-1,self.triplane_reso,self.triplane_reso,self.triplane_reso) #B,C,1,64**3

        grid_feat=self.vox_process(grid_feat)
        xzy_grid=grid_feat.permute(0,4,2,3,1)
        xz_as_query=xz_feat.permute(0,2,3,1).reshape(batch_size*self.triplane_reso**2,1,-1)
        xz_as_key=xzy_grid.reshape(batch_size*self.triplane_reso**2,self.triplane_reso,-1)

        xyz_grid=grid_feat.permute(0,3,2,4,1)
        xy_as_query=xy_feat.permute(0,2,3,1).reshape(batch_size*self.triplane_reso**2,1,-1)
        xy_as_key = xyz_grid.reshape(batch_size * self.triplane_reso ** 2, self.triplane_reso, -1)

        yzx_grid = grid_feat.permute(0, 4, 3, 2, 1)
        yz_as_query = yz_feat.permute(0,2,3,1).reshape(batch_size*self.triplane_reso**2,1,-1)
        yz_as_key = yzx_grid.reshape(batch_size * self.triplane_reso ** 2, self.triplane_reso, -1)

        query=self.q(torch.cat([xz_as_query,xy_as_query,yz_as_query],dim=0))
        key=self.k(torch.cat([xz_as_key,xy_as_key,yz_as_key],dim=0))+self.condition_pe.to(xz_as_key.device)
        value=self.v(torch.cat([xz_as_key,xy_as_key,yz_as_key],dim=0))+self.condition_pe.to(xz_as_key.device)

        attn,_=self.attn(query,key,value)
        xz_plane,xy_plane,yz_plane=torch.split(attn,dim=0,split_size_or_sections=batch_size*self.triplane_reso**2)
        xz_plane=xz_plane.reshape(batch_size,self.triplane_reso,self.triplane_reso,-1).permute(0,3,1,2)
        xy_plane = xy_plane.reshape(batch_size, self.triplane_reso, self.triplane_reso, -1).permute(0, 3, 1, 2)
        yz_plane = yz_plane.reshape(batch_size, self.triplane_reso, self.triplane_reso, -1).permute(0, 3, 1, 2)

        triplane_wImg=torch.cat([xz_plane,xy_plane,yz_plane],dim=2)
        triplane_wImg=self.proj_out(triplane_wImg)
        #print(triplane_wImg.shape)

        return triplane_wImg

class Image_Direct_AttenwMask_Sampler(nn.Module):
    def __init__(self,reso,vit_reso=16,padding=0.1,triplane_in_channels=64,
                 img_in_channels=1280,inner_channel=128,out_channels=64,n_heads=8):
        super().__init__()
        self.triplane_reso=reso
        self.vit_reso=vit_reso
        self.padding=padding
        self.n_heads=n_heads
        self.get_plane_expand_coord()
        self.get_vit_coords()
        self.out_channels=out_channels
        self.kernel_func=mask_kernel
        self.k=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.q=nn.Linear(in_features=triplane_in_channels,out_features=inner_channel)
        self.v=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)

        self.proj_out=nn.Linear(in_features=inner_channel,out_features=out_channels)
        self.image_pe = position_encoding(inner_channel, self.vit_reso**2+1).unsqueeze(0).cuda().float() #1,n_img*reso*reso,channel
        self.triplane_pe = position_encoding(inner_channel, 3*self.triplane_reso**2).unsqueeze(0).cuda().float()
    def get_plane_expand_coord(self):
        x = torch.arange(self.triplane_reso)/(self.triplane_reso-1)
        y = torch.arange(self.triplane_reso)/(self.triplane_reso-1)
        z = torch.arange(self.triplane_reso)/(self.triplane_reso-1)

        first,second,third=torch.meshgrid(x,y,z,indexing='xy')
        xyz_coords=torch.stack([first,second,third],dim=-1)#reso,reso,reso,3
        xyz_coords=(xyz_coords-0.5)*2*(1+self.padding+10e-6) #ordering yxz ->xyz
        xzy_coords=xyz_coords.clone().permute(2,1,0,3) #ordering zxy ->xzy
        yzx_coords=xyz_coords.clone().permute(2,0,1,3) #ordering zyx ->yzx

        # print(xyz_coords[0,0,0],xyz_coords[0,0,1],xyz_coords[1,0,0],xyz_coords[0,1,0])
        # print(xzy_coords[0, 0, 0], xzy_coords[0, 0, 1], xzy_coords[1, 0, 0], xzy_coords[0, 1, 0])
        # print(yzx_coords[0, 0, 0], yzx_coords[0, 0, 1], yzx_coords[1, 0, 0], yzx_coords[0, 1, 0])

        xyz_coords=xyz_coords.reshape(self.triplane_reso**3,-1)
        xzy_coords=xzy_coords.reshape(self.triplane_reso**3,-1)
        yzx_coords=yzx_coords.reshape(self.triplane_reso**3,-1)

        coords=torch.cat([xzy_coords,xyz_coords,yzx_coords],dim=0)
        self.plane_coords=coords.cuda().float()
        # self.xzy_coords=xzy_coords.cuda().float() #reso**3，3
        # self.xyz_coords=xyz_coords.cuda().float() #reso**3，3
        # self.yzx_coords=yzx_coords.cuda().float() #reso**3，3

    def get_vit_coords(self):
        x=torch.arange(self.vit_reso)
        y=torch.arange(self.vit_reso)

        X,Y=torch.meshgrid(x,y,indexing='xy')
        vit_coords=torch.stack([X,Y],dim=-1)
        self.vit_coords=vit_coords.view(self.vit_reso**2,2).cuda().float()

    def get_attn_mask(self,coords_proj,vit_coords,kernel_size=1.0):
        '''
        :param coords_proj: B,reso**3,2, in range of 0~1
        :param vit_coords: B,vit_reso**2,2, in range of 0~vit_reso
        :param kernel_size: 0.5, so that only one pixel will be select
        :return:
        '''
        bs=coords_proj.shape[0]
        coords_proj=coords_proj*(self.vit_reso-1)
        #print(torch.amin(coords_proj[0,0:self.triplane_reso**3]),torch.amax(coords_proj[0,0:self.triplane_reso**3]))
        dist=torch.cdist(coords_proj.float(),vit_coords.float())
        mask=self.kernel_func(dist,sigma=kernel_size).float() #True if valid, B,3*reso**3,vit_reso**2
        mask=mask.reshape(bs,3*self.triplane_reso**2,self.triplane_reso,self.vit_reso**2)
        mask=torch.sum(mask,dim=2)
        attn_mask=(mask==0)
        return attn_mask

    def forward(self,triplane_feat,image_feat,proj_mat):
        #xz_feat,xy_feat,yz_feat=torch.split(triplane_feat,triplane_feat.shape[2]//3,dim=2) #B,C,64,64
        batch_size=image_feat.shape[0]
        #print(self.plane_coords.shape)
        coords=self.plane_coords.unsqueeze(0).expand(batch_size,-1,-1)

        coords_homo=torch.cat([coords,torch.ones(batch_size,self.triplane_reso**3*3,1).float().cuda()],dim=-1)
        coords_inimg=torch.einsum('bhc,bck->bhk',coords_homo,proj_mat.transpose(1,2))
        coords_x=coords_inimg[:,:,0]/coords_inimg[:,:,2]/(224.0-1) #0~1
        coords_y=coords_inimg[:,:,1]/coords_inimg[:,:,2]/(224.0-1) #0~1
        coords_x=torch.clamp(coords_x,min=0.0,max=1.0)
        coords_y=torch.clamp(coords_y,min=0.0,max=1.0)
        #print(torch.amin(coords_x),torch.amax(coords_x))
        coords_proj=torch.stack([coords_x,coords_y],dim=-1)
        vit_coords=self.vit_coords.unsqueeze(0).expand(batch_size,-1,-1)
        attn_mask=torch.repeat_interleave(
            self.get_attn_mask(coords_proj,vit_coords,kernel_size=1.0),self.n_heads, 0
        )
        attn_mask = torch.cat([torch.zeros([attn_mask.shape[0], attn_mask.shape[1], 1]).cuda().bool(), attn_mask],
                              dim=-1)  # add global token
        #print(attn_mask.shape,torch.sum(attn_mask.float()))
        triplane_feat=triplane_feat.permute(0,2,3,1).view(batch_size,3*self.triplane_reso**2,-1)
        #print(triplane_feat.shape,self.triplane_pe.shape)
        query=self.q(triplane_feat)+self.triplane_pe
        key=self.k(image_feat)+self.image_pe
        value=self.v(image_feat)+self.image_pe
        #print(query.shape,key.shape,value.shape)
        attn,_=self.attn(query,key,value,attn_mask=attn_mask)
        #print(attn.shape)
        output=self.proj_out(attn).transpose(1,2).reshape(batch_size,-1,3*self.triplane_reso,self.triplane_reso)

        return output

class MultiImage_Direct_AttenwMask_Sampler(nn.Module):
    def __init__(self,reso,vit_reso=16,padding=0.1,triplane_in_channels=64,
                 img_in_channels=1280,inner_channel=128,out_channels=64,n_heads=8,max_nimg=5):
        super().__init__()
        self.triplane_reso=reso
        self.vit_reso=vit_reso
        self.padding=padding
        self.n_heads=n_heads
        self.get_plane_expand_coord()
        self.get_vit_coords()
        self.out_channels=out_channels
        self.kernel_func=mask_kernel
        self.k=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.q=nn.Linear(in_features=triplane_in_channels,out_features=inner_channel)
        self.v=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)

        self.proj_out=nn.Linear(in_features=inner_channel,out_features=out_channels)
        self.image_pe = position_encoding(inner_channel, max_nimg*(self.vit_reso**2+1)).unsqueeze(0).cuda().float()
        self.triplane_pe = position_encoding(inner_channel, 3*self.triplane_reso**2).unsqueeze(0).cuda().float()
    def get_plane_expand_coord(self):
        x = torch.arange(self.triplane_reso)/(self.triplane_reso-1)
        y = torch.arange(self.triplane_reso)/(self.triplane_reso-1)
        z = torch.arange(self.triplane_reso)/(self.triplane_reso-1)

        first,second,third=torch.meshgrid(x,y,z,indexing='xy')
        xyz_coords=torch.stack([first,second,third],dim=-1)#reso,reso,reso,3
        xyz_coords=(xyz_coords-0.5)*2*(1+self.padding+10e-6) #ordering yxz ->xyz
        xzy_coords=xyz_coords.clone().permute(2,1,0,3) #ordering zxy ->xzy
        yzx_coords=xyz_coords.clone().permute(2,0,1,3) #ordering zyx ->yzx

        xyz_coords=xyz_coords.reshape(self.triplane_reso**3,-1)
        xzy_coords=xzy_coords.reshape(self.triplane_reso**3,-1)
        yzx_coords=yzx_coords.reshape(self.triplane_reso**3,-1)

        coords=torch.cat([xzy_coords,xyz_coords,yzx_coords],dim=0)
        self.plane_coords=coords.cuda().float()
        # self.xzy_coords=xzy_coords.cuda().float() #reso**3，3
        # self.xyz_coords=xyz_coords.cuda().float() #reso**3，3
        # self.yzx_coords=yzx_coords.cuda().float() #reso**3，3

    def get_vit_coords(self):
        x=torch.arange(self.vit_reso)
        y=torch.arange(self.vit_reso)

        X,Y=torch.meshgrid(x,y,indexing='xy')
        vit_coords=torch.stack([X,Y],dim=-1)
        self.vit_coords=vit_coords.view(self.vit_reso**2,2).cuda().float()

    def get_attn_mask(self,coords_proj,vit_coords,valid_frames,kernel_size=1.0):
        '''
        :param coords_proj: B,n_img,3*reso**3,2, in range of 0~vit_reso
        :param vit_coords:  B,n_img,vit_reso**2,2, in range of 0~vit_reso
        :param kernel_size: 0.5, so that only one pixel will be select
        :return:
        '''
        bs,n_img=coords_proj.shape[0],coords_proj.shape[1]
        coords_proj_flat=coords_proj.reshape(bs*n_img,3*self.triplane_reso**3,2)
        vit_coords_flat=vit_coords.reshape(bs*n_img,self.vit_reso**2,2)
        dist=torch.cdist(coords_proj_flat.float(),vit_coords_flat.float())
        mask=self.kernel_func(dist,sigma=kernel_size).float() #True if valid, B*n_img,3*reso**3,vit_reso**2
        mask=mask.reshape(bs,n_img,3*self.triplane_reso**2,self.triplane_reso,self.vit_reso**2)
        mask=torch.sum(mask,dim=3) #B,n_img,3*reso**2,vit_reso**2
        mask=torch.cat([torch.ones(size=mask.shape[0:3]).unsqueeze(3).float().cuda(),mask],dim=-1) #B,n_img,3*reso**2,vit_reso**2+1, add global mask
        mask[valid_frames == 0, :, :] = False
        mask=mask.permute(0,2,1,3).reshape(bs,3*self.triplane_reso**2,-1) #B,3*reso**2,n_img*(vit_resso**2+1)
        attn_mask=(mask==0) #invert the mask, False indicates valid, True indicates invalid
        return attn_mask

    def forward(self,triplane_feat,image_feat,proj_mat,valid_frames):
        '''image feat is bs,n_img,length,channel'''
        batch_size,n_img=image_feat.shape[0],image_feat.shape[1]
        img_length=image_feat.shape[2]
        image_feat_flat=image_feat.view(batch_size,n_img*img_length,-1)
        coords=self.plane_coords.unsqueeze(0).unsqueeze(1).expand(batch_size,n_img,-1,-1)

        coord_homo=torch.cat([coords,torch.ones(batch_size,n_img,self.triplane_reso**3*3,1).float().cuda()],dim=-1)
        #print(coord_homo.shape,proj_mat.shape)
        coord_inimg = torch.einsum('bjnc,bjck->bjnk', coord_homo, proj_mat.transpose(2, 3))
        x = coord_inimg[:, :, :, 0] / coord_inimg[:, :, :, 2]
        y = coord_inimg[:, :, :, 1] / coord_inimg[:, :, :, 2]
        x = x/(224.0-1)
        y = y/(224.0-1)
        coords_x=torch.clamp(x,min=0.0,max=1.0)*(self.vit_reso-1)
        coords_y=torch.clamp(y,min=0.0,max=1.0)*(self.vit_reso-1)
        coords_proj=torch.stack([coords_x,coords_y],dim=-1)
        vit_coords=self.vit_coords.unsqueeze(0).unsqueeze(1).expand(batch_size,n_img,-1,-1)
        attn_mask=torch.repeat_interleave(
            self.get_attn_mask(coords_proj,vit_coords,valid_frames,kernel_size=1.0),self.n_heads, 0
        )
        triplane_feat=triplane_feat.permute(0,2,3,1).view(batch_size,3*self.triplane_reso**2,-1)
        query=self.q(triplane_feat)+self.triplane_pe
        key=self.k(image_feat_flat)+self.image_pe
        value=self.v(image_feat_flat)+self.image_pe
        attn,_=self.attn(query,key,value,attn_mask=attn_mask)
        output=self.proj_out(attn).transpose(1,2).reshape(batch_size,-1,3*self.triplane_reso,self.triplane_reso)

        return output

class MultiImage_Fuse_Sampler(nn.Module):
    def __init__(self,reso,vit_reso=16,padding=0.1,triplane_in_channels=64,
                 img_in_channels=1280,inner_channel=128,out_channels=64,n_heads=8):
        super().__init__()
        self.triplane_reso=reso
        self.vit_reso=vit_reso
        self.inner_channel=inner_channel
        self.padding=padding
        self.n_heads=n_heads
        self.get_vox_coord()
        self.get_vit_coords()
        self.out_channels=out_channels
        self.kernel_func=mask_kernel
        self.image_unflatten=nn.Unflatten(2,(vit_reso,vit_reso))
        self.k=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.q=nn.Linear(in_features=triplane_in_channels*3,out_features=inner_channel)
        self.v=nn.Linear(in_features=img_in_channels,out_features=inner_channel)

        #self.cross_attn=CrossAttention(query_dim=inner_channel,heads=8,dim_head=inner_channel//8)
        self.cross_attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)
        self.proj_out=nn.Linear(in_features=inner_channel,out_features=out_channels)
        self.image_pe = position_encoding(inner_channel, self.vit_reso**2)[None,None,:,:].cuda().float() #1,1,length,channel
        #self.image_pe = self.image_pe.reshape(1,max_nimg,self.vit_reso,self.vit_reso,inner_channel)
        self.triplane_pe = position_encoding(inner_channel, self.triplane_reso ** 3).unsqueeze(0).cuda().float()

    def get_vit_coords(self):
        x = torch.arange(self.vit_reso)
        y = torch.arange(self.vit_reso)

        X, Y = torch.meshgrid(x, y, indexing='xy')
        vit_coords = torch.stack([X, Y], dim=-1)
        self.vit_coords = vit_coords.cuda().float() #reso,reso,2

    def get_vox_coord(self):
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)

        X, Y, Z = torch.meshgrid(x, y, z, indexing='ij')
        vox_coor = torch.cat([X[:, :, :, None], Y[:, :, :, None], Z[:, :, :, None]], dim=-1)
        self.vox_index = vox_coor.view(-1, 3).long().cuda()

        vox_coor = self.vox_index.float() / (self.triplane_reso - 1)
        vox_coor = (vox_coor - 0.5) * 2 * (1 + self.padding + 10e-6)
        self.vox_coor = vox_coor.view(-1, 3).float().cuda()

    def get_attn_mask(self,valid_frames):
        '''
        :param valid_frames: of shape B,n_img
        '''
        #print(valid_frames)
        #bs,n_img=valid_frames.shape[0:2]
        attn_mask=(valid_frames.float()==0)
        #attn_mask=attn_mask.unsqueeze(1).unsqueeze(2).expand(-1,self.triplane_reso**3,-1,-1) #B,1,n_img
        #attn_mask=attn_mask.reshape(bs*self.triplane_reso**3,-1,n_img).bool()
        attn_mask=torch.repeat_interleave(attn_mask.unsqueeze(1),self.triplane_reso**3,0)
        # print(attn_mask[self.triplane_reso**3*1+10])
        # print(attn_mask[self.triplane_reso ** 3 * 2+10])
        # print(attn_mask[self.triplane_reso ** 3 * 3+10])
        return attn_mask

    def forward(self,triplane_feat,image_feat,proj_mat,valid_frames):
        '''image feat is bs,n_img,length,channel'''
        batch_size,n_img=image_feat.shape[0],image_feat.shape[1]
        image_feat=image_feat[:,:,1:,:] #discard global feature

        #image_feat=image_feat.permute(0,1,3,4,2) #B,n_img,h,w,c
        image_k=self.k(image_feat)+self.image_pe #B,n_img,h,w,c
        image_v=self.v(image_feat)+self.image_pe #B,n_img,h,w,c
        image_k_v=torch.cat([image_k,image_v],dim=-1) #B,n_img,h,w,c
        unflat_k_v=self.image_unflatten(image_k_v).permute(0,4,1,2,3) #Bs,channel,n_img,reso,reso
        #unflat_k_v=image_k_v.permute(0,4,1,2,3)
        #vit_coords=self.vit_coords[None,None].expand(batch_size,n_img,-1,-1,-1) #Bs,n_img,reso,reso,2

        coords=self.vox_coor.unsqueeze(0).unsqueeze(1).expand(batch_size,n_img,-1,-1)
        coord_homo=torch.cat([coords,torch.ones(batch_size,n_img,self.triplane_reso**3,1).float().cuda()],dim=-1)
        coord_inimg = torch.einsum('bjnc,bjck->bjnk', coord_homo, proj_mat.transpose(2, 3))
        x = coord_inimg[:, :, :, 0] / coord_inimg[:, :, :, 2]
        y = coord_inimg[:, :, :, 1] / coord_inimg[:, :, :, 2]
        x = x/(224.0-1) #0~1
        y = y/(224.0-1)
        coords_proj=torch.stack([x,y],dim=-1)
        coords_proj=(coords_proj-0.5)*2
        img_index=((torch.arange(n_img)[None,:,None,None].expand(
            batch_size,-1,self.triplane_reso**3,-1).float().cuda()/(n_img-1))-0.5)*2 #Bs,n_img,64**3,1

        # img_index_feat=torch.arange(n_img)[None,:,None,None,None].expand(
        #     batch_size,-1,self.vit_reso,self.vit_reso,-1).float().cuda() #Bs,n_img,reso,reso,1
        #coords_feat=torch.cat([vit_coords,img_index_feat],dim=-1).permute(0,4,1,2,3)#Bs,n_img,reso,reso,3
        grid=torch.cat([coords_proj,img_index],dim=-1) #x,y,index
        grid=torch.clamp(grid,min=-1.0,max=1.0)
        sample_k_v = torch.nn.functional.grid_sample(unflat_k_v, grid.unsqueeze(1), align_corners=True, mode='bilinear').squeeze(2) #B,C,n_img,64**3
        xz_feat, xy_feat, yz_feat = torch.split(triplane_feat, split_size_or_sections=triplane_feat.shape[2] // 3,
                                                dim=2)  # B,C,64,64
        xz_vox_feat=xz_feat.unsqueeze(4).expand(-1,-1,-1,-1,self.triplane_reso)#.permute(0,1,3,4,2).reshape(batch_size,-1,self.triplane_reso**3).transpose(1,2) #zxy
        xz_vox_feat=rearrange(xz_vox_feat, 'b c z x y -> b (x y z) c')
        xy_vox_feat=xy_feat.unsqueeze(4).expand(-1,-1,-1,-1,self.triplane_reso)#.permute(0,1,3,2,4).reshape(batch_size,-1,self.triplane_reso**3).transpose(1,2) #yxz
        xy_vox_feat=rearrange(xy_vox_feat, 'b c y x z -> b (x y z) c')
        yz_vox_feat=yz_feat.unsqueeze(4).expand(-1,-1,-1,-1,self.triplane_reso)#.permute(0,1,4,3,2).reshape(batch_size,-1,self.triplane_reso**3).transpose(1,2) #zyx
        yz_vox_feat=rearrange(yz_vox_feat, 'b c z y x -> b (x y z) c')
        #xz_vox_feat = xz_feat[:, :, vox_index[:, 2], vox_index[:, 0]].transpose(1, 2)  # B,C,64*64*64
        #xy_vox_feat = xy_feat[:, :, vox_index[:, 1], vox_index[:, 0]].transpose(1, 2)
        #yz_vox_feat = yz_feat[:, :, vox_index[:, 2], vox_index[:, 1]].transpose(1, 2)

        triplane_expand_feat = torch.cat([xz_vox_feat, xy_vox_feat, yz_vox_feat], dim=-1)  # B,64*64*64,3*C
        triplane_query = self.q(triplane_expand_feat) + self.triplane_pe
        k_v=rearrange(sample_k_v, 'b c n k -> (b k) n c')
        #k_v=sample_k_v.permute(0,3,2,1).reshape(batch_size*self.triplane_reso**3,n_img,-1) #B*64**3,n_img,C
        k=k_v[:,:,0:self.inner_channel]
        v=k_v[:,:,self.inner_channel:]
        q=rearrange(triplane_query,'b k c -> (b k) 1 c')
        #q=triplane_query.view(batch_size*self.triplane_reso**3,1,-1)
        #k,v is of shape, B*reso**3,k,channel, q is of shape B*reso**3,1,channel
        #attn mask should be B*reso**3*n_heads,1,k
        #attn_mask=torch.repeat_interleave(self.get_attn_mask(valid_frames),self.n_heads,0)
        #print(q.shape,k.shape,v.shape)
        attn_out,_=self.cross_attn(q,k,v)#attn_mask=attn_mask) #fuse multi-view feature
        #volume=attn_out.view(batch_size,self.triplane_reso,self.triplane_reso,self.triplane_reso,-1) #B,reso,reso,reso,channel
        #print(attn_out.shape)
        volume=rearrange(attn_out,'(b x y z) 1 c -> b x y z c',x=self.triplane_reso,y=self.triplane_reso,z=self.triplane_reso)
        #xz_feat = torch.mean(volume, dim=2).transpose(1,2) #B,reso,reso,C
        xz_feat = reduce(volume, "b x y z c -> b z x c", 'mean')
        #xy_feat = torch.mean(volume, dim=3).transpose(1,2) #B,reso,reso,C
        xy_feat= reduce(volume, 'b x y z c -> b y x c', 'mean')
        #yz_feat = torch.mean(volume, dim=1).transpose(1,2) #B,reso,reso,C
        yz_feat=reduce(volume, 'b x y z c -> b z y c', 'mean')
        triplane_out = torch.cat([xz_feat, xy_feat, yz_feat], dim=1) #B,reso*3,reso,C
        #print(triplane_out.shape)
        triplane_out = self.proj_out(triplane_out)
        triplane_out = triplane_out.permute(0,3,1,2)
        #print(triplane_out.shape)
        return triplane_out

class MultiImage_TriFuse_Sampler(nn.Module):
    def __init__(self,reso,vit_reso=16,padding=0.1,triplane_in_channels=64,
                 img_in_channels=1280,inner_channel=128,out_channels=64,n_heads=8,max_nimg=5):
        super().__init__()
        self.triplane_reso=reso
        self.vit_reso=vit_reso
        self.inner_channel=inner_channel
        self.padding=padding
        self.n_heads=n_heads
        self.get_triplane_coord()
        self.get_vit_coords()
        self.out_channels=out_channels
        self.kernel_func=mask_kernel
        self.image_unflatten=nn.Unflatten(2,(vit_reso,vit_reso))
        self.k=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.q=nn.Linear(in_features=triplane_in_channels,out_features=inner_channel)
        self.v=nn.Linear(in_features=img_in_channels,out_features=inner_channel)

        self.cross_attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)
        self.proj_out=nn.Conv2d(in_channels=inner_channel,out_channels=out_channels,kernel_size=1)
        self.image_pe = position_encoding(inner_channel, self.vit_reso**2)[None,None,:,:].expand(-1,max_nimg,-1,-1).cuda().float() #B,n_img,length,channel
        self.triplane_pe = position_encoding(inner_channel, self.triplane_reso ** 2*3).unsqueeze(0).cuda().float()

    def get_vit_coords(self):
        x = torch.arange(self.vit_reso)
        y = torch.arange(self.vit_reso)

        X, Y = torch.meshgrid(x, y, indexing='xy')
        vit_coords = torch.stack([X, Y], dim=-1)
        self.vit_coords = vit_coords.cuda().float() #reso,reso,2

    def get_triplane_coord(self):
        '''xz plane firstly, z is at the '''
        x = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)
        X, Z = torch.meshgrid(x, z, indexing='xy')
        xz_coords = torch.cat(
            [X[:, :, None], torch.ones_like(X[:, :, None]) * (self.triplane_reso - 1) / 2, Z[:, :, None]],
            dim=-1)  # in xyz order

        '''xy plane'''
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        X, Y = torch.meshgrid(x, y, indexing='xy')
        xy_coords = torch.cat(
            [X[:, :, None], Y[:, :, None], torch.ones_like(X[:, :, None]) * (self.triplane_reso - 1) / 2],
            dim=-1)  # in xyz order

        '''yz plane'''
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)
        Y, Z = torch.meshgrid(y, z, indexing='xy')
        yz_coords = torch.cat(
            [torch.ones_like(Y[:, :, None]) * (self.triplane_reso - 1) / 2, Y[:, :, None], Z[:, :, None]], dim=-1)

        triplane_coords = torch.cat([xz_coords, xy_coords, yz_coords], dim=0)
        triplane_coords = triplane_coords / (self.triplane_reso - 1)
        triplane_coords = (triplane_coords - 0.5) * 2 * (1 + self.padding + 10e-6)
        self.triplane_coords = triplane_coords.view(-1,3).float().cuda()
    def forward(self,triplane_feat,image_feat,proj_mat,valid_frames):
        '''image feat is bs,n_img,length,channel'''
        batch_size,n_img=image_feat.shape[0],image_feat.shape[1]
        image_feat=image_feat[:,:,1:,:] #discard global feature
        #print(image_feat.shape)

        #image_feat=image_feat.permute(0,1,3,4,2) #B,n_img,h,w,c
        image_k=self.k(image_feat)+self.image_pe #B,n_img,h,w,c
        image_v=self.v(image_feat)+self.image_pe #B,n_img,h,w,c
        image_k_v=torch.cat([image_k,image_v],dim=-1) #B,n_img,h,w,c
        unflat_k_v=self.image_unflatten(image_k_v).permute(0,4,1,2,3) #Bs,channel,n_img,reso,reso

        coords=self.triplane_coords.unsqueeze(0).unsqueeze(1).expand(batch_size,n_img,-1,-1)
        coord_homo=torch.cat([coords,torch.ones(batch_size,n_img,self.triplane_reso**2*3,1).float().cuda()],dim=-1)
        coord_inimg = torch.einsum('bjnc,bjck->bjnk', coord_homo, proj_mat.transpose(2, 3))
        x = coord_inimg[:, :, :, 0] / coord_inimg[:, :, :, 2]
        y = coord_inimg[:, :, :, 1] / coord_inimg[:, :, :, 2]
        x = x/(224.0-1) #0~1
        y = y/(224.0-1)
        coords_proj=torch.stack([x,y],dim=-1)
        coords_proj=(coords_proj-0.5)*2
        img_index=((torch.arange(n_img)[None,:,None,None].expand(
            batch_size,-1,self.triplane_reso**2*3,-1).float().cuda()/(n_img-1))-0.5)*2 #Bs,n_img,64**3,1

        grid=torch.cat([coords_proj,img_index],dim=-1) #x,y,index
        grid=torch.clamp(grid,min=-1.0,max=1.0)
        sample_k_v = torch.nn.functional.grid_sample(unflat_k_v, grid.unsqueeze(1), align_corners=True, mode='bilinear').squeeze(2) #B,C,n_img,64**3

        triplane_flat_feat=rearrange(triplane_feat,'b c h w -> b (h w) c')
        triplane_query = self.q(triplane_flat_feat) + self.triplane_pe

        k_v=rearrange(sample_k_v, 'b c n k -> (b k) n c')
        k=k_v[:,:,0:self.inner_channel]
        v=k_v[:,:,self.inner_channel:]
        q=rearrange(triplane_query,'b k c -> (b k) 1 c')
        attn_out,_=self.cross_attn(q,k,v)
        triplane_out=rearrange(attn_out,'(b h w) 1 c -> b c h w',b=batch_size,h=self.triplane_reso*3,w=self.triplane_reso)
        triplane_out = self.proj_out(triplane_out)
        return triplane_out


class MultiImage_Global_Sampler(nn.Module):
    def __init__(self,reso,vit_reso=16,padding=0.1,triplane_in_channels=64,
                 img_in_channels=1280,inner_channel=128,out_channels=64,n_heads=8,max_nimg=5):
        super().__init__()
        self.triplane_reso=reso
        self.vit_reso=vit_reso
        self.inner_channel=inner_channel
        self.padding=padding
        self.n_heads=n_heads
        self.out_channels=out_channels
        self.k=nn.Linear(in_features=img_in_channels,out_features=inner_channel)
        self.q=nn.Linear(in_features=triplane_in_channels,out_features=inner_channel)
        self.v=nn.Linear(in_features=img_in_channels,out_features=inner_channel)

        self.cross_attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)
        self.proj_out=nn.Linear(in_features=inner_channel,out_features=out_channels)
        self.image_pe = position_encoding(inner_channel, self.vit_reso**2)[None,None,:,:].expand(-1,max_nimg,-1,-1).cuda().float() #B,n_img,length,channel
        self.triplane_pe = position_encoding(inner_channel, self.triplane_reso**2*3).unsqueeze(0).cuda().float()
    def forward(self,triplane_feat,image_feat,proj_mat,valid_frames):
        '''image feat is bs,n_img,length,channel
            triplane feat is bs,C,H*3,W
        '''
        batch_size,n_img=image_feat.shape[0],image_feat.shape[1]
        L=image_feat.shape[2]-1
        image_feat=image_feat[:,:,1:,:] #discard global feature

        image_k=self.k(image_feat)+self.image_pe #B,n_img,h*w,c
        image_v=self.v(image_feat)+self.image_pe #B,n_img,h*w,c
        image_k=image_k.view(batch_size,n_img*L,-1)
        image_v=image_v.view(batch_size,n_img*L,-1)

        triplane_flat_feat=rearrange(triplane_feat,"b c h w -> b (h w) c")
        triplane_query = self.q(triplane_flat_feat) + self.triplane_pe
        #print(triplane_query.shape,image_k.shape,image_v.shape)
        attn_out,_=self.cross_attn(triplane_query,image_k,image_v)
        triplane_flat_out = self.proj_out(attn_out)
        triplane_out=rearrange(triplane_flat_out,"b (h w) c -> b c h w",h=self.triplane_reso*3,w=self.triplane_reso)

        return triplane_out

class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads

        if context_dim is None:
            context_dim = query_dim

        self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    def forward(self, q,k,v):
        h = self.heads

        q, k, v = map(lambda t: rearrange(
            t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale

        # attention, what we cannot get enough of
        attn = sim.softmax(dim=-1)

        out = torch.einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        return self.to_out(out)

class Image_Vox_Local_Sampler_Pooling(nn.Module):
    def __init__(self,reso,padding=0.1,in_channels=1280,inner_channel=128,out_channels=64,stride=4):
        super().__init__()
        self.triplane_reso=reso
        self.padding=padding
        self.get_vox_coord()
        self.out_channels=out_channels
        self.img_proj=nn.Conv2d(in_channels=in_channels,out_channels=inner_channel,kernel_size=1)

        self.vox_process=nn.Sequential(
            nn.Conv3d(in_channels=inner_channel,out_channels=inner_channel,kernel_size=3,padding=1)
        )
        self.xz_conv=nn.Sequential(
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
            nn.AvgPool3d((1,stride,1),stride=(1,stride,1)), #8
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
            nn.AvgPool3d((1,stride,1), stride=(1,stride,1)), #2
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
        )
        self.xy_conv = nn.Sequential(
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
            nn.AvgPool3d((1, 1, stride), stride=(1, 1, stride)),  # 8
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
            nn.AvgPool3d((1, 1, stride), stride=(1, 1, stride)),  # 2
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
        )
        self.yz_conv = nn.Sequential(
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
            nn.AvgPool3d((stride, 1, 1), stride=(stride, 1, 1)),  # 8
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
            nn.AvgPool3d((stride, 1, 1), stride=(stride, 1, 1)),  # 2
            nn.BatchNorm3d(inner_channel),
            nn.ReLU(),
            nn.Conv3d(in_channels=inner_channel, out_channels=inner_channel, kernel_size=3, padding=1),
        )
        self.roll_out_conv=RollOut_Conv(in_channels=inner_channel,out_channels=out_channels)
        #self.proj_out=nn.Conv2d(in_channels=inner_channel,out_channels=out_channels,kernel_size=1)
    def get_vox_coord(self):
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)

        X,Y,Z=torch.meshgrid(x,y,z,indexing='ij')
        vox_coor=torch.cat([X[:,:,:,None],Y[:,:,:,None],Z[:,:,:,None]],dim=-1)
        vox_coor=vox_coor/(self.triplane_reso-1)
        vox_coor=(vox_coor-0.5)*2*(1+self.padding+10e-6)
        self.vox_coor=vox_coor.view(-1,3).float().cuda()


    def forward(self,image_feat,proj_mat):
        image_feat=self.img_proj(image_feat)
        batch_size=image_feat.shape[0]
        vox_coords=self.vox_coor.unsqueeze(0).expand(batch_size,-1,-1) #B,64*64*64,3
        vox_homo=torch.cat([vox_coords,torch.ones((batch_size,self.triplane_reso**3,1)).float().cuda()],dim=-1)
        coord_inimg=torch.einsum('bhc,bck->bhk',vox_homo,proj_mat.transpose(1,2))
        x=coord_inimg[:,:,0]/coord_inimg[:,:,2]
        y=coord_inimg[:,:,1]/coord_inimg[:,:,2]
        x=(x/(224.0-1.0)-0.5)*2 #-1~1
        y=(y/(224.0-1.0)-0.5)*2 #-1~1

        xy=torch.cat([x[:,:,None],y[:,:,None]],dim=-1).unsqueeze(1).contiguous() #B, 1,64**3,2
        #print(image_feat.shape,xy.shape)
        grid_feat=torch.nn.functional.grid_sample(image_feat,xy,align_corners=True,mode='bilinear').squeeze(2).\
            view(batch_size,-1,self.triplane_reso,self.triplane_reso,self.triplane_reso) #B,C,1,64**3

        grid_feat=self.vox_process(grid_feat)
        xz_feat=torch.mean(self.xz_conv(grid_feat),dim=3).transpose(2,3)
        xy_feat=torch.mean(self.xy_conv(grid_feat),dim=4).transpose(2,3)
        yz_feat=torch.mean(self.yz_conv(grid_feat),dim=2).transpose(2,3)
        triplane_wImg=torch.cat([xz_feat,xy_feat,yz_feat],dim=2)
        #print(triplane_wImg.shape)

        return self.roll_out_conv(triplane_wImg)

class Image_ExpandVox_attn_Sampler(nn.Module):
    def __init__(self,reso,vit_reso=16,padding=0.1,triplane_in_channels=64,img_in_channels=1280,inner_channel=128,out_channels=64,n_heads=8):
        super().__init__()
        self.triplane_reso=reso
        self.padding=padding
        self.vit_reso=vit_reso
        self.get_vox_coord()
        self.get_vit_coords()
        self.out_channels=out_channels
        self.n_heads=n_heads

        self.kernel_func = mask_kernel_close_false
        self.k = nn.Linear(in_features=img_in_channels, out_features=inner_channel)
        # self.q_xz = nn.Conv2d(in_channels=triplane_in_channels,out_channels=inner_channel,kernel_size=1)
        # self.q_xy = nn.Conv2d(in_channels=triplane_in_channels,out_channels=inner_channel,kernel_size=1)
        # self.q_yz = nn.Conv2d(in_channels=triplane_in_channels,out_channels=inner_channel,kernel_size=1)
        self.q=nn.Linear(in_features=triplane_in_channels*3,out_features=inner_channel)

        self.v = nn.Linear(in_features=img_in_channels, out_features=inner_channel)
        self.attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channel, num_heads=n_heads, batch_first=True)
        self.out_proj=nn.Linear(in_features=inner_channel,out_features=out_channels)

        self.triplane_pe = position_encoding(inner_channel, self.triplane_reso ** 3).unsqueeze(0).cuda().float()
        self.image_pe = position_encoding(inner_channel, self.vit_reso ** 2+1).unsqueeze(0).cuda().float()
    def get_vox_coord(self):
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)

        X,Y,Z=torch.meshgrid(x,y,z,indexing='ij')
        vox_coor=torch.cat([X[:,:,:,None],Y[:,:,:,None],Z[:,:,:,None]],dim=-1)
        self.vox_index=vox_coor.view(-1,3).long().cuda()


        vox_coor = self.vox_index.float() / (self.triplane_reso - 1)
        vox_coor = (vox_coor - 0.5) * 2 * (1 + self.padding + 10e-6)
        self.vox_coor = vox_coor.view(-1, 3).float().cuda()
        # print(self.vox_coor[0])
        # print(self.vox_coor[self.triplane_reso**2])#x should increase
        # print(self.vox_coor[self.triplane_reso]) #y should increase
        # print(self.vox_coor[1])#z should increase

    def get_vit_coords(self):
        x=torch.arange(self.vit_reso)
        y=torch.arange(self.vit_reso)

        X,Y=torch.meshgrid(x,y,indexing='xy')
        vit_coords=torch.stack([X,Y],dim=-1)
        self.vit_coords=vit_coords.view(self.vit_reso**2,2).cuda().float()

    def compute_attn_mask(self,proj_coords,vit_coords,kernel_size=1.0):
        dist = torch.cdist(proj_coords.float(), vit_coords.float())
        mask = self.kernel_func(dist, sigma=kernel_size)  # True if valid, B,reso**3,vit_reso**2
        return mask


    def forward(self,triplane_feat,image_feat,proj_mat):
        xz_feat, xy_feat, yz_feat = torch.split(triplane_feat, split_size_or_sections=triplane_feat.shape[2] // 3, dim=2)  # B,C,64,64
        #xz_feat=self.q_xz(xz_feat)
        #xy_feat=self.q_xy(xy_feat)
        #yz_feat=self.q_yz(yz_feat)
        batch_size=image_feat.shape[0]
        vox_index=self.vox_index #64*64*64,3
        xz_vox_feat=xz_feat[:,:,vox_index[:,2],vox_index[:,0]].transpose(1,2) #B,C,64*64*64
        xy_vox_feat=xy_feat[:,:,vox_index[:,1],vox_index[:,0]].transpose(1,2)
        yz_vox_feat=yz_feat[:,:,vox_index[:,2],vox_index[:,1]].transpose(1,2)
        triplane_expand_feat=torch.cat([xz_vox_feat,xy_vox_feat,yz_vox_feat],dim=-1)#B,C,64*64*64,3
        triplane_query=self.q(triplane_expand_feat)+self.triplane_pe

        '''compute projection'''
        vox_coords=self.vox_coor.unsqueeze(0).expand(batch_size,-1,-1) #
        vox_homo = torch.cat([vox_coords, torch.ones((batch_size, self.triplane_reso ** 3, 1)).float().cuda()], dim=-1)
        coord_inimg = torch.einsum('bhc,bck->bhk', vox_homo, proj_mat.transpose(1, 2))
        x = coord_inimg[:, :, 0] / coord_inimg[:, :, 2]
        y = coord_inimg[:, :, 1] / coord_inimg[:, :, 2]
        #
        x = x / (224.0 - 1.0)  * (self.vit_reso-1)  # 0~self.vit_reso-1
        y = y / (224.0 - 1.0)  * (self.vit_reso-1)  # 0~self.vit_reso-1 #B,N
        xy=torch.stack([x,y],dim=-1) #B,64*64*64,2
        xy=torch.clamp(xy,min=0,max=self.vit_reso-1)
        vit_coords=self.vit_coords.unsqueeze(0).expand(batch_size,-1,-1) #B, 16*16,2
        attn_mask=torch.repeat_interleave(self.compute_attn_mask(xy,vit_coords,kernel_size=0.5),
                                          self.n_heads,0) #B*n_heads, reso**3, vit_reso**2

        k=self.k(image_feat)+self.image_pe
        v=self.v(image_feat)+self.image_pe
        attn_mask=torch.cat([torch.zeros([attn_mask.shape[0],attn_mask.shape[1],1]).cuda().bool(),attn_mask],dim=-1) #add empty token to each key and value
        vox_feat,_=self.attn(triplane_query,k,v,attn_mask=attn_mask) #B,reso**3,C
        feat_volume=self.out_proj(vox_feat).transpose(1,2).reshape(batch_size,-1,self.triplane_reso,
                                                                   self.triplane_reso,self.triplane_reso)
        xz_feat=torch.mean(feat_volume,dim=3).transpose(2,3)
        xy_feat=torch.mean(feat_volume,dim=4).transpose(2,3)
        yz_feat=torch.mean(feat_volume,dim=2).transpose(2,3)
        triplane_out=torch.cat([xz_feat,xy_feat,yz_feat],dim=2)
        return triplane_out

class Multi_Image_Fusion(nn.Module):
    def __init__(self,reso,image_reso=16,padding=0.1,img_channels=1280,triplane_channel=64,inner_channels=128,output_channel=64,n_heads=8):
        super().__init__()
        self.triplane_reso=reso
        self.image_reso=image_reso
        self.padding=padding
        self.get_triplane_coord()
        self.get_vit_coords()
        self.img_proj=nn.Conv3d(in_channels=img_channels,out_channels=512,kernel_size=1)
        self.kernel_func=mask_kernel

        self.q = nn.Linear(in_features=triplane_channel, out_features=inner_channels, bias=False)
        self.k = nn.Linear(in_features=512, out_features=inner_channels)
        self.v = nn.Linear(in_features=512, out_features=inner_channels)

        self.attn = torch.nn.MultiheadAttention(
            embed_dim=inner_channels, num_heads=n_heads, batch_first=True)
        self.out_proj=nn.Linear(in_features=inner_channels,out_features=output_channel)
        self.n_heads=n_heads

    def get_triplane_coord(self):
        '''xz plane firstly, z is at the '''
        x=torch.arange(self.triplane_reso)
        z=torch.arange(self.triplane_reso)
        X,Z=torch.meshgrid(x,z,indexing='xy')
        xz_coords=torch.cat([X[:,:,None],torch.ones_like(X[:,:,None])*(self.triplane_reso-1)/2,Z[:,:,None]],dim=-1) #in xyz order

        '''xy plane'''
        x = torch.arange(self.triplane_reso)
        y = torch.arange(self.triplane_reso)
        X, Y = torch.meshgrid(x, y, indexing='xy')
        xy_coords = torch.cat([X[:, :, None],  Y[:, :, None],torch.ones_like(X[:, :, None])*(self.triplane_reso-1)/2], dim=-1)  # in xyz order

        '''yz plane'''
        y = torch.arange(self.triplane_reso)
        z = torch.arange(self.triplane_reso)
        Y,Z = torch.meshgrid(y,z,indexing='xy')
        yz_coords= torch.cat([torch.ones_like(Y[:, :, None])*(self.triplane_reso-1)/2,Y[:,:,None],Z[:,:,None]], dim=-1)

        triplane_coords=torch.cat([xz_coords,xy_coords,yz_coords],dim=0)
        triplane_coords=triplane_coords/(self.triplane_reso-1)
        triplane_coords=(triplane_coords-0.5)*2*(1 + self.padding + 10e-6)
        self.triplane_coords=triplane_coords.float().cuda()

    def get_vit_coords(self):
        x=torch.arange(self.image_reso)
        y=torch.arange(self.image_reso)
        X,Y=torch.meshgrid(x,y,indexing='xy')
        vit_coords=torch.cat([X[:,:,None],Y[:,:,None]],dim=-1)
        self.vit_coords=vit_coords.float().cuda() #in x,y order

    def compute_attn_mask(self,proj_coord,vit_coords,valid_frames,kernel_size=2.0):
        '''
        :param proj_coord: B,K,H,W,2
        :param vit_coords: H,W,2
        :return:
        '''
        B,K=proj_coord.shape[0:2]
        vit_coords_expand=vit_coords[None,None,:,:,:].expand(B,K,-1,-1,-1)

        proj_coord=proj_coord.view(B*K,proj_coord.shape[2]*proj_coord.shape[3],proj_coord.shape[4])
        vit_coords_expand=vit_coords_expand.view(B*K,self.image_reso*self.image_reso,2)
        attn_mask=self.kernel_func(torch.cdist(proj_coord,vit_coords_expand),sigma=float(kernel_size))
        attn_mask=attn_mask.reshape(B,K,proj_coord.shape[1],vit_coords_expand.shape[1])
        valid_expand=valid_frames[:,:,None,None]
        attn_mask[valid_frames>0,:,:]=True
        attn_mask=attn_mask.permute(0,2,1,3)
        attn_mask=attn_mask.reshape(B,proj_coord.shape[1],K*vit_coords_expand.shape[1])
        atten_index=torch.where(attn_mask[0,0]==False)
        return attn_mask


    def forward(self,triplane_feat,image_feat,proj_mat,valid_frames):
        '''
        :param image_feat: B,C,K,16,16
        :param proj_mat: B,K,4,4
        :param valid_frames: B,K, true if have image, used to compute attn_mask for transformer
        :return:
        '''
        image_feat=self.img_proj(image_feat)
        batch_size=image_feat.shape[0] #K is number of frames
        K=image_feat.shape[2]
        triplane_coords=self.triplane_coords.unsqueeze(0).unsqueeze(1).expand(batch_size,K,-1,-1,-1) #B,K,192,64,3
        #print(torch.amin(triplane_coords),torch.amax(triplane_coords))
        coord_homo=torch.cat([triplane_coords,torch.ones((batch_size,K,triplane_coords.shape[2],triplane_coords.shape[3],1)).float().cuda()],dim=-1)
        #print(coord_homo.shape,proj_mat.shape)
        coord_inimg=torch.einsum('bjhwc,bjck->bjhwk',coord_homo,proj_mat.transpose(2,3))
        x=coord_inimg[:,:,:,:,0]/coord_inimg[:,:,:,:,2]
        y=coord_inimg[:,:,:,:,1]/coord_inimg[:,:,:,:,2]
        x=x/(224.0-1.0)*(self.image_reso-1)
        y=y/(224.0-1.0)*(self.image_reso-1)

        xy=torch.cat([x[...,None],y[...,None]],dim=-1) #B,K,H,W,2
        image_value=image_feat.view(image_feat.shape[0],image_feat.shape[1],-1).transpose(1,2)
        triplane_query=triplane_feat.view(triplane_feat.shape[0],triplane_feat.shape[1],-1).transpose(1,2)
        valid_frames=1.0-valid_frames.float()
        attn_mask=torch.repeat_interleave(self.compute_attn_mask(xy,self.vit_coords,valid_frames),
                                          self.n_heads,dim=0)

        q=self.q(triplane_query)
        k=self.k(image_value)
        v=self.v(image_value)
        #print(q.shape,k.shape,v.shape)

        attn,_=self.attn(q,k,v,attn_mask=attn_mask)
        #print(attn.shape)
        output=self.out_proj(attn).transpose(1,2).reshape(batch_size,-1,triplane_feat.shape[2],triplane_feat.shape[3])
        #print(output.shape)
        return output


if __name__=="__main__":
    # import sys
    # sys.path.append("../..")
    # from datasets.SingleView_dataset import Object_PartialPoints_Img
    # from datasets.transforms import Aug_with_Tran
    # #sampler=#Image_Vox_Local_Sampler_Pooling(reso=64,padding=0.1,out_channels=64,stride=4).cuda().float()
    # sampler=Image_ExpandVox_attn_Sampler(reso=32,vit_reso=16,padding=0.1,img_in_channels=1280,triplane_in_channels=64,inner_channel=64
    #                                         ,out_channels=64,n_heads=8).cuda().float()
    # # sampler=Image_Direct_AttenwMask_Sampler(reso=32,vit_reso=16,padding=0.1,img_in_channels=1280,triplane_in_channels=128,inner_channel=128
    # #                                          ,out_channels=64,n_heads=8).cuda().float()
    # dataset_config = {
    #     "data_path": "/data1/haolin/datasets",
    #     "surface_size": 20000,
    #     "par_pc_size": 4096,
    #     "load_proj_mat": True,
    # }
    # transform = Aug_with_Tran()
    # datasets = Object_PartialPoints_Img(dataset_config['data_path'], split_filename="val_par_img.json", split='val',
    #                                         transform=transform, sampling=False,
    #                                         num_samples=1024, return_surface=True, ret_sample=True,
    #                                         surface_sampling=True, par_pc_size=dataset_config['par_pc_size'],
    #                                         surface_size=dataset_config['surface_size'],
    #                                         load_proj_mat=dataset_config['load_proj_mat'], load_image=True,
    #                                         load_org_img=False, load_triplane=True, replica=1)
    #
    # dataloader = torch.utils.data.DataLoader(
    #     datasets=datasets,
    #     batch_size=64,
    #     shuffle=True
    # )
    # iterator = dataloader.__iter__()
    # data_batch = iterator.next()
    # unflatten = torch.nn.Unflatten(1, (16, 16))
    # image = data_batch['image'][:,:,:].cuda().float()
    # #image=unflatten(image).permute(0,3,1,2)
    # proj_mat = data_batch['proj_mat'].cuda().float()
    # triplane_feat=torch.randn((64,64,32*3,32)).cuda().float()
    # sampler(triplane_feat,image,proj_mat)
    # memory_usage=torch.cuda.max_memory_allocated() / MB
    # print("memory usage %f mb"%(memory_usage))


    import sys
    sys.path.append("../..")
    from datasets.SingleView_dataset import Object_PartialPoints_MultiImg
    from datasets.transforms import Aug_with_Tran

    dataset_config = {
        "data_path": "/data1/haolin/datasets",
        "surface_size": 20000,
        "par_pc_size": 4096,
        "load_proj_mat": True,
    }
    transform = Aug_with_Tran()
    dataset = Object_PartialPoints_MultiImg(dataset_config['data_path'], split_filename="train_par_img.json", split='train',
                                            transform=transform, sampling=False,
                                            num_samples=1024, return_surface=True, ret_sample=True,
                                            surface_sampling=True, par_pc_size=dataset_config['par_pc_size'],
                                            surface_size=dataset_config['surface_size'],
                                            load_proj_mat=dataset_config['load_proj_mat'], load_image=True,
                                            load_org_img=False, load_triplane=True, replica=1)

    dataloader = torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=10,
        shuffle=False
    )
    iterator = dataloader.__iter__()
    data_batch = iterator.next()
    #unflatten = torch.nn.Unflatten(2, (16, 16))
    image = data_batch['image'][:,:,:,:].cuda().float()
    #image=unflatten(image).permute(0,4,1,2,3)
    proj_mat = data_batch['proj_mat'].cuda().float()
    valid_frames = data_batch['valid_frames'].cuda().float()
    triplane_feat=torch.randn((10,128,32*3,32)).cuda().float()

    # fusion_module=MultiImage_Fuse_Sampler(reso=32,vit_reso=16,padding=0.1,img_in_channels=1280,triplane_in_channels=128,inner_channel=128
    #                                           ,out_channels=64,n_heads=8).cuda().float()
    fusion_module=MultiImage_Global_Sampler(reso=32,vit_reso=16,padding=0.1,img_in_channels=1280,triplane_in_channels=128,inner_channel=128
                                               ,out_channels=64,n_heads=8).cuda().float()
    fusion_module(triplane_feat,image,proj_mat,valid_frames)
    memory_usage=torch.cuda.max_memory_allocated() / MB
    print("memory usage %f mb"%(memory_usage))