Spaces:
Running on Zero
Running on Zero
File size: 4,852 Bytes
740d966 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | import torch,os,sys
code_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(f'{code_dir}/../../')
import torch.nn.functional as F
import numpy as np
class InputPadder:
""" Pads images such that dimensions are divisible by 8 """
def __init__(self, dims, mode='sintel', divis_by=8, force_square=False):
self.ht, self.wd = dims[-2:]
if force_square:
max_side = max(self.ht, self.wd)
pad_ht = ((max_side // divis_by) + 1) * divis_by - self.ht
pad_wd = ((max_side // divis_by) + 1) * divis_by - self.wd
else:
pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by
pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by
if mode == 'sintel':
self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
else:
self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
def pad(self, *inputs):
assert all((x.ndim == 4) for x in inputs)
return [F.pad(x, self._pad, mode='replicate') for x in inputs]
def unpad(self, x):
assert x.ndim == 4
ht, wd = x.shape[-2:]
c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
return x[..., c[0]:c[1], c[2]:c[3]]
def bilinear_sampler1d(img, x_coords, mode='bilinear', align_corners=True):
"""
1D bilinear sampling along width dimension only (for stereo applications)
Much faster than grid_sample for stereo where y is constant
Args:
img: (B, C, 1, W) input tensor
x_coords: (B, 1, W_out, 1) x coordinates in pixel space [0, W-1]
mode: interpolation mode ('bilinear' or 'nearest')
align_corners: if True, corner pixels are aligned (like grid_sample)
Returns:
sampled: (B, C, 1, W_coords) sampled tensor
mask: (B, 1, H, W) validity mask (if mask=True)
"""
B, C, H_img, W = img.shape
x = x_coords.reshape(B,-1) # (B, W_out)
if align_corners:
# align_corners=True: coordinate range [0, W-1] maps to pixel centers
# This matches grid_sample with align_corners=True behavior
x_normalized = x
else:
# align_corners=False: coordinate range [0, W-1] maps to pixel edges
# Need to adjust coordinates to match grid_sample with align_corners=False
# grid_sample maps [-1, 1] to [0, W-1] when align_corners=False
# So our [0, W-1] input should be treated as [0.5, W-0.5] in pixel space
x_normalized = x + 0.5
if mode == 'nearest':
# Nearest neighbor sampling with zero padding outside [0, W-1]
if align_corners:
x_nearest = torch.round(x_normalized).long()
else:
x_nearest = torch.floor(x_normalized).long()
valid = (x_nearest >= 0) & (x_nearest < W) # (B, W_out)
x_index = torch.clamp(x_nearest, 0, W-1)
sampled = torch.gather(img, 3, x_index.view(B,1,1,-1).expand(B,C,1,-1))
sampled = sampled * valid.view(B,1,1,-1).to(img.dtype)
else: # bilinear
# Get integer and fractional parts
x_floor = torch.floor(x_normalized)
x_ceil = x_floor + 1
x_frac = x_normalized - x_floor # (B, W_out)
# Zero padding behavior: mark validity and zero-out invalid contributions
valid_floor = (x_floor >= 0) & (x_floor < W)
valid_ceil = (x_ceil >= 0) & (x_ceil < W)
x_floor_clamped = torch.clamp(x_floor, 0, W-1)
x_ceil_clamped = torch.clamp(x_ceil, 0, W-1)
# Create index tensors
batch_idx = torch.arange(B, device=img.device).view(B, 1)
img_floor = torch.gather(img, 3, x_floor_clamped.view(B,1,1,-1).expand(B,C,1,-1).long())
img_ceil = torch.gather(img, 3, x_ceil_clamped.view(B,1,1,-1).expand(B,C,1,-1).long())
# Apply validity masks (zero out-of-bounds samples)
img_floor = img_floor * valid_floor.view(B,1,1,-1).to(img.dtype)
img_ceil = img_ceil * valid_ceil.view(B,1,1,-1).to(img.dtype)
# Linear interpolation
x_frac = x_frac.view(B,1,1,-1)
sampled = img_floor * (1 - x_frac) + img_ceil * x_frac
return sampled
def bilinear_sampler(img, coords, mode='bilinear', mask=False, low_memory=False, use1d=False):
""" Wrapper for grid_sample, uses pixel coordinates """
H, W = img.shape[-2:]
coords[...,0] = 2*coords[...,0]/(W-1) - 1
if low_memory:
B = img.shape[0]
out = []
bs = 102400
for b in np.arange(0,B,bs):
tmp = F.grid_sample(img[b:b+bs], coords[b:b+bs], align_corners=True)
out.append(tmp)
img = torch.cat(out, dim=0)
else:
img = F.grid_sample(img, coords, align_corners=True)
if mask:
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
return img, mask.float()
return img
|