File size: 729 Bytes
5db43ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
import torch
import numpy as np
def diff_tex_sampling(uv_map: np.ndarray, tex: torch.Tensor) -> torch.Tensor:
assert tex.shape[2] == 3
height = uv_map.shape[0]
width = uv_map.shape[1]
grid = torch.from_numpy(uv_map.astype(np.float32)/255.0).unsqueeze(0)#B Hout Wout 2
grid = grid*2.0 -1.0
input = tex.permute(2,0,1).unsqueeze(0)
#print(input.shape)
#print(grid.shape)
output = torch.nn.functional.grid_sample(input.clone().cuda(), grid.cuda(), mode='bilinear', padding_mode='zeros', align_corners=False)
return output
if __name__ == '__main__':
uv_map = np.ones([512,512,2],np.uint8)
tex = torch.ones((256,256,3))
img = diff_tex_sampling(uv_map,tex)
print(img.shape)
|