def get_random_batch(self, batch_size):
data_ids_a = np.zeros(batch_size, 'int32')
data_ids_b = np.zeros(batch_size, 'int32')
viewpoint_ids_a = jt.zeros(batch_size)
viewpoint_ids_b = jt.zeros(batch_size)
for i in range(batch_size):
class_id = np.random.choice(self.class_ids)
object_id = np.random.randint(0, self.num_data[class_id])
viewpoint_id_a = np.random.randint(0, 24)
viewpoint_id_b = np.random.randint(0, 24)
data_id_a = (object_id + self.pos[class_id]) * 24 + viewpoint_id_a
data_id_b = (object_id + self.pos[class_id]) * 24 + viewpoint_id_b
data_ids_a[i] = data_id_a
data_ids_b[i] = data_id_b
viewpoint_ids_a[i] = viewpoint_id_a
viewpoint_ids_b[i] = viewpoint_id_b
images_a = jt.array((self.images[data_ids_a].astype('float32') / 255.))
images_b = jt.array((self.images[data_ids_b].astype('float32') / 255.))
distances = jt.ones(batch_size).float() * self.distance
elevations_a = jt.ones(batch_size).float() * self.elevation
elevations_b = jt.ones(batch_size).float() * self.elevation
viewpoints_a = jr.get_points_from_angles(distances, elevations_a,
-viewpoint_ids_a * 15)
viewpoints_b = jr.get_points_from_angles(distances, elevations_b,
-viewpoint_ids_b * 15)
return images_a, images_b, viewpoints_a, viewpoints_b
def create_mask_montage(self, image, predictions):
"""
Create a montage showing the probability heatmaps for each one one of the
detected objects
Arguments:
image (np.ndarray): an image as returned by OpenCV
predictions (BoxList): the result of the computation by the model.
It should contain the field `mask`.
"""
masks = predictions.get_field("mask")
masks_per_dim = self.masks_per_dim
masks = L.interpolate(
masks.float(), scale_factor=1 / masks_per_dim
).byte()
height, width = masks.shape[-2:]
max_masks = masks_per_dim ** 2
masks = masks[:max_masks]
# handle case where we have less detections than max_masks
if len(masks) < max_masks:
masks_padded = jt.zeros((max_masks, 1, height, width)).uint8()
masks_padded[: len(masks)] = masks
masks = masks_padded
masks = masks.reshape(masks_per_dim, masks_per_dim, height, width)
result = jt.zeros(
(masks_per_dim * height, masks_per_dim * width)
).uint8()
for y in range(masks_per_dim):
start_y = y * height
end_y = (y + 1) * height
for x in range(masks_per_dim):
start_x = x * width
end_x = (x + 1) * width
result[start_y:end_y, start_x:end_x] = masks[y, x]
return cv2.applyColorMap(result.numpy(), cv2.COLORMAP_JET)
def draw_lr_schedule():
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
param1 = jt.zeros([3, 3, 256, 256]).float32()
param2 = jt.zeros([3, 3, 256, 256]).float32()
base_lr = 1e-4
params = [{
'params': [param1],
'lr': base_lr * 1.
}, {
'params': [param2],
'lr': base_lr * 2.,
'weight_decay': 0.
}]
#params = [param1, param2]
optimizer = jt.optim.Adam(params, base_lr, weight_decay=0.0001)
iterations = range(100000)
lrs = []
for iteration in iterations:
lr = adjust_learning_rate(optimizer,
iteration,
BASE_LR=1e-4,
WARM_UP_FACTOR=1.0 / 3.0,
WARM_UP_ITERS=5000,
STEPS=[0, 60000, 80000],
GAMMA=0.1)
lrs.append(lr)
plt.figure()
plt.plot(iterations, lrs)
plt.savefig("lr_schedule.jpg")
def get_sample_region(self, gt, strides, num_points_per, gt_xs, gt_ys, radius=1):
num_gts = gt.shape[0]
K = len(gt_xs)
gt = gt[None].expand((K, num_gts, 4))
center_x = (gt[..., 0] + gt[..., 2]) / 2
center_y = (gt[..., 1] + gt[..., 3]) / 2
center_gt = jt.zeros(gt.shape,dtype=gt.dtype)
# no gt
if center_x[..., 0].sum() == 0:
return jt.zeros(gt_xs.shape, dtype='uint8')
beg = 0
for level, n_p in enumerate(num_points_per):
end = beg + n_p
stride = strides[level] * radius
xmin = center_x[beg:end] - stride
ymin = center_y[beg:end] - stride
xmax = center_x[beg:end] + stride
ymax = center_y[beg:end] + stride
# limit sample region in gt
center_gt[beg:end, :, 0] = jt.ternary(xmin > gt[beg:end, :, 0], xmin, gt[beg:end, :, 0])
center_gt[beg:end, :, 1] = jt.ternary(ymin > gt[beg:end, :, 1], ymin, gt[beg:end, :, 1])
center_gt[beg:end, :, 2] = jt.ternary(xmax > gt[beg:end, :, 2], gt[beg:end, :, 2], xmax)
center_gt[beg:end, :, 3] = jt.ternary(ymax > gt[beg:end, :, 3], gt[beg:end, :, 3], ymax)
beg = end
left = gt_xs[:, None] - center_gt[..., 0]
right = center_gt[..., 2] - gt_xs[:, None]
top = gt_ys[:, None] - center_gt[..., 1]
bottom = center_gt[..., 3] - gt_ys[:, None]
center_bbox = jt.stack((left, top, right, bottom), -1)
inside_gt_bbox_mask = center_bbox.min(-1) > 0
return inside_gt_bbox_mask
def grad(self, grad_soft_colors):
face_vertices, textures, soft_colors, faces_info, aggrs_info = self.save_vars
image_size = self.image_size
background_color = self.background_color
near = self.near
far = self.far
eps = self.eps
sigma_val = self.sigma_val
dist_eps = self.dist_eps
gamma_val = self.gamma_val
func_dist_type = self.func_dist_type
func_rgb_type = self.func_rgb_type
func_alpha_type = self.func_alpha_type
texture_type = self.texture_type
fill_back = self.fill_back
grad_faces = jt.zeros((face_vertices.shape))
grad_textures = jt.zeros((textures.shape))
grad_faces, grad_textures = \
soft_rasterize_cuda.backward_soft_rasterize(face_vertices, textures, soft_colors,
faces_info, aggrs_info,
grad_faces, grad_textures, grad_soft_colors,
image_size, near, far, eps,
sigma_val, func_dist_type, dist_eps,
gamma_val, func_rgb_type, func_alpha_type,
texture_type, int(fill_back))
# print(grad_faces, grad_textures)
return grad_faces, grad_textures
def __init__(self, num_features):
super().__init__()
self.num_features = num_features
self.weight = jt.zeros((num_features))
self.bias = jt.zeros((num_features))
init.uniform_(self.weight, 0, 1)
init.constant_(self.bias, 0.0)
def add_param_group(self, group):
values = group["values"] = []
m = group["m"] = []
for p in group["params"]:
values.append(jt.zeros(p.shape, p.dtype).stop_grad())
m.append(jt.zeros(p.shape, p.dtype).stop_grad())
self.param_groups.append(group)
def __init__(self,
parameters,
lr,
eps=1e-8,
betas=(0.9, 0.999),
weight_decay=0,
param_sync_iter=10000):
self.lr = lr
self.eps = eps
self.betas = betas
# self.weight_decay = weight_decay
assert weight_decay == 0, "weight_decay is not supported yet"
self.adam_step = 0
self.param_sync_iter = param_sync_iter
self.no_grad_parameters = []
self.parameters = []
self.values = []
self.m = []
for p in parameters:
if jt.mpi:
p.assign(p.mpi_broadcast().detach())
if p.is_stop_grad():
self.no_grad_parameters.append(p)
continue
self.parameters.append(p)
self.values.append(
jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad())
self.m.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad())
def __init__(self,
img_size=224,
patch_size=16,
in_chans=3,
num_classes=1000,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4.,
qkv_bias=False,
qk_scale=None,
drop_rate=0.,
attn_drop_rate=0.,
drop_path_rate=0.,
hybrid_backbone=None,
norm_layer=nn.LayerNorm):
super(VisionTransformer, self).__init__()
if hybrid_backbone is not None:
self.patch_embed = HybridEmbed(hybrid_backbone,
img_size=img_size,
in_chans=in_chans,
embed_dim=embed_dim)
else:
self.patch_embed = PatchEmbed(img_size=img_size,
patch_size=patch_size,
in_chans=in_chans,
embed_dim=embed_dim)
num_patches = self.patch_embed.num_patches
self.cls_token = jt.zeros((1, 1, embed_dim))
self.pos_embed = jt.zeros((1, num_patches + 1, embed_dim))
self.pos_drop = nn.Dropout(drop_rate)
dpr = [x.item() for x in np.linspace(0, drop_path_rate, depth)
] # stochastic depth decay rule
self.blocks = nn.ModuleList([
Block(dim=embed_dim,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop=drop_rate,
attn_drop=attn_drop_rate,
drop_path=dpr[i],
norm_layer=norm_layer) for i in range(depth)
])
self.norm = norm_layer(embed_dim)
# NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here
#self.repr = nn.Linear(embed_dim, representation_size)
#self.repr_act = nn.Tanh()
# Classifier head
self.head = nn.Linear(embed_dim, num_classes)
self.pos_embed = trunc_normal(self.pos_embed, std=.02)
self.cls_token = trunc_normal(self.cls_token, std=.02)
self.apply(self._init_weights)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--filename-input', type=str,
default=os.path.join(data_dir, 'obj/spot/spot_triangulated.obj'))
parser.add_argument('-o', '--output-dir', type=str,
default=os.path.join(data_dir, 'results/output_render'))
args = parser.parse_args()
# other settings
camera_distance = 2.732
elevation = 30
azimuth = 0
# load from Wavefront .obj file
mesh = jr.Mesh.from_obj(args.filename_input, load_texture=True, texture_res=5 ,texture_type='surface', dr_type='softras')
# create renderer with SoftRas
renderer = jr.Renderer(dr_type='softras')
os.makedirs(args.output_dir, exist_ok=True)
#0.5 0.4
metallic_textures = jt.zeros((1, mesh.faces.shape[1], 5 * 5, 1)).float32() + 0.5
roughness_textures = jt.zeros((1, mesh.faces.shape[1], 5 * 5, 1)).float32() + 0.4
# draw object from different view
loop = tqdm.tqdm(list(range(0, 360, 4)))
writer = imageio.get_writer(os.path.join(args.output_dir, 'rotation.gif'), mode='I')
imgs = []
from PIL import Image
for num, azimuth in enumerate(loop):
# rest mesh to initial state
mesh.reset_()
loop.set_description('Drawing rotation')
renderer.transform.set_eyes_from_angles(camera_distance, elevation, azimuth)
rgb = renderer(mesh.vertices, mesh.faces, textures=mesh.textures, metallic_textures=metallic_textures, roughness_textures=roughness_textures)
image = rgb.numpy()[0].transpose((1, 2, 0))
writer.append_data((255*image).astype(np.uint8))
writer.close()
# draw object from different sigma and gamma
loop = tqdm.tqdm(list(np.arange(-4, -2, 0.2)))
renderer.transform.set_eyes_from_angles(camera_distance, elevation, 45)
writer = imageio.get_writer(os.path.join(args.output_dir, 'bluring.gif'), mode='I')
for num, gamma_pow in enumerate(loop):
# rest mesh to initial state
mesh.reset_()
renderer.set_gamma(10**gamma_pow)
renderer.set_sigma(10**(gamma_pow - 1))
loop.set_description('Drawing blurring')
images = renderer(mesh.vertices, mesh.faces, textures=mesh.textures, metallic_textures=metallic_textures, roughness_textures=roughness_textures)
image = images.numpy()[0].transpose((1, 2, 0)) # [image_size, image_size, RGB]
writer.append_data((255*image).astype(np.uint8))
writer.close()
# save to textured obj
mesh.reset_()
mesh.save_obj(os.path.join(args.output_dir, 'saved_spot.obj'))
def execute(self, face_vertices, textures):
# face_vertices: [nb, nf, 9]
# textures: [nb, nf, 9]
func_dist_map = {'hard': 0, 'barycentric': 1, 'euclidean': 2}
func_rgb_map = {'hard': 0, 'softmax': 1}
func_alpha_map = {'hard': 0, 'sum': 1, 'prod': 2}
func_map_sample = {'surface': 0, 'vertex': 1}
image_size = self.image_size
background_color = self.background_color
near = self.near
far = self.far
eps = self.eps
sigma_val = self.sigma_val
gamma_val = self.gamma_val
dist_eps = self.dist_eps
fill_back = self.fill_back
dist_func = self.dist_func
aggr_func_rgb = self.aggr_func_rgb
aggr_func_alpha = self.aggr_func_alpha
texture_type = self.texture_type
self.func_dist_type = func_dist_map[dist_func]
self.func_rgb_type = func_rgb_map[aggr_func_rgb]
self.func_alpha_type = func_alpha_map[aggr_func_alpha]
self.texture_type = func_map_sample[texture_type]
face_vertices = face_vertices.clone()
textures = textures.clone()
self.batch_size, self.num_faces = face_vertices.shape[:2]
faces_info = jt.zeros((self.batch_size, self.num_faces,
9 * 3)) # [inv*9, sym*9, obt*3, 0*6]
aggrs_info = jt.zeros(
(self.batch_size, 2, self.image_size, self.image_size))
soft_colors = jt.ones(
(self.batch_size, 4, self.image_size, self.image_size))
soft_colors[:, 0, :, :] *= background_color[0]
soft_colors[:, 1, :, :] *= background_color[1]
soft_colors[:, 2, :, :] *= background_color[2]
# TODO: soft_colors do not init in cuda
faces_info, aggrs_info, soft_colors = \
soft_rasterize_cuda.forward_soft_rasterize(face_vertices, textures,
faces_info, aggrs_info,
soft_colors,
image_size, near, far, eps,
sigma_val, self.func_dist_type, self.dist_eps,
gamma_val, self.func_rgb_type, self.func_alpha_type,
self.texture_type, int(fill_back))
self.save_vars = face_vertices, textures, soft_colors, faces_info, aggrs_info
return soft_colors
def execute(self, mesh, eyes=None):
if self.Gbuffer == "albedo":
return mesh
if self.Gbuffer == "normal" or self.Gbuffer == "depth":
mesh.textures = jt.ones_like(mesh.textures)
if self.light_mode == 'surface':
diffuseLight = jt.zeros(mesh.faces.shape)
specularLight = jt.zeros(mesh.faces.shape)
diffuseLight = self.ambient(diffuseLight)
for directional in self.directionals:
[diffuseLight, specularLight] = directional(
diffuseLight, specularLight, mesh.surface_normals,
(jt.sum(mesh.face_vertices, dim=2) / 3.0), eyes,
mesh.with_specular, mesh.metallic_textures,
mesh.roughness_textures)
if len(mesh.textures.shape) == 4:
mesh.textures = jt.clamp(
mesh.textures * diffuseLight.unsqueeze(2) +
jt.ones_like(mesh.textures) * specularLight.unsqueeze(2),
0.0, 1.0)
elif len(mesh.textures.shape) == 6:
mesh.textures = jt.clamp(
mesh.textures *
diffuseLight.unsqueeze(2).unsqueeze(2).unsqueeze(2) +
jt.ones_like(mesh.textures) *
specularLight.unsqueeze(2).unsqueeze(2).unsqueeze(2), 0.0,
1.0)
elif self.light_mode == 'vertex':
diffuseLight = jt.zeros(mesh.vertices.shape)
specularLight = jt.zeros(mesh.vertices.shape)
diffuseLight = self.ambient(diffuseLight)
for directional in self.directionals:
[diffuseLight, specularLight
] = directional(diffuseLight, specularLight,
mesh.vertex_normals, mesh.vertices, eyes,
mesh.with_specular, mesh.metallic_textures,
mesh.roughness_textures)
if len(mesh.textures.shape) == 4:
mesh.textures = jt.clamp(
mesh.textures * diffuseLight.unsqueeze(2) +
jt.ones_like(mesh.textures) * specularLight.unsqueeze(2),
0.0, 1.0)
elif len(mesh.textures.shape) == 6:
mesh.textures = jt.clamp(
mesh.textures *
diffuseLight.unsqueeze(2).unsqueeze(2).unsqueeze(2) +
jt.ones_like(mesh.textures) *
specularLight.unsqueeze(2).unsqueeze(2).unsqueeze(2), 0.0,
1.0)
return mesh
def __init__(self, params, lr, eps=1e-8, betas=(0.9, 0.999), weight_decay=0):
super().__init__(params, lr)
self.eps = eps
self.betas = betas
# self.weight_decay = weight_decay
assert weight_decay==0, "weight_decay is not supported yet"
# initialize required arguments for each param_groups
for pg in self.param_groups:
values = pg["values"] = []
m = pg["m"] = []
for p in pg["params"]:
values.append(jt.zeros(p.shape, p.dtype).stop_grad())
m.append(jt.zeros(p.shape, p.dtype).stop_grad())
def loc2bbox(src_bbox,loc):
if src_bbox.shape[0] == 0:
return jt.zeros((0, 4), dtype=loc.dtype)
src_width = src_bbox[:, 2:3] - src_bbox[:, 0:1]
src_height = src_bbox[:, 3:4] - src_bbox[:, 1:2]
src_center_x = src_bbox[:, 0:1] + 0.5 * src_width
src_center_y = src_bbox[:, 1:2] + 0.5 * src_height
dx = loc[:, 0:1]
dy = loc[:, 1:2]
dw = loc[:, 2:3]
dh = loc[:, 3:4]
center_x = dx*src_width+src_center_x
center_y = dy*src_height+src_center_y
w = jt.exp(dw.minimum(20.0)) * src_width
h = jt.exp(dh.minimum(20.0)) * src_height
x1,y1,x2,y2 = center_x-0.5*w, center_y-0.5*h, center_x+0.5*w, center_y+0.5*h
dst_bbox = jt.contrib.concat([x1,y1,x2,y2],dim=1)
return dst_bbox
def gcn_norm(edge_index,
edge_weight=None,
num_nodes=None,
improved=False,
add_self_loops=True,
dtype=None):
fill_value = 2. if improved else 1.
if isinstance(edge_index, Var):
num_nodes = maybe_num_nodes(edge_index, num_nodes)
if edge_weight is None:
edge_weight = jt.ones((edge_index.size(1), ))
if add_self_loops:
edge_index, tmp_edge_weight = add_remaining_self_loops(
edge_index, edge_weight, fill_value, num_nodes)
assert tmp_edge_weight is not None
edge_weight = tmp_edge_weight
row, col = edge_index[0], edge_index[1]
shape = list(edge_weight.shape)
shape[0] = num_nodes
deg = jt.zeros(shape)
deg = jt.scatter(deg, 0, col, src=edge_weight, reduce='add')
deg_inv_sqrt = deg.pow(-0.5)
deg_inv_sqrt.masked_fill(deg_inv_sqrt == float('inf'), 0)
return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
def test_setitem_(self):
arr0 = jt.random((4, 2, 2))
data0 = jt.ones((2, 2))
arr0[1] = data0
arr0.sync()
data0.data[0, 0] = 0
assert arr0[1, 0, 0] == 0
arr00 = jt.random((4, 2, 2))
data00 = jt.ones((2, 2))
# share memory will fail if d has an edge to other nodes.
tmp = data00 + 1
arr00[1] = data00
arr00.sync()
data00.data[0, 0] = 0
assert arr00[1, 0, 0] == 0
arr1 = jt.random((4, 2, 2))
data1 = jt.zeros((2, 2))
arr1[3, :, :] = data1
arr1.sync()
data1.data[0, 0] = 1
assert arr1[3, 0, 0] == 1
arr21 = jt.ones((2, 2))
arr22 = jt.ones((2, 2)) * 2
arr2 = jt.contrib.concat([arr21, arr22], dim=0)
arr2.sync()
arr21.data[0, 0] = 3
arr22.data[0, 0] = 4
assert arr2[0, 0] == 3
assert arr2[2, 0] == 4
def execute(self, pc1, pc2, reduction='mean', dims='BNC'):
assert dims in ['BNC', 'BCN']
if dims == 'BCN':
pc1, pc2 = pc1.permute(0, 2, 1), pc2.permute(0, 2, 1)
batch_size_1, N, _ = pc1.shape
batch_size_2, M, _ = pc2.shape
assert batch_size_1 == batch_size_2
batch_size = batch_size_1
temp = jt.zeros([batch_size, (N + M) * 2], pc1.dtype)
match = jt.code(
shape=[batch_size, M, N],
dtype=pc1.dtype,
inputs=[pc1, pc2, temp],
cuda_header=EMD_gpu_header,
cuda_src=approxmatch_gpu_src,
)
emd = jt.code(
shape=[batch_size],
dtype=pc1.dtype,
inputs=[pc1, pc2, match],
cuda_header=EMD_gpu_header,
cuda_src=matchcost_gpu_src,
)
self.saved_vars = (pc1, pc2, match, reduction)
if reduction is None:
return emd
elif reduction == 'sum':
return emd.sum()
elif reduction == 'mean':
return emd.mean()
def __call__(self, proposals, mask_logits, targets):
"""
Arguments:
proposals (list[BoxList])
mask_logits (Tensor)
targets (list[BoxList])
Return:
mask_loss (Tensor): scalar tensor containing the loss
"""
labels, mask_targets, mask_ratios = self.prepare_targets(
proposals, targets)
labels = cat(labels, dim=0)
mask_targets = cat(mask_targets, dim=0)
positive_inds = jt.nonzero(labels > 0).squeeze(1)
labels_pos = labels[positive_inds]
# accept empty tensors, so handle it separately
if mask_targets.numel() == 0:
if not self.maskiou_on:
return mask_logits.sum() * 0
else:
selected_index = jt.arange(mask_logits.shape[0])
selected_mask = mask_logits[selected_index, labels]
mask_num, mask_h, mask_w = selected_mask.shape
selected_mask = selected_mask.reshape(mask_num, 1, mask_h,
mask_w)
return mask_logits.sum() * 0, selected_mask, labels, None
if self.maskiou_on:
mask_ratios = cat(mask_ratios, dim=0)
value_eps = 1e-10 * jt.ones((mask_targets.shape[0], ))
mask_ratios = jt.maximum(mask_ratios, value_eps)
pred_masks = mask_logits[positive_inds, labels_pos]
pred_masks[:] = pred_masks > 0
mask_targets_full_area = mask_targets.sum(
dims=[1, 2]) / mask_ratios
mask_ovr = pred_masks * mask_targets
mask_ovr_area = mask_ovr.sum(dims=[1, 2])
mask_union_area = pred_masks.sum(
dims=[1, 2]) + mask_targets_full_area - mask_ovr_area
value_1 = jt.ones((pred_masks.shape[0], ))
value_0 = jt.zeros((pred_masks.shape[0], ))
mask_union_area = jt.maximum(mask_union_area, value_1)
mask_ovr_area = jt.maximum(mask_ovr_area, value_0)
maskiou_targets = mask_ovr_area / mask_union_area
binary_cross_entropy_with_logits = nn.BCEWithLogitsLoss()
mask_loss = binary_cross_entropy_with_logits(
mask_logits[positive_inds, labels_pos], mask_targets)
if not self.maskiou_on:
return mask_loss
else:
selected_index = jt.index((mask_logits.shape[0], ), dim=0)
selected_mask = mask_logits[selected_index, labels]
mask_num, mask_h, mask_w = selected_mask.shape
selected_mask = selected_mask.reshape(mask_num, 1, mask_h, mask_w)
selected_mask = selected_mask.sigmoid()
return mask_loss, selected_mask, labels, maskiou_targets
def __init__(self, template_path):
super(Model, self).__init__()
# set template mesh
self.template_mesh = jr.Mesh.from_obj(template_path, dr_type='n3mr')
self.vertices = (self.template_mesh.vertices * 0.5).stop_grad()
self.faces = self.template_mesh.faces.stop_grad()
self.textures = self.template_mesh.textures.stop_grad()
# optimize for displacement map and center
self.displace = jt.zeros(self.template_mesh.vertices.shape)
self.center = jt.zeros((1, 1, 3))
# define Laplacian and flatten geometry constraints
self.laplacian_loss = LaplacianLoss(self.vertices[0], self.faces[0])
self.flatten_loss = FlattenLoss(self.faces[0])
def execute(self, mesh):
if self.light_mode == 'surface':
light = jt.zeros(mesh.faces.shape)
light = self.ambient(light)
for directional in self.directionals:
light = directional(light, mesh.surface_normals)
mesh.textures = mesh.textures * light.unsqueeze(2)
elif self.light_mode == 'vertex':
light = jt.zeros(mesh.vertices.shape)
light = self.ambient(light)
for directional in self.directionals:
light = directional(light, mesh.vertex_normals)
mesh.textures = mesh.textures * light
return mesh
def test_log_capture(self):
LOG.log_capture_start()
with jt.var_scope(log_v=1000, log_vprefix=""):
LOG.v("1")
LOG.vv("2")
LOG.i("3")
LOG.w("4")
LOG.e("5")
a = jt.zeros([10])
a.sync()
LOG.log_capture_stop()
# TODO: why need manually delete this variable?
del a
logs = LOG.log_capture_read()
logs2 = LOG.log_capture_read()
assert len(logs2) == 0
for i in range(5):
assert logs[i]['msg'] == str(i + 1)
assert logs[i]['level'] == 'iiiwe'[i]
assert logs[i]['name'] == 'test_log.py'
finished_log = [
l["msg"] for l in logs
if l["name"] == "executor.cc" and "return vars:" in l["msg"]
]
assert len(finished_log) == 1 and "[10,]" in finished_log[0]
def create_texture_image(textures, texture_res=16):
num_faces = textures.shape[0]
tile_width = int((num_faces - 1.)**0.5) + 1
tile_height = int((num_faces - 1.) / tile_width) + 1
image = jt.ones((tile_height * texture_res, tile_width * texture_res, 3))
vertices = jt.zeros((num_faces, 3, 2)) # [:, :, UV]
face_nums = jt.array(np.arange(num_faces))
column = face_nums % tile_width
row = face_nums / tile_width
vertices[:, 0, 0] = column * texture_res + texture_res / 2
vertices[:, 0, 1] = row * texture_res + 1
vertices[:, 1, 0] = column * texture_res + 1
vertices[:, 1, 1] = (row + 1) * texture_res - 1 - 1
vertices[:, 2, 0] = (column + 1) * texture_res - 1 - 1
vertices[:, 2, 1] = (row + 1) * texture_res - 1 - 1
image = create_texture_image_cuda.create_texture_image(
vertices, textures, image, 1e-5)
vertices[:, :, 0] /= (image.shape[1] - 1)
vertices[:, :, 1] /= (image.shape[0] - 1)
image = image.numpy()
vertices = vertices.numpy()
image = image[::-1, ::1]
return image, vertices
def make_base_grid_5D(theta,N,C,D,H,W,align_corners):
base_grid = jt.zeros((N, D, H, W, 4), dtype=theta.dtype)
base_grid[...,0] = linspace_from_neg_one(theta, W, align_corners)
base_grid[...,1] = jt.unsqueeze(linspace_from_neg_one(theta, H, align_corners),-1)
base_grid[...,2] = jt.unsqueeze(jt.unsqueeze(linspace_from_neg_one(theta, D, align_corners),-1),-1)
base_grid[...,-1] = 1
return base_grid
def fuse_conv_and_bn(conv, bn):
# Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
fusedconv = nn.Conv2d(conv.in_channels,
conv.out_channels,
kernel_size=conv.kernel_size,
stride=conv.stride,
padding=conv.padding,
groups=conv.groups,
bias=True)
# prepare filters
w_conv = conv.weight.clone().view(conv.out_channels, -1)
w_bn = jt.diag(bn.weight / (jt.sqrt(bn.eps + bn.running_var)))
fusedconv.weight.assign(
jt.matmul(w_bn, w_conv).view(fusedconv.weight.shape))
# prepare spatial bias
b_conv = jt.zeros(
(conv.weight.shape[0], )) if conv.bias is None else conv.bias
b_bn = bn.bias - bn.weight * bn.running_mean / jt.sqrt(bn.running_var +
bn.eps)
fusedconv.bias.assign(
jt.matmul(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
return fusedconv
def __init__(self,
parameters,
lr,
momentum=0,
weight_decay=0,
dampening=0,
nesterov=False,
param_sync_iter=10000):
self.lr = lr
self.momentum = momentum
self.weight_decay = weight_decay
self.dampening = dampening
self.nesterov = nesterov
self.sgd_step = 0
self.param_sync_iter = param_sync_iter
self.no_grad_parameters = []
self.parameters = []
self.values = []
for p in parameters:
# broadcast parameter from 0 node when init
if jt.mpi:
p.assign(p.mpi_broadcast().detach())
if p.is_stop_grad():
self.no_grad_parameters.append(p)
continue
self.parameters.append(p)
self.values.append(
jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad())
def __init__(self, filename_obj, filename_ref):
super(Model, self).__init__()
# set template mesh
texture_size = 4
self.template_mesh = jr.Mesh.from_obj(filename_obj,
texture_res=texture_size,
load_texture=True,
dr_type='softras')
self.vertices = (self.template_mesh.vertices).stop_grad()
self.faces = self.template_mesh.faces.stop_grad()
self.textures = self.template_mesh.textures.stop_grad()
self.metallic_textures = jt.zeros(
(1, self.faces.shape[1], texture_size * texture_size,
1)).float32() + 0.4
self.metallic_textures = self.metallic_textures.stop_grad()
self.roughness_textures = jt.ones(
(1, self.faces.shape[1], texture_size * texture_size,
1)).float32()
# load reference image
self.image_ref = jt.array(
imread(filename_ref).astype('float32') / 255.).permute(
2, 0, 1).unsqueeze(0).stop_grad()
# setup renderer
self.renderer = jr.Renderer(dr_type='softras')
def __init__(self,
backbone,
img_size=224,
feature_size=None,
in_chans=3,
embed_dim=768):
super(HybridEmbed, self).__init__()
assert isinstance(backbone, nn.Module)
img_size = _pair(img_size)
self.img_size = img_size
self.backbone = backbone
if feature_size is None:
with jt.no_grad():
# FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
# map for all networks, the feature metadata has reliable channel and stride info, but using
# stride to calc feature dim requires info about padding of each stage that isn't captured.
training = backbone.is_training()
if training:
backbone.eval()
o = self.backbone(
jt.zeros((1, in_chans, img_size[0], img_size[1])))[-1]
feature_size = o.shape[-2:]
feature_dim = o.shape[1]
backbone.train()
else:
feature_size = _pair(feature_size)
feature_dim = self.backbone.feature_info.channels()[-1]
self.num_patches = feature_size[0] * feature_size[1]
self.proj = nn.Linear(feature_dim, embed_dim)
def __init__(self, C, kernel_size, stride, padding):
super(Shift, self).__init__()
self.C = C
kernel = jt.zeros((C, 1, kernel_size, kernel_size)).float32()
ch_idx = 0
assert stride in [1, 2]
self.stride = stride
self.padding = padding
self.kernel_size = kernel_size
self.dilation = 1
hks = kernel_size // 2
ksq = kernel_size ** 2
for i in range(kernel_size):
for j in range(kernel_size):
if i == hks and j == hks:
num_ch = C // ksq + C % ksq
else:
num_ch = C // ksq
kernel[ch_idx : ch_idx + num_ch, 0, i, j] = 1
ch_idx += num_ch
self.bias = None
self.kernel = kernel
self.kernel.stop_grad()
def dis_loss(self, real_samps, fake_samps, height, alpha):
# small assertion:
# assert real_samps.device == fake_samps.device, \
# "Real and Fake samples are not on the same device"
# device for computations:
# device = fake_samps.device
# predictions for real images and fake images separately :
r_preds = self.dis(real_samps, height, alpha)
f_preds = self.dis(fake_samps, height, alpha)
# calculate the real loss:
# real_loss = self.criterion(
# torch.squeeze(r_preds),
# torch.ones(real_samps.shape[0]).to(device))
real_loss = self.criterion(jt.squeeze(r_preds),
jt.ones(real_samps.shape[0]))
# calculate the fake loss:
# fake_loss = self.criterion(
# torch.squeeze(f_preds),
# torch.zeros(fake_samps.shape[0]).to(device))
fake_loss = self.criterion(jt.squeeze(f_preds),
jt.zeros(fake_samps.shape[0]))
# return final losses
return (real_loss + fake_loss) / 2
def select_over_all_levels(self, boxlists):
num_images = len(boxlists)
# different behavior during training and during testing:
# during training, post_nms_top_n is over *all* the proposals combined, while
# during testing, it is over the proposals for each image
# NOTE: it should be per image, and not per batch. However, to be consistent
# with Detectron, the default is per batch (see Issue #672)
if self.is_training() and self.fpn_post_nms_per_batch:
objectness = jt.contrib.concat(
[boxlist.get_field("objectness") for boxlist in boxlists], dim=0
)
box_sizes = [len(boxlist) for boxlist in boxlists]
post_nms_top_n = min(self.fpn_post_nms_top_n, len(objectness))
_, inds_sorted = jt.topk(objectness, post_nms_top_n, dim=0, sorted=True)
inds_mask = jt.zeros(objectness.shape).bool()
inds_mask[inds_sorted] = 1
inds_mask = inds_mask.split(box_sizes,dim=0)
for i in range(num_images):
boxlists[i] = boxlists[i][inds_mask[i]]
else:
for i in range(num_images):
objectness = boxlists[i].get_field("objectness")
post_nms_top_n = min(self.fpn_post_nms_top_n, objectness.shape[0])
_, inds_sorted = jt.topk(
objectness, post_nms_top_n, dim=0, sorted=True
)
boxlists[i] = boxlists[i][inds_sorted]
return boxlists
