代码拉取完成，页面将自动刷新
import torch
import numpy as np
import gurobipy as gp
from gurobipy import GRB
import copy
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.patches as patches
class QuadTree():
def __init__(self, sidelength, patch_size):
self.sidelength = sidelength
self.patch_size = patch_size
self.min_patch_size = np.min(patch_size)
self.max_patch_size = np.min(sidelength)
self.aspect_ratio = np.array(patch_size) / np.min(patch_size)
# how many levels of quadtree are there
self.min_quadtree_level = int(np.log2(np.min(self.sidelength) // self.max_patch_size))
self.max_quadtree_level = int(np.log2(np.min(sidelength) // self.min_patch_size))
self.num_scales = self.max_quadtree_level - self.min_quadtree_level + 1
# optimization model
self.optim_model = gp.Model()
self.optim_model.setParam('OutputFlag', 0)
self.c_max_patches = None
# initialize tree
self.root = self.init_root(self.max_quadtree_level)
# populate tree nodes with coordinate values, metadata
self.populate_tree()
def __deepcopy__(self, memo):
deep_copied_obj = QuadTree(self.sidelength, self.patch_size)
for k, v in self.__dict__.items():
if k in ['optim_model', 'c_max_patches']:
# setattr(deep_copied_obj, k, v)
del(deep_copied_obj.__dict__[k])
else:
setattr(deep_copied_obj, k, copy.deepcopy(v, memo))
return deep_copied_obj
def __getstate__(self):
state = self.__dict__.copy()
for k, v in self.__dict__.items():
if k in ['optim_model', 'c_max_patches']:
del(state[k])
return state
def __load__(self, obj):
for k, v in obj.__dict__.items():
if k == 'root':
continue
setattr(self, k, v)
self.root = self.init_root(obj.max_quadtree_level)
self.populate_tree()
def _load_helper(curr_patch, curr_obj_patch):
curr_patch.__load__(curr_obj_patch)
for child, obj_child in zip(curr_patch.children, curr_obj_patch.children):
_load_helper(child, obj_child)
return
_load_helper(self.root, obj.root)
def __str__(self, level=0):
def _str_helper(curr_patch, level):
ret = "\t"*level+repr(curr_patch.active)+"\n"
for child in curr_patch.children:
ret += _str_helper(child, level+1)
return ret
return _str_helper(self.root, 0)
def populate_tree(self):
# get block coords for patches at each scale
patch_sizes = []
curr_size = self.max_patch_size
while True:
patch_sizes.append(curr_size)
curr_size //= 2
if curr_size == self.min_patch_size:
patch_sizes.append(curr_size)
break
elif curr_size < self.min_patch_size:
raise ValueError('Patch sizes and resolution are incompatible')
block_coords = [self.get_block_coords(patch_size=patch_size, include_ends=True) for patch_size in patch_sizes]
block_sizes = [block[1, 1, :] - block[0, 0, :] for block in block_coords]
block_coords = [block[:-1, :-1, :] for block in block_coords]
# create sampling grids for training
num_samples = self.min_patch_size * self.aspect_ratio
row_posts = torch.linspace(-1, 1, int(self.min_patch_size*self.aspect_ratio[0])+1)[:-1]
col_posts = torch.linspace(-1, 1, int(self.min_patch_size*self.aspect_ratio[1])+1)[:-1]
row_coords, col_coords = torch.meshgrid(row_posts, col_posts)
row_coords = row_coords.flatten()
col_coords = col_coords.flatten()
# create sampling grids for evaluation
# here we need to sample every pixel within each block
row_posts = [torch.linspace(-1, 1, int(pixel_size*self.aspect_ratio[0])+1)[:-1] for pixel_size in patch_sizes]
col_posts = [torch.linspace(-1, 1, int(pixel_size*self.aspect_ratio[1])+1)[:-1] for pixel_size in patch_sizes]
eval_coords = [torch.meshgrid(row_post, col_post) for row_post, col_post in zip(row_posts, col_posts)]
eval_row_coords = [eval_coord[0].flatten() for eval_coord in eval_coords]
eval_col_coords = [eval_coord[1].flatten() for eval_coord in eval_coords]
def _populate_tree_helper(patch, idx):
# get block scale idx
scale_idx = len(idx) - (self.min_quadtree_level)
# do we have patches at this level?
if scale_idx >= 0:
# set patch parameters
coords = block_coords[scale_idx]
coord_idx = _index_block_coord(idx, coords.shape[0], coord=[0, 0])
patch.block_coord = coords[coord_idx[0], coord_idx[1]]
patch.block_size = block_sizes[scale_idx]
patch.scale = scale_idx
patch.pixel_size = patch_sizes[scale_idx]
patch.num_samples = num_samples
patch.row_coords = row_coords
patch.col_coords = col_coords
patch.eval_row_coords = eval_row_coords[scale_idx]
patch.eval_col_coords = eval_col_coords[scale_idx]
if not patch.children:
return
# recurse
for i in range(4):
child = patch.children[i]
_populate_tree_helper(child, [*idx, i])
return
# given list of tree idxs in {0,1,2,3}^N, retrieve the block coordinate
def _index_block_coord(tree_idx, length, coord=[0, 0]):
if length == 1:
return coord
if tree_idx[0] == 0:
pass
elif tree_idx[0] == 1:
coord[1] += length//2
elif tree_idx[0] == 2:
coord[0] += length//2
elif tree_idx[0] == 3:
coord[0] += length//2
coord[1] += length//2
else:
raise ValueError("Unexpected child value, should be 0, 1, 2, or 3")
return _index_block_coord(tree_idx[1:], length//2, coord)
# done with setup, now actually populate the tree
_populate_tree_helper(self.root, [])
def init_root(self, max_level):
def _init_root_helper(curr_patch, curr_level, max_level, optim_model):
if curr_level == max_level:
return
curr_patch.children = [Patch(optim_model) for _ in range(4)]
for patch in curr_patch.children:
patch.parent = curr_patch
_init_root_helper(patch, curr_level+1, max_level, optim_model)
return
# create root node
root = Patch(self.optim_model)
_init_root_helper(root, 0, max_level, self.optim_model)
return root
def get_block_coords(self, flatten=False, include_ends=False, patch_size=None):
patch_size = patch_size * self.aspect_ratio
# get size of each block
block_size = (2 / (self.sidelength[0]-1) * patch_size[0], 2 / (self.sidelength[1]-1) * patch_size[1])
# get block begin/end coordinates
if include_ends:
block_coords_y = torch.arange(-1, 1+block_size[0], block_size[0])
block_coords_x = torch.arange(-1, 1+block_size[1], block_size[1])
else:
block_coords_y = torch.arange(-1, 1, block_size[0])
block_coords_x = torch.arange(-1, 1, block_size[1])
# repeat for every single block
block_coords = torch.meshgrid(block_coords_y, block_coords_x)
block_coords = torch.stack((block_coords[0], block_coords[1]), dim=-1)
if flatten:
block_coords = block_coords.reshape(-1, 2)
return block_coords
def get_patches_at_level(self, level):
# level is the image scale: 0-> coarsest, N->finest
if level == -1:
level = self.max_quadtree_level
# what quadtree level do our patches start at?
# check input, too
target_level = level + self.min_quadtree_level
assert level <= (self.max_quadtree_level - self.min_quadtree_level), \
"invalid 'level' input to get_blocks_at_level"
def _get_patches_at_level_helper(curr_patch, curr_level, patches):
if curr_level > target_level:
return
for patch in curr_patch.children:
_get_patches_at_level_helper(patch, curr_level+1, patches)
if curr_level == target_level:
patches.append(curr_patch)
return patches
return _get_patches_at_level_helper(self.root, 0, [])
def get_frozen_patches(self):
def _get_frozen_patches_helper(curr_patch, patches):
if curr_patch.frozen and curr_patch.active:
patches.append(curr_patch)
for patch in curr_patch.children:
_get_frozen_patches_helper(patch, patches)
return patches
return _get_frozen_patches_helper(self.root, [])
def get_active_patches(self, include_frozen_patches=False):
def _get_active_patches_helper(curr_patch, patches):
if curr_patch.active and \
(include_frozen_patches or
(not include_frozen_patches and not curr_patch.frozen)):
patches.append(curr_patch)
for patch in curr_patch.children:
_get_active_patches_helper(patch, patches)
return patches
return _get_active_patches_helper(self.root, [])
def activate_random(self):
def _activate_random_helper(curr_patch):
if not curr_patch.children:
curr_patch.activate()
return
elif (curr_patch.scale is not None) and (torch.rand(1).item() < 0.2):
curr_patch.activate()
return
for patch in curr_patch.children:
_activate_random_helper(patch)
return
_activate_random_helper(self.root)
def synchronize(self, master):
# set active/inactive nodes to be the same as master
# for now just toggle the flags without worrying about the gurobi variables
def _synchronize_helper(curr_patch, curr_patch_master):
curr_patch.active = curr_patch_master.active
if not curr_patch.children:
return
for patch, patch_master in zip(curr_patch.children, curr_patch_master.children):
_synchronize_helper(patch, patch_master)
return
_synchronize_helper(self.root, master.root)
def get_frozen_samples(self):
patches = self.get_frozen_patches()
if not patches:
return None, None, None, None
rel_coords, abs_coords, vals = [], [], []
patch_idx = []
for idx, p in enumerate(patches):
rel_samp, abs_samp = p.get_stratified_samples(jitter=False, eval=True)
rel_samp = rel_samp.reshape(-1, int(np.prod(self.min_patch_size * self.aspect_ratio)), 2)
abs_samp = abs_samp.reshape(-1, int(np.prod(self.min_patch_size * self.aspect_ratio)), 2)
patch_idx.extend(rel_samp.shape[0] * [idx, ])
rel_coords.append(rel_samp)
abs_coords.append(abs_samp)
# values have the same size as rel_samp but last dim is a scalar
vals.append(p.value*torch.ones(abs_samp.shape[:-1] + (1,)))
return torch.cat(rel_coords, dim=0), torch.cat(abs_coords, dim=0), \
torch.cat(vals, dim=0), patch_idx
def get_stratified_samples(self, jitter=True, eval=False):
patches = self.get_active_patches()
rel_coords, abs_coords = [], []
patch_idx = []
for idx, p in enumerate(patches):
rel_samp, abs_samp = p.get_stratified_samples(jitter=jitter, eval=eval)
# always batch the coordinates in groups of a specific patch size
# so we can process them in parallel
rel_samp = rel_samp.reshape(-1, int(np.prod(self.min_patch_size * self.aspect_ratio)), 2)
abs_samp = abs_samp.reshape(-1, int(np.prod(self.min_patch_size * self.aspect_ratio)), 2)
# since patch samples could be split across batches,
# keep track of which batch idx maps to which patch idx
patch_idx.extend(rel_samp.shape[0] * [idx, ])
rel_coords.append(rel_samp)
abs_coords.append(abs_samp)
return torch.cat(rel_coords, dim=0), torch.cat(abs_coords, dim=0), patch_idx
def solve_optim(self, max_num_patches=1024):
patches = self.get_active_patches()
assert (len(patches) <= max_num_patches), \
"You are trying to solve a model which is infeasible: " \
"Number of active patches > Max number of patches"
if self.c_max_patches is not None:
self.optim_model.remove(self.c_max_patches)
# global "knapsack" constraint
expr_sum_patches = [p.update_merge() for p in patches]
self.c_max_patches = self.optim_model.addConstr(gp.quicksum(expr_sum_patches) <= max_num_patches)
# objective
self.optim_model.setObjective(gp.quicksum([p.get_cost() for p in patches]), GRB.MINIMIZE)
self.optim_model.optimize()
obj_val = self.optim_model.objVal
if self.optim_model.Status == GRB.INFEASIBLE:
print("----------- Model is infeasible")
self.optim_model.computeIIS()
self.optim_model.write("model.ilp")
# split and merge
merged = 0
split = 0
none = 0
for p in patches:
# print(p)
if p.has_split() and p.scale < self.max_quadtree_level:
p.deactivate()
for child in p.get_children():
child.activate()
split += 1
elif p.has_merged() and p.scale >= self.min_quadtree_level and p.scale > 0:
# we first check if it is active,
# since we could have already been activated by a neighbor
if p.active:
for neighbor in p.get_neighbors():
neighbor.deactivate()
p.parent.activate()
merged += 1
else:
p.update()
none += 1
stats_dict = {'merged': merged,
'splits': split,
'none': none,
'obj': obj_val}
print(f"============================= Total patches:{len(patches)}, split/merge:{split}/{merged}")
return stats_dict
def draw(self):
fig, ax = plt.subplots(1, figsize=(5, 5))
depth = 1 + self.max_quadtree_level - self.min_quadtree_level
sidelen = 4**(depth-1) // 2**(depth-1)
# calculate scale
patch_list = self.get_active_patches()
patches_err = [p.err for p in patch_list]
max_err = np.max(patches_err)
min_err = np.min(patches_err)
cmap = plt.cm.get_cmap('viridis')
def _draw_level(patch, curr_level, ax, sidelen, offset, scale):
if curr_level > self.max_quadtree_level:
return ax
scale = scale/2.
for i, child in enumerate(patch.children):
if i == 0:
new_offset = (offset[0], offset[1])
elif i == 1:
new_offset = (offset[0] + scale * sidelen, offset[1])
elif i == 2:
new_offset = (offset[0], offset[1] + scale * sidelen)
else:
new_offset = (offset[0] + scale * sidelen, offset[1] + scale * sidelen)
if child.active:
norm_err = (child.err-min_err)/(max_err-min_err)
if child.frozen:
facecolor = 'white'
edgecolor = 'red'
else:
facecolor = cmap(norm_err)
edgecolor = 'black'
rect = patches.Rectangle(new_offset, scale * sidelen, scale * sidelen, linewidth=1,
edgecolor=edgecolor,
facecolor=facecolor, fill=True)
ax.add_patch(rect)
else:
ax = _draw_level(child, curr_level+1, ax, sidelen, new_offset, scale)
return ax
ax = _draw_level(self.root, self.min_quadtree_level, ax, sidelen, (0., 0.), 1.)
ax.set_aspect('equal')
plt.xlim(-1, sidelen + 1)
plt.ylim(-1, sidelen + 1)
plt.gca().invert_yaxis() # we want 0,0 to be on top-left
norm = colors.Normalize(vmin=min_err, vmax=max_err)
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
plt.colorbar(sm)
return fig
# patch class
class Patch():
def __init__(self, optim_model=None, block_coord=None, scale=None):
self.active = False
self.parent = None
self.children = []
# absolute block coordinate
self.block_coord = block_coord
# size of block in absolute coord frame
self.block_size = None
# scale level of block
self.scale = scale
# num samples to be generated for this block
self.num_samples = None
# num pixels in this patch
self.pixel_size = None
# optimization model
self.optim = optim_model
# row/column coords for sampling at test time
# initialized by set_samples() function
self.row_coords = None
self.col_coords = None
self.eval_row_coords = None
self.eval_col_coords = None
# error for doing nothing, merging, splitting
self.err = 0.
self.last_updated = 0.
self._nocopy = ['optim', 'I_grp', 'I_split', 'I_none',
'I_merge', 'c_joinable', 'c_merge_split']
self._pickle_vars = ['parent', 'children', 'active', 'err', 'last_updated']
self.spec_cstrs = []
# options for pruning
self.frozen = False
self.value = 0.0
def __str__(self):
str = f"Patch id={id(self)}\n" \
f" . active={self.active}\n" \
f" . level={self.scale}\n" \
f" . model={self.optim}"
if self.active:
str += f"\n . g={self.I_grp.x}, s={self.I_split.x}, n={self.I_none.x}"
return str
# override deep copy to copy undeepcopyable objects by reference
def __deepcopy__(self, memo):
deep_copied_obj = Patch()
for k, v in self.__dict__.items():
if k in self._nocopy:
# setattr(deep_copied_obj, k, None)
if k in deep_copied_obj.__dict__.keys():
del(deep_copied_obj.__dict__[k])
else:
setattr(deep_copied_obj, k, copy.deepcopy(v, memo))
return deep_copied_obj
def __getstate__(self):
state = self.__dict__.copy()
for k, v in self.__dict__.items():
if k in self._nocopy:
# if k not in self._pickle_vars:
del(state[k])
return state
def __load__(self, obj):
for k, v in obj.__dict__.items():
if k in ['children', 'parent']:
continue
setattr(self, k, v)
if self.active:
self.activate()
def update(self):
self.deactivate()
self.activate()
def activate(self):
self.active = True
# indicator variables
self.I_grp = self.optim.addVar(vtype=GRB.BINARY)
self.I_split = self.optim.addVar(vtype=GRB.BINARY)
self.I_none = self.optim.addVar(vtype=GRB.BINARY)
self.I_merge = gp.LinExpr(0.0)
# local constraint "merge/none/split"
self.c_joinable = self.optim.addConstr(self.I_grp + self.I_none + self.I_split == 1)
# local constraint "merge-split"
self.c_merge_split = None
def deactivate(self):
self.active = False
self.optim.remove(self.I_grp)
self.optim.remove(self.I_split)
self.optim.remove(self.I_none)
self.I_merge = gp.LinExpr(0.0)
self.optim.remove(self.c_joinable)
if self.c_merge_split is not None:
self.optim.remove(self.c_merge_split)
for cstr in self.spec_cstrs:
self.optim.remove(cstr)
self.spec_cstrs = []
def is_mergeable(self):
siblings = self.parent.children
return np.all(np.all([sib.active for sib in siblings]))
def set_sample_params(self, num_samples):
self.num_samples = num_samples
posts = torch.linspace(-1, 1, self.num_samples+1)[:-1]
row_coords, col_coords = torch.meshgrid(posts, posts)
self.row_coords = row_coords.flatten()
self.col_coords = col_coords.flatten()
def must_split(self):
self.spec_cstrs.append(
self.optim.addConstr(self.I_split == 1)
)
def must_merge(self):
self.spec_cstrs.append(
self.optim.addConstr(self.I_grp == 1)
)
def has_split(self):
return self.I_split.x == 1
def has_merged(self):
return self.I_grp.x == 1
# return self.I_none.x==0 and self.I_split.x==0
def has_done_nothing(self):
return self.I_none.x == 1
def get_cost(self):
area = self.block_size[0]**2
alpha = 0.2 # how much worse we expect the error to be when merging
beta = -0.02 # how much better we expect the error to be when splitting
# == Merge
if self.scale > 0: # it should never be root, but still..
err_merge = (4+alpha) * area * self.err
if self.parent.last_updated:
parent_area = self.parent.block_size[0]**2
err_merge = parent_area * self.parent.err # can multiply by 1/4 as in paper to make merging more aggressive
else:
err_merge = self.err
# == Split
if self.children:
err_split = (0.25+beta) * area * self.err
if self.children[0].last_updated:
err_children = np.sum([child.err for child in self.children])
err_split = area * err_children
else:
err_split = 1. # in case you don't have children, high to avoid splitting
# == None
err_none = area * self.err
return err_none * self.I_none \
+ err_split * self.I_split \
+ err_merge * self.I_grp
def update_merge(self):
if self.parent is None: # if root
return gp.LinExpr(0)
siblings = self.parent.children
if np.all([sib.active for sib in siblings]):
I_grp_neighs = [s.I_grp for s in siblings]
self.I_merge = gp.quicksum(I_grp_neighs)
# local constraint "joinable"
self.c_merge_split = self.optim.addConstr(self.I_none + self.I_split + .25*self.I_merge == 1)
expr_max_patches = 4 * self.I_split + 1 * self.I_none + .25 * self.I_grp
return expr_max_patches
def get_neighbors(self):
return self.parent.children
def get_children(self):
return self.children
def get_parent(self):
return self.parent
def is_joinable(self):
# test if siblings are all leaf nodes
siblings = self.parent.children
return np.all([sib.active for sib in siblings])
def get_block_coord(self):
return self.block_coord
def get_scale(self):
return self.scale
def update_error(self, error, iter):
self.err = error
self.last_updated = iter
def get_stratified_samples(self, jitter=True, eval=False):
# Block coords are always aligned to the pixel grid,
# e.g., they align with pixels 0, 8, 16, 24, etc. for
# patch size 8
#
# To normalize the coordinates between (-1, 1), consider
# we have an image of 64x64 and patch size 8x8.
# The block coordinate (-1, -1) aligns with pixel (0, 0)
# and coordinate (1, 1) aligns with pixel (63, 63)
#
# Absolute coordinates within a block should stretch all the way
# from the absolute position of one block coordinate to another.
# Say each block contains 8x8 pixels and we use a feature grid
# of 8x8 features to interpolate values within a block.
# This means is that the feature positions are not actually
# aligned to the pixel positions. The features are positioned
# on a grid stretching from one block coord to another whereas
# the pixel grid ends just short of the next block coordinate
#
# Example patch (x = pixel position, B = block coordinate position)
# and relative coordinate positions.
#
# -1 ^ B x x x x x x x B
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# 1 v B x x x x x x x B
# <--------------->
# -1 1
#
# When we generate samples for a patch, we sample an
# 8x8 grid that extends between block coords, i.e.
# between the arrows above
#
if eval:
row_coords = self.eval_row_coords.flatten()
col_coords = self.eval_col_coords.flatten()
else:
row_coords = self.row_coords
col_coords = self.col_coords
if jitter:
row_coords = self.row_coords + torch.rand_like(self.row_coords) * 2./self.num_samples[0]
col_coords = self.col_coords + torch.rand_like(self.col_coords) * 2./self.num_samples[1]
rel_samples = torch.stack((row_coords, col_coords), dim=-1)
abs_samples = self.block_coord[None, :] + self.block_size[None, :] * (rel_samples+1)/2
return rel_samples, abs_samples
class OctTree():
def __init__(self, sidelength, min_octant_size, bounds=((-1, 1), (-1, 1), (-1, 1)), mesh_kd_tree=None):
self.sidelength = sidelength
self.min_octant_size = min_octant_size
self.max_octant_size = sidelength[0]
# how many levels of quadtree are there
self.min_octtree_level = int(np.log2(self.sidelength[0] // self.max_octant_size))
self.max_octtree_level = int(np.log2(sidelength[0] // min_octant_size))
self.num_scales = self.max_octtree_level - self.min_octtree_level + 1
# optimization model
self.optim_model = gp.Model()
self.optim_model.setParam('OutputFlag', 0)
self.c_max_octants = None
# set bounds
self.z_min, self.z_max = bounds[0]
self.y_min, self.y_max = bounds[1]
self.x_min, self.x_max = bounds[2]
# KD tree that stores points on mesh surface
self.surface_tree = mesh_kd_tree
# initialize tree
self.root = self.init_root(self.max_octtree_level)
# populate tree nodes with coordinate values, metadata
self.populate_tree()
def __deepcopy__(self, memo):
deep_copied_obj = OctTree(self.sidelength, self.min_octant_size)
for k, v in self.__dict__.items():
if k in ['optim_model', 'c_max_octants']:
setattr(deep_copied_obj, k, v)
else:
setattr(deep_copied_obj, k, copy.deepcopy(v, memo))
return deep_copied_obj
def __getstate__(self):
state = self.__dict__.copy()
for k, v in self.__dict__.items():
if k in ['optim_model', 'c_max_octants']:
del(state[k])
return state
def __load__(self, obj):
for k, v in obj.__dict__.items():
if k == 'root':
continue
setattr(self, k, v)
self.root = self.init_root(obj.max_octtree_level)
self.populate_tree()
def _load_helper(curr_patch, curr_obj_patch):
curr_patch.__load__(curr_obj_patch)
for child, obj_child in zip(curr_patch.children, curr_obj_patch.children):
_load_helper(child, obj_child)
return
_load_helper(self.root, obj.root)
def __str__(self, level=0):
def _str_helper(curr_octant, level):
ret = "\t"*level+repr(curr_octant.active)+"\n"
for child in curr_octant.children:
ret += _str_helper(child, level+1)
return ret
return _str_helper(self.root, 0)
def populate_tree(self):
# maximum octant scale
max_octant_scale = int(np.log2(self.max_octant_size))
min_octant_scale = int(np.log2(self.min_octant_size))
# get block coords for octants at each scale
octant_sizes = [2**s for s in range(min_octant_scale, max_octant_scale+1)]
octant_sizes.reverse()
block_coords = [self.get_block_coords(octant_size=octant_size, include_ends=True) for octant_size in octant_sizes]
block_sizes = [block[1, 1, 1, :] - block[0, 0, 0, :] for block in block_coords]
block_coords = [block[:-1, :-1, :-1, :] for block in block_coords]
# create sampling grids for training
num_samples = self.min_octant_size
posts = torch.linspace(-1, 1, self.min_octant_size+1)[:-1]
row_coords, col_coords, dep_coords = torch.meshgrid(posts, posts, posts)
row_coords = row_coords.flatten()
col_coords = col_coords.flatten()
dep_coords = dep_coords.flatten()
# create sampling grids for evaluation
# here we need to sample every voxel within each block
posts = [torch.linspace(-1, 1, voxel_size+1)[:-1] + (1/voxel_size)/2 for voxel_size in octant_sizes]
eval_coords = [torch.meshgrid(post, post, post) for post in posts]
eval_row_coords = [eval_coord[0].flatten() for eval_coord in eval_coords]
eval_col_coords = [eval_coord[1].flatten() for eval_coord in eval_coords]
eval_dep_coords = [eval_coord[2].flatten() for eval_coord in eval_coords]
def _populate_tree_helper(octant, idx):
# get block scale idx
scale_idx = len(idx) - (self.min_octtree_level)
# do we have octants at this level?
if scale_idx >= 0:
# set patch parameters
coords = block_coords[scale_idx]
coord_idx = _index_block_coord(idx, coords.shape[0], coord=[0, 0, 0])
octant.block_coord = coords[coord_idx[0], coord_idx[1], coord_idx[2]]
octant.block_size = block_sizes[scale_idx]
octant.scale = scale_idx
octant.voxel_size = octant_sizes[scale_idx]
octant.num_samples = num_samples
octant.row_coords = row_coords
octant.col_coords = col_coords
octant.dep_coords = dep_coords
octant.eval_row_coords = eval_row_coords[scale_idx]
octant.eval_col_coords = eval_col_coords[scale_idx]
octant.eval_dep_coords = eval_dep_coords[scale_idx]
octant.surface_tree = self.surface_tree
if not octant.children:
return
# recurse
for i in range(8):
child = octant.children[i]
_populate_tree_helper(child, [*idx, i])
return
# given list of tree idxs in {0,1,2,3}^N, retrieve the block coordinate
def _index_block_coord(tree_idx, length, coord=[0, 0, 0]):
if length == 1:
return coord
# depth 0
if tree_idx[0] == 0:
pass
elif tree_idx[0] == 1:
coord[1] += length//2
elif tree_idx[0] == 2:
coord[0] += length//2
elif tree_idx[0] == 3:
coord[0] += length//2
coord[1] += length//2
# depth 1
elif tree_idx[0] == 4:
coord[2] += length//2
elif tree_idx[0] == 5:
coord[1] += length//2
coord[2] += length//2
elif tree_idx[0] == 6:
coord[0] += length//2
coord[2] += length//2
elif tree_idx[0] == 7:
coord[0] += length//2
coord[1] += length//2
coord[2] += length//2
else:
raise ValueError("Unexpected child value, should be in{0,...7}")
return _index_block_coord(tree_idx[1:], length//2, coord)
# done with setup, now actually populate the tree
_populate_tree_helper(self.root, [])
def init_root(self, max_level):
def _init_root_helper(curr_octant, curr_level, max_level, optim_model):
if curr_level == max_level:
return
curr_octant.children = [Octant(optim_model) for _ in range(8)]
for octant in curr_octant.children:
octant.parent = curr_octant
_init_root_helper(octant, curr_level+1, max_level, optim_model)
return
# create root node
root = Octant(self.optim_model)
_init_root_helper(root, 0, max_level, self.optim_model)
return root
def get_block_coords(self, flatten=False, include_ends=False, octant_size=None):
# use finest scale patch by default
if octant_size is None:
octant_size = self.min_octant_size # TODO: ?? verify
# get size of each block
z_len = self.z_max - self.z_min
y_len = self.y_max - self.y_min
x_len = self.x_max - self.x_min
block_size = (z_len / (self.sidelength[0]) * octant_size,
y_len / (self.sidelength[1]) * octant_size,
x_len / (self.sidelength[2]) * octant_size)
# get block begin/end coordinates
if include_ends:
block_coords_z = torch.arange(self.z_min, self.z_max + block_size[0], block_size[0])
block_coords_y = torch.arange(self.y_min, self.y_max + block_size[1], block_size[1])
block_coords_x = torch.arange(self.x_min, self.x_max + block_size[2], block_size[2])
else:
block_coords_z = torch.arange(self.z_min, self.z_max, block_size[0])
block_coords_y = torch.arange(self.y_min, self.y_max, block_size[1])
block_coords_x = torch.arange(self.x_min, self.x_max, block_size[2])
# repeat for every single block
block_coords = torch.meshgrid(block_coords_z, block_coords_y, block_coords_x)
block_coords = torch.stack((block_coords[0], block_coords[1], block_coords[2]), dim=-1)
if flatten:
block_coords = block_coords.reshape(-1, 3)
return block_coords
def get_octants_at_level(self, level):
# level is the image scale: 0-> coarsest, N->finest
# what quadtree level do our octants start at?
# check input, too
target_level = level + self.min_octtree_level
assert level <= (self.max_octtree_level - self.min_octtree_level), \
"invalid 'level' input to get_blocks_at_level"
def _get_octants_at_level_helper(curr_octant, curr_level, octants):
if curr_level > target_level:
return
for octant in curr_octant.children:
_get_octants_at_level_helper(octant, curr_level+1, octants)
if curr_level == target_level:
octants.append(curr_octant)
return octants
return _get_octants_at_level_helper(self.root, 0, [])
def get_frozen_octants(self):
def _get_frozen_octants_helper(curr_octant, octants):
if curr_octant.frozen and curr_octant.active:
octants.append(curr_octant)
for octant in curr_octant.children:
_get_frozen_octants_helper(octant, octants)
return octants
return _get_frozen_octants_helper(self.root, [])
def get_active_octants(self, include_frozen_octants=False):
def _get_active_octants_helper(curr_octant, octants):
if curr_octant.active and \
(include_frozen_octants or
(not include_frozen_octants and not curr_octant.frozen)):
octants.append(curr_octant)
return octants
for octant in curr_octant.children:
_get_active_octants_helper(octant, octants)
return octants
return _get_active_octants_helper(self.root, [])
def activate_random(self):
def _activate_random_helper(curr_octant):
if not curr_octant.children:
curr_octant.activate()
return
elif (curr_octant.scale is not None) and (torch.rand(1).item() < 0.2):
curr_octant.activate()
return
for patch in curr_octant.children:
_activate_random_helper(patch)
return
_activate_random_helper(self.root)
def synchronize(self, master):
# set active/inactive nodes to be the same as master
# for now just toggle the flags without worrying about the gurobi variables
def _synchronize_helper(curr_octant, curr_octant_master):
curr_octant.active = curr_octant_master.active
curr_octant.frozen = curr_octant_master.frozen
curr_octant.value = curr_octant_master.value
if not curr_octant.children:
return
for octant, octant_master in zip(curr_octant.children, curr_octant_master.children):
_synchronize_helper(octant, octant_master)
return
_synchronize_helper(self.root, master.root)
def get_frozen_samples(self, oversample):
octants = self.get_frozen_octants()
if not octants:
return None, None, None, None
rel_coords, abs_coords, vals = [], [], []
octant_idx = []
for idx, p in enumerate(octants):
rel_samp, abs_samp, _ = p.get_stratified_samples(jitter=False, eval=True, oversample=oversample)
rel_samp = rel_samp.reshape(-1, self.min_octant_size**3, 3)
abs_samp = abs_samp.reshape(-1, self.min_octant_size**3, 3)
octant_idx.extend(rel_samp.shape[0] * [idx, ])
rel_coords.append(rel_samp)
abs_coords.append(abs_samp)
# values have the same size as rel_samp but last dim is a scalar
vals.append(p.value*torch.ones(abs_samp.shape[:-1] + (1,)))
return torch.cat(rel_coords, dim=0), torch.cat(abs_coords, dim=0), \
torch.cat(vals, dim=0), octant_idx
def get_stratified_samples(self, jitter=True, eval=False, oversample=1):
octants = self.get_active_octants()
rel_coords, abs_coords = [], []
all_global_indices = []
octant_idx = []
for idx, p in enumerate(octants):
rel_samp, abs_samp, global_indices = p.get_stratified_samples(jitter=jitter, eval=eval, oversample=oversample)
# always batch the coordinates in groups of a specific patch size
# so we can process them in parallel
rel_samp = rel_samp.reshape(-1, int(self.min_octant_size * oversample)**3, 3)
abs_samp = abs_samp.reshape(-1, int(self.min_octant_size * oversample)**3, 3)
if global_indices is not None:
global_indices = global_indices.reshape(-1, int(self.min_octant_size * oversample)**3, 1)
# since patch samples could be split across batches,
# keep track of which batch idx maps to which patch idx
octant_idx.extend(rel_samp.shape[0] * [idx, ])
rel_coords.append(rel_samp)
abs_coords.append(abs_samp)
if global_indices is not None:
all_global_indices.append(global_indices)
return torch.cat(rel_coords, dim=0), torch.cat(abs_coords, dim=0), octant_idx, None
def solve_optim(self, Max_Num_Octants=150):
octants = self.get_active_octants()
assert (len(octants) <= Max_Num_Octants), \
"You are trying to solve a model which is infeasible: " \
"Number of active octants > Max number of octants"
if self.c_max_octants is not None:
self.optim_model.remove(self.c_max_octants)
# global "knapsack" constraint
expr_sum_octants = [p.update_merge() for p in octants]
self.c_max_octants = self.optim_model.addConstr(gp.quicksum(expr_sum_octants) <= Max_Num_Octants)
# objective
self.optim_model.setObjective(gp.quicksum([p.get_cost() for p in octants]), GRB.MINIMIZE)
self.optim_model.optimize()
obj_val = self.optim_model.objVal
if self.optim_model.Status == GRB.INFEASIBLE:
print("----------- Model is infeasible")
self.optim_model.computeIIS()
self.optim_model.write("model.ilp")
# split and merge
merged = 0
split = 0
none = 0
for p in octants:
# print(p)
if p.has_split() and p.scale < self.max_octtree_level:
p.deactivate()
for child in p.get_children():
child.activate()
split += 1
elif p.has_merged() and p.scale >= self.min_octtree_level and p.scale > 0:
# we first check if it is active,
# since we could have already been activated by a neighbor
if p.active:
for neighbor in p.get_neighbors():
neighbor.deactivate()
p.parent.activate()
merged += 1
else:
p.update()
none += 1
stats_dict = {'merged': merged,
'splits': split,
'none': none,
'obj': obj_val}
print(f"============================= Total octants:{len(octants)}, split/merge:{split}/{merged}")
print(f"Vars={len(self.optim_model.getVars())}, Cstrs={len(self.optim_model.getConstrs())}")
return stats_dict
def draw(self, color_by_scale=False, save_fig=True):
fig = plt.figure(figsize=(5, 5))
ax = fig.gca(projection='3d')
depth = 1 + self.max_octtree_level - self.min_octtree_level
sidelen = 4**(depth-1) // 2**(depth-1)
# calculate scale
octant_list = self.get_active_octants()
octants_err = [p.err for p in octant_list]
max_err = np.max(octants_err)
min_err = np.min(octants_err)
cmap = plt.cm.get_cmap('viridis')
def cuboid_data(pos, size=(1, 1, 1)):
eps = 0.2
o = pos + (eps, eps, eps)
# get the length, width, and height
l, w, h = size
l -= eps
w -= eps
h -= eps
x = [[o[0], o[0] + l, o[0] + l, o[0], o[0]],
[o[0], o[0] + l, o[0] + l, o[0], o[0]],
[o[0], o[0] + l, o[0] + l, o[0], o[0]],
[o[0], o[0] + l, o[0] + l, o[0], o[0]]]
y = [[o[1], o[1], o[1] + w, o[1] + w, o[1]],
[o[1], o[1], o[1] + w, o[1] + w, o[1]],
[o[1], o[1], o[1], o[1], o[1]],
[o[1] + w, o[1] + w, o[1] + w, o[1] + w, o[1] + w]]
z = [[o[2], o[2], o[2], o[2], o[2]],
[o[2] + h, o[2] + h, o[2] + h, o[2] + h, o[2] + h],
[o[2], o[2], o[2] + h, o[2] + h, o[2]],
[o[2], o[2], o[2] + h, o[2] + h, o[2]]]
return np.array(x), np.array(y), np.array(z)
def draw_cube_at(pos=(0, 0, 0), size=(1, 1, 1),
color='b', edgecolor='b', alpha=1., ax=None):
# Plotting a cube element at position pos
if ax is not None:
Z, Y, X = cuboid_data(pos, size)
ax.plot_surface(X, Y, Z, color=color, rstride=1, cstride=1, alpha=alpha,
edgecolors=edgecolor, linewidth=0.1)
def _draw_level_2(octant, curr_level,
ax, sidelen, offset, scale):
if curr_level > self.max_octtree_level:
return ax
scale = scale/2.
for i, child in enumerate(octant.children):
# depth 0
if i == 0:
new_offset = (offset[0],
offset[1],
offset[2])
elif i == 1:
new_offset = (offset[0] + scale * sidelen,
offset[1],
offset[2])
elif i == 2:
new_offset = (offset[0],
offset[1] + scale * sidelen,
offset[2])
elif i == 3:
new_offset = (offset[0] + scale * sidelen,
offset[1] + scale * sidelen,
offset[2])
# depth 1
elif i == 4:
new_offset = (offset[0],
offset[1],
offset[2] + scale * sidelen)
elif i == 5:
new_offset = (offset[0] + scale * sidelen,
offset[1],
offset[2] + scale * sidelen)
elif i == 6:
new_offset = (offset[0],
offset[1] + scale * sidelen,
offset[2] + scale * sidelen)
else:
new_offset = (offset[0] + scale * sidelen,
offset[1] + scale * sidelen,
offset[2] + scale * sidelen)
if child.active:
norm_err = (child.err-min_err)/(max_err-min_err)
sz = scale*sidelen
if child.frozen:
color = [0.5, 0.5, 0.5]
edgecolor = [1.0, 0.0, 0.0, 0.1]
alpha = 0.
else:
color = cmap(norm_err)[0:3]
edgecolor = 'none'
alpha = 0.1
draw_cube_at(pos=new_offset, size=(sz, sz, sz),
color=color, edgecolor=edgecolor,
alpha=alpha, ax=ax)
else:
ax = _draw_level_2(child, curr_level+1,
ax, sidelen, new_offset, scale)
return ax
ax = _draw_level_2(self.root, self.min_octtree_level,
ax, sidelen,
(0., 0., 0.), 1.)
ax = fig.gca(projection='3d')
ax.grid(False)
ax.set_xlim(-1, sidelen+1)
ax.set_ylim(-1, sidelen+1)
ax.set_zlim(-1, sidelen+1)
ax.set_xticks([-1, sidelen+1])
ax.set_xticklabels([-1, 1])
ax.set_yticks([-1, sidelen+1])
ax.set_yticklabels([-1, 1])
ax.set_zticks([-1, sidelen+1])
ax.set_zticklabels([-1, 1])
return fig
class Octant():
def __init__(self, optim_model=None, block_coord=None, scale=None, gamma=0.95):
self.active = False
self.parent = None
self.children = []
# absolute block coordinate
self.block_coord = block_coord
# size of block in absolute coord frame
self.block_size = None
self.old_block_coord = None
self.old_block_size = None
# scale level of block
self.scale = scale
# num samples to be generated for this block
self.num_samples = None
# num pixels in this patch
self.voxel_size = None
# optimization model
self.optim = optim_model
# row/column coords for sampling at test time
# initialized by set_samples() function
self.row_coords = None
self.col_coords = None
self.dep_coords = None
self.near_mesh_abs_samples = None
self.near_mesh_rel_samples = None
# error for doing nothing, merging, splitting
self.err = 0.
self.last_updated = 0.
self.gamma = gamma
self._nocopy = ['optim', 'I_grp', 'I_split', 'I_none',
'I_merge', 'c_joinable', 'c_merge_split',
'children', 'parent', 'loss', 'loss_iter',
'err']
self.spec_cstrs = []
self._pickle_vars = ['parent', 'children', 'active', 'err', 'last_updated', 'frozen', 'value']
# options for pruning
self.frozen = False
self.value = 0
def __str__(self):
str = f"Octant id={id(self)}\n" \
f" . active={self.active}\n" \
f" . level={self.scale}\n" \
f" . model={self.optim}"
if self.active:
str += f"\n . g={self.I_grp.x}, s={self.I_split.x}, n={self.I_none.x}"
return str
# override deep copy to copy undeepcopyable objects by reference
def __deepcopy__(self, memo):
deep_copied_obj = Octant()
for k, v in self.__dict__.items():
if k in self._nocopy:
setattr(deep_copied_obj, k, v)
else:
setattr(deep_copied_obj, k, copy.deepcopy(v, memo))
return deep_copied_obj
def __getstate__(self):
state = self.__dict__.copy()
for k, v in self.__dict__.items():
if k not in self._pickle_vars:
del(state[k])
return state
def __load__(self, obj):
for k, v in obj.__dict__.items():
if k in ['children', 'parent']:
continue
setattr(self, k, v)
if self.active:
self.activate()
def update(self):
self.deactivate()
self.activate()
def activate(self):
self.active = True
# indicator variables
self.I_grp = self.optim.addVar(vtype=GRB.BINARY)
self.I_split = self.optim.addVar(vtype=GRB.BINARY)
self.I_none = self.optim.addVar(vtype=GRB.BINARY)
self.I_merge = gp.LinExpr(0.0)
# local constraint "merge/none/split"
self.c_joinable = self.optim.addConstr(self.I_grp + self.I_none + self.I_split == 1)
# local constraint "merge-split"
self.c_merge_split = None
def deactivate(self):
self.active = False
self.optim.remove(self.I_grp)
self.optim.remove(self.I_split)
self.optim.remove(self.I_none)
self.I_merge = gp.LinExpr(0.0)
self.optim.remove(self.c_joinable)
if self.c_merge_split is not None:
self.optim.remove(self.c_merge_split)
for cstr in self.spec_cstrs:
self.optim.remove(cstr)
self.spec_cstrs = []
def is_mergeable(self):
siblings = self.parent.children
return np.all(np.all([sib.active for sib in siblings]))
def set_sample_params(self, num_samples):
self.num_samples = num_samples
posts = torch.linspace(-1, 1, self.num_samples+1)[:-1]
row_coords, col_coords, dep_coords = torch.meshgrid(posts, posts, posts)
self.row_coords = row_coords.flatten()
self.col_coords = col_coords.flatten()
self.dep_coords = dep_coords.flatten()
def must_split(self):
self.spec_cstrs.append(
self.optim.addConstr(self.I_split == 1)
)
def must_merge(self):
self.spec_cstrs.append(
self.optim.addConstr(self.I_grp == 1)
)
def has_split(self):
return self.I_split.x == 1
def has_merged(self):
return self.I_grp.x == 1
# return self.I_none.x==0 and self.I_split.x==0
def has_done_nothing(self):
return self.I_none.x == 1
def get_cost(self):
area = self.block_size[0]**2
alpha = 0.2 # how much worse we expect the error to be when merging
beta = -0.02 # how much better we expect the error to be when splitting
# == Merge
if self.scale > 0: # it should never be root, but still..
err_merge = (8+alpha) * area * self.err
if self.parent.last_updated:
parent_area = self.parent.block_size[0]**2
err_merge = parent_area * self.parent.err # can multiply by 1/8 as in paper to make merging more aggressive
else:
err_merge = self.err
# == Split
if self.children:
err_split = (0.125+beta) * area * self.err
if self.children[0].last_updated:
err_children = np.sum([child.err for child in self.children])
err_split = area * err_children
else:
err_split = 1. # in case you don't have children, high to avoid splitting
err_none = area * self.err
return err_none * self.I_none \
+ err_split * self.I_split \
+ err_merge * self.I_grp
def update_merge(self):
if self.parent is None: # if root
return gp.LinExpr(0)
siblings = self.parent.children
if np.all([sib.active for sib in siblings]):
I_grp_neighs = [s.I_grp for s in siblings]
self.I_merge = gp.quicksum(I_grp_neighs)
# local constraint "joinable"
self.c_merge_split = self.optim.addConstr(self.I_none + self.I_split + .125*self.I_merge == 1)
expr_max_patches = 8 * self.I_split + 1 * self.I_none + .125 * self.I_grp
return expr_max_patches
def get_neighbors(self):
return self.parent.children
def get_children(self):
return self.children
def get_parent(self):
return self.parent
def is_joinable(self):
# test if siblings are all leaf nodes
siblings = self.parent.children
return np.all([sib.active for sib in siblings])
def get_block_coord(self):
return self.block_coord
def get_scale(self):
return self.scale
def update_error(self, error, iter):
self.err = self.gamma*self.err + (1-self.gamma)*error
self.last_updated = iter
def get_block_coords(self, flatten=False, include_ends=False, octant_size=None):
# get size of each block
z_len = 2
y_len = 2
x_len = 2
sidelength = 256
block_size = (z_len / (sidelength) * octant_size,
y_len / (sidelength) * octant_size,
x_len / (sidelength) * octant_size)
# get block begin/end coordinates
if include_ends:
block_coords_z = torch.arange(-1, -1 + block_size[0], block_size[0])
block_coords_y = torch.arange(-1, 1 + block_size[1], block_size[1])
block_coords_x = torch.arange(-1, 1 + block_size[2], block_size[2])
else:
block_coords_z = torch.arange(-1, 1, block_size[0])
block_coords_y = torch.arange(-1, 1, block_size[1])
block_coords_x = torch.arange(-1, 1, block_size[2])
# repeat for every single block
block_coords = torch.meshgrid(block_coords_z, block_coords_y, block_coords_x)
block_coords = torch.stack((block_coords[0], block_coords[1], block_coords[2]), dim=-1)
if flatten:
block_coords = block_coords.reshape(-1, 3)
return block_coords
def get_stratified_samples(self, jitter=True, eval=False, oversample=1., kd_tree=None):
# Block coords are always aligned to the pixel grid,
# e.g., they align with pixels 0, 8, 16, 24, etc. for
# patch size 8
#
# To normalize the coordinates between (-1, 1), consider
# we have an image of 64x64 and patch size 8x8.
# The block coordinate (-1, -1) aligns with pixel (0, 0)
# and coordinate (1, 1) aligns with pixel (63, 63)
#
# Absolute coordinates within a block should stretch all the way
# from the absolute position of one block coordinate to another.
# Say each block contains 8x8 pixels and we use a feature grid
# of 8x8 features to interpolate values within a block.
# This means is that the feature positions are not actually
# aligned to the pixel positions. The features are positioned
# on a grid stretching from one block coord to another whereas
# the pixel grid ends just short of the next block coordinate
#
# Example patch (x = pixel position, B = block coordinate position)
# and relative coordinate positions.
#
# -1 ^ B x x x x x x x B
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# | x x x x x x x x x
# 1 v B x x x x x x x B
# <--------------->
# -1 1
#
# When we generate samples for a patch, we sample an
# 8x8 grid that extends between block coords, i.e.
# between the arrows above
#
if eval:
if True or oversample != 1.:
post = torch.linspace(-1, 1, int(self.voxel_size * oversample + 1))[:-1]
post += (post[1] - post[0])/2
eval_coords = torch.meshgrid(post, post, post)
row_coords = eval_coords[0].flatten()
col_coords = eval_coords[1].flatten()
dep_coords = eval_coords[2].flatten()
else:
row_coords = self.eval_row_coords.flatten()
col_coords = self.eval_col_coords.flatten()
dep_coords = self.eval_dep_coords.flatten()
rel_samples = torch.stack((row_coords, col_coords, dep_coords), dim=-1)
abs_samples = self.block_coord[None, :] + (self.block_size[None, :]) * (rel_samples+1)/2
return rel_samples, abs_samples, None
else:
if (self.old_block_coord is None or self.old_block_size is None or
(self.old_block_coord != self.block_coord).any()
or (self.old_block_size != self.block_size).any()):
self.old_block_coord = self.block_coord
self.old_block_size = self.block_size
self.update_surface_coords()
return self.select_near_mesh(self.num_samples**3)
def update_surface_coords(self):
center = (self.block_coord + 0.5 * self.block_size).cpu().numpy()
side_length = float(self.block_size[0])
search_radius = (side_length/2)*(3**0.5)
indices = np.array(self.surface_tree.query_ball_point(center, search_radius))
if indices.shape[0] == 0:
self.near_mesh_abs_samples = np.zeros((0, 3))
self.near_mesh_rel_samples = np.zeros((0, 3))
else:
coordinates = self.surface_tree.data[indices]
in_cube = np.linalg.norm(coordinates - center, ord=np.inf, axis=1) < (side_length/2)
self.near_mesh_abs_samples = torch.FloatTensor(coordinates[in_cube])
self.near_mesh_rel_samples = 2 * (self.near_mesh_abs_samples - self.block_coord[None, :])/self.block_size[None, :] - 1
def select_near_mesh(self, num_samples, jitter=True):
if np.random.rand() > 0.9 and self.near_mesh_abs_samples.shape[0] > 0: # Near surface
selection_indices = torch.randint(self.near_mesh_abs_samples.shape[0], (num_samples,))
rel_samples_mesh = self.near_mesh_rel_samples[selection_indices]
abs_samples_mesh = self.near_mesh_abs_samples[selection_indices]
rel_samples_mesh += torch.randn_like(rel_samples_mesh) * 0.05 * self.block_size
abs_samples_mesh = self.block_coord[None, :] + (self.block_size[None, :]) * (rel_samples_mesh+1)/2
return rel_samples_mesh, abs_samples_mesh, None
else: # Uniform
row_coords = self.row_coords + torch.rand_like(self.row_coords) * 2./self.num_samples
col_coords = self.col_coords + torch.rand_like(self.col_coords) * 2./self.num_samples
dep_coords = self.dep_coords + torch.rand_like(self.dep_coords) * 2./self.num_samples
rel_samples = torch.stack((row_coords, col_coords, dep_coords), dim=-1)
abs_samples = self.block_coord[None, :] + (self.block_size[None, :]) * (rel_samples+1)/2
return rel_samples, abs_samples, None
此处可能存在不合适展示的内容，页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容，可点击提交进行申诉，我们将尽快为您处理。