wehub-resource-sync
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import megengine as mge
|
|
import megengine.functional as F
|
|
import numpy as np
|
|
|
|
|
|
def bilinear_sampler(img, coords, mode="bilinear", mask=False):
|
|
"""Wrapper for grid_sample, uses pixel coordinates"""
|
|
H, W = img.shape[-2:]
|
|
|
|
img = F.remap(img, coords, border_mode="constant")
|
|
|
|
if mask:
|
|
mask = (
|
|
(coords[:, :, :, 0:1] < 0)
|
|
| (coords[:, :, :, 0:1] > W - 1)
|
|
| (coords[:, :, :, 1:2] < 0)
|
|
| (coords[:, :, :, 1:2] > H - 1)
|
|
)
|
|
mask = F.logical_not(mask)
|
|
return img, mask.astype("float32")
|
|
|
|
return img
|
|
|
|
|
|
def coords_grid(batch, ht, wd):
|
|
x_grid, y_grid = np.meshgrid(np.arange(wd), np.arange(ht))
|
|
y_grid, x_grid = mge.tensor(y_grid, dtype="float32"), mge.tensor(
|
|
x_grid, dtype="float32"
|
|
)
|
|
coords = F.stack([x_grid, y_grid], axis=0)
|
|
coords = F.repeat(F.expand_dims(coords, axis=0), batch, axis=0)
|
|
return coords
|
|
|
|
|
|
def manual_pad(x, pady, padx):
|
|
if pady > 0:
|
|
u = F.repeat(x[:, :, 0:1, :], pady, axis=2)
|
|
d = F.repeat(x[:, :, -1:, :], pady, axis=2)
|
|
x = F.concat([u, x, d], axis=2)
|
|
if padx > 0:
|
|
l = F.repeat(x[:, :, :, 0:1], padx, axis=3)
|
|
r = F.repeat(x[:, :, :, -1:], padx, axis=3)
|
|
x = F.concat([l, x, r], axis=3)
|
|
return x
|