class ModuleISEA(torch.nn.Module):
Cache = SphericalMappingCache()
Verts2Order = dict([(compute_num_vertices(o), o) for o in range(8)])
Order2Verts = dict([(o, compute_num_vertices(o)) for o in range(8)])
def __init__(self):
super(ModuleISEA, self).__init__()
def spherical_mapped_deconv(features: torch.Tensor, cache: Cache,
out_size: int, src_order: int, tgt_order: int,
kernel_size: int):
_, c, __, ___ = features.size()
dev = features.get_device()
conv = MappedTransposedConvolution(c, out_size, kernel_size * kernel_size)
conv = conv.to(dev) if dev >= 0 else conv
m, w = cache.get_kernels(src_order, tgt_order, kernel_size)
m = m.to(dev) if dev >= 0 else m
w = w.to(dev) if dev >= 0 else w
return conv.forward(features, 1, compute_num_vertices(tgt_order), m, w)
def icosahedron_resampling(img: torch.Tensor, cache: Cache, order: int):
b, c, h, w = img.size()
dev = img.get_device()
m, w = cache.get_image_samplers([h, w], order)
m = m.to(dev) if dev >= 0 else m
w = w.to(dev) if dev >= 0 else w
resample = Resample('bispherical')
vertices = resample(img, m, (1, compute_num_vertices(order)), w)
sum_weights = torch.zeros(vertices.shape[-1])
sum_weights = sum_weights.to(dev) if dev >= 0 else sum_weights
sum_weights.index_add_(0, m[..., 0].long().view(-1), w.view(-1))
vertices /= (sum_weights + 1e-12)
return vertices
kernel_sizes = [7.0, 9.0]
low_res_ws = [16]
high_res_ws = [128]
stds = [3.5, 7.0]
corners = 4
for in_order, out_order in tqdm.tqdm(zip(in_orders, out_orders),
desc="Order Loop"):
res_desc = "Order{}->{}".format(in_order, out_order)
for low_res_w, high_res_w in tqdm.tqdm(zip(low_res_ws, high_res_ws),
desc=res_desc):
res_desc = "Res = {}->{}".format(low_res_w, high_res_w)
for kernel_size in tqdm.tqdm(kernel_sizes, desc=res_desc):
ks_desc = "Kernel Size = {}".format(kernel_size)
for std in tqdm.tqdm(stds, desc=ks_desc):
kernel = gkern(kernlen=int(kernel_size), std=std)
verts = compute_num_vertices(in_order)
rm, rw = sphere_to_image_resample_map(
in_order, [low_res_w // 2, low_res_w])
ico = generate_icosphere(out_order)
vm, vw = vertex_to_vertex_kernel_map(
ico, int(kernel_size), int(kernel_size), in_order)
conv = MappedConvolution(1,
1,
kernel_size=int(kernel_size *
kernel_size),
interpolation='bispherical',
bias=False)
# conv.weight = torch.nn.Parameter(torch.ones_like(conv.weight) / (kernel_size * kernel_size))
conv.weight = torch.nn.Parameter(
torch.from_numpy(kernel.reshape(1, 1, -1)).float())
um, uw = sphere_to_image_resample_map(
__RESOLUTION_FACTOR__ = 8
__SPHERE2IMAGE_SAMPLE_MAPS__ = {}
__SPHERE2IMAGE_WEIGHTS__ = {}
__START_KERNEL_SIZE__ = 3
__MAX_KERNEL_SIZE__ = 7
__KERNEL_SIZE_LIST__ = [1, 3, 4, 5, 7]
__KERNEL_SAMPLE_MAPS__ = {}
__KERNEL_INTERP_WEIGHTS__ = {}
if __name__ == "__main__":
# init all icospheres up to an oder
for order in range(__MAX_ICOSPHERE_ORDER__ + 1):
__ICOSPHERES__[order] = generate_icosphere(order)
__MAX_VERTICES__[order] = compute_num_vertices(order)
# create sample maps for the given resolutions for each order
for s in range(__RESOLUTION_FACTOR__ + 1):
w = __START_WIDTH__ * (2**s)
h = w // 2
for o, _ in __ICOSPHERES__.items():
rersample_map, resample_weights =\
sphere_to_image_resample_map(o, [h, w], nearest=False)
key = "RES:{}x{}_ORDER:{}".format(w, h, o)
__SPHERE2IMAGE_SAMPLE_MAPS__[key] = rersample_map.cpu().numpy()
__SPHERE2IMAGE_WEIGHTS__[key] = resample_weights.cpu().numpy()
# create kernel sample maps and interpolation weights for each order and kernel size
# for k in range(__START_KERNEL_SIZE__, __MAX_KERNEL_SIZE__ + 1):
for k in __KERNEL_SIZE_LIST__:
for source_order in range(__MAX_ICOSPHERE_ORDER__ + 1):
target_order_min = max(source_order - 1, 0)
kernel_size: int):
_, c, __, ___ = features.size()
dev = features.get_device()
conv = MappedTransposedConvolution(c, out_size, kernel_size * kernel_size)
conv = conv.to(dev) if dev >= 0 else conv
m, w = cache.get_kernels(src_order, tgt_order, kernel_size)
m = m.to(dev) if dev >= 0 else m
w = w.to(dev) if dev >= 0 else w
return conv.forward(features, 1, compute_num_vertices(tgt_order), m, w)
if __name__ == "__main__":
cache = SphericalMappingCache()
verts_to_order = {}
for o in range(8):
verts_to_order[compute_num_vertices(o)] = o
order = 5
kernel_size = 3
out_features = 32
# resolution = [256, 512]
# map, wgts = cache.get_image_samplers(resolution, order)
# print(map.shape)
# print(wgts.shape)
# kernel_size = 3
# source_icosphere_order = 5
# target_icosphere_order = 4
# samples, params = cache.get_kernels(source_icosphere_order, target_icosphere_order, kernel_size)
# print(samples.shape)
# print(params.shape)
#TODO: create mappedconv, mappeddownconv, mappedupconv, mappeddeconv, mappedupdeconv