Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,539 Bytes
c295391 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
"""
Scalable, Detailed and Mask-free Universal Photometric Stereo Network (CVPR2023)
# Copyright (c) 2023 Satoshi Ikehata
# All rights reserved.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
def divide_tensor_spatial(x, block_size=256, method='tile_stride'):
assert x.dim() == 4, "Input tensor must have 4 dimensions [B, C, H, W]"
B, C, H, W = x.shape
assert H == W, "Height and Width must be equal"
assert H % block_size == 0 and W % block_size ==0, "The tensor size cannot be divided by the block size"
mosaic_scale = H // block_size
if method == 'tile_stride':
""" decomposing x into K x K of (Hc, Wc) non-overlapped blocks (grid)"""
K = mosaic_scale * mosaic_scale
fold_params_grid = dict(kernel_size=(mosaic_scale, mosaic_scale), stride=(mosaic_scale, mosaic_scale), padding=(0,0), dilation=(1,1))
unfold_grid = nn.Unfold(**fold_params_grid)
tensor_grids = unfold_grid(x) # (B, C * K, Hm * Hm)
tensor_grids = tensor_grids.reshape(B, C, K, block_size, block_size).permute(0, 2, 1, 3, 4) # (B, K, C, Hm, Hm)
return tensor_grids
if method == 'tile_block':
tensor_blocks = x.view(B, C, mosaic_scale, block_size, mosaic_scale, block_size)
tensor_blocks = tensor_blocks.permute(0, 2, 4, 1, 3, 5) # (B, mc, mc, C, Hm, Wm)
tensor_blocks = tensor_blocks.contiguous().view(B, mosaic_scale**2, C, block_size, block_size) ## (B, K, C, Hm, Hm)
return tensor_blocks
return -1
def merge_tensor_spatial(x, method='tile_stride'):
K, N, feat_dim, Hm, Wm = x.shape
mosaic_scale = int(math.sqrt(K))
if method == 'tile_stride':
x = x.reshape(K, N, feat_dim, -1)
fold_params_grid = dict(kernel_size=(mosaic_scale, mosaic_scale), stride=(mosaic_scale, mosaic_scale), padding=(0,0), dilation=(1,1))
fold_grid = nn.Fold(output_size=(Hm * mosaic_scale, Wm * mosaic_scale), **fold_params_grid) # downsample based on the encoder
x = x.permute(1, 2, 0, 3).reshape(N, feat_dim * K, -1)
x = fold_grid(x)
return x
if method == 'tile_block':
x = x.permute(1, 0, 2, 3, 4).reshape(N, mosaic_scale, mosaic_scale, feat_dim, Hm, Wm)
x = x.permute(0, 3, 1, 4, 2, 5)
x = x.reshape(N, feat_dim, mosaic_scale * Hm, mosaic_scale * Wm)
return x
def divide_overlapping_patches(input_tensor, patch_size, margin):
B, C, W, _ = input_tensor.shape
stride = patch_size - margin
padded_W = ((W - patch_size + stride - 1) // stride) * stride + patch_size
pad = padded_W - W
padded_tensor = F.pad(input_tensor, (0, pad, 0, pad), mode='constant', value=0)
patches = F.unfold(padded_tensor, kernel_size=patch_size, stride=stride)
patches = patches.view(B, C, patch_size, patch_size, -1).permute(0, 4, 1, 2, 3)
return patches
def merge_overlappnig_patches(patches, patch_size, margin, original_size):
B, _, C, _, _ = patches.shape
stride = patch_size - margin
W = original_size[2]
patches = patches.permute(0, 2, 3, 4, 1).contiguous().view(B, C * patch_size * patch_size, -1)
output = F.fold(patches, (W, W), kernel_size=patch_size, stride=stride)
weight = torch.ones(patches.size()).to(patches.device)
weight = F.fold(weight, (W, W), kernel_size=patch_size, stride=stride)
output = output / weight
return output
|