def run(input, scalefactor): """ :param input: input mesh to reconstruct optimally. :return: final reconstruction after optimisation """ input, translation = center(input) if not global_variables.opt.HR: mesh_ref = global_variables.mesh_ref_LR else: mesh_ref = global_variables.mesh_ref ## Extract points and put them on GPU points = input.vertices random_sample = np.random.choice(np.shape(points)[0], size=10000) points = torch.from_numpy(points.astype( np.float32)).contiguous().unsqueeze(0) points = Variable(points) points = points.transpose(2, 1).contiguous() points = points.cuda() # Get a low resolution PC to find the best reconstruction after a rotation on the Y axis points_LR = torch.from_numpy(input.vertices[random_sample].astype( np.float32)).contiguous().unsqueeze(0) points_LR = Variable(points_LR) points_LR = points_LR.transpose(2, 1).contiguous() points_LR = points_LR.cuda() theta = 0 bestLoss = 10 # print("size: ", points_LR.size()) pointsReconstructed = global_variables.network(points_LR) dist1, dist2 = distChamfer( points_LR.transpose(2, 1).contiguous(), pointsReconstructed) loss_net = (torch.mean(dist1)) + (torch.mean(dist2)) # print("loss : ", loss_net.data[0], 0) # ---- Search best angle for best reconstruction on the Y axis--- for theta in np.linspace(-np.pi / 2, np.pi / 2, global_variables.opt.num_angles): if global_variables.opt.num_angles == 1: theta = 0 # Rotate mesh by theta and renormalise rot_matrix = np.array([[np.cos(theta), 0, np.sin(theta)], [0, 1, 0], [-np.sin(theta), 0, np.cos(theta)]]) rot_matrix = Variable(torch.from_numpy(rot_matrix).float()).cuda() points2 = torch.matmul(rot_matrix, points_LR) mesh_tmp = trimesh.Trimesh(process=False, use_embree=False, vertices=points2[0].transpose( 1, 0).data.cpu().numpy(), faces=global_variables.network.mesh.faces) #bbox bbox = np.array([[ np.max(mesh_tmp.vertices[:, 0]), np.max(mesh_tmp.vertices[:, 1]), np.max(mesh_tmp.vertices[:, 2]) ], [ np.min(mesh_tmp.vertices[:, 0]), np.min(mesh_tmp.vertices[:, 1]), np.min(mesh_tmp.vertices[:, 2]) ]]) norma = Variable( torch.from_numpy((bbox[0] + bbox[1]) / 2).float().cuda()) norma2 = norma.unsqueeze(1).expand(3, points2.size(2)).contiguous() points2[0] = points2[0] - norma2 mesh_tmp = trimesh.Trimesh(process=False, use_embree=False, vertices=points2[0].transpose( 1, 0).data.cpu().numpy(), faces=np.array([[0, 0, 0]])) # reconstruct rotated mesh pointsReconstructed = global_variables.network(points2) dist1, dist2 = distChamfer( points2.transpose(2, 1).contiguous(), pointsReconstructed) loss_net = (torch.mean(dist1)) + (torch.mean(dist2)) if loss_net < bestLoss: bestLoss = loss_net best_theta = theta # unrotate the mesh norma3 = norma.unsqueeze(0).expand(pointsReconstructed.size(1), 3).contiguous() pointsReconstructed[0] = pointsReconstructed[0] + norma3 rot_matrix = np.array([[np.cos(-theta), 0, np.sin(-theta)], [0, 1, 0], [-np.sin(-theta), 0, np.cos(-theta)]]) rot_matrix = Variable(torch.from_numpy(rot_matrix).float()).cuda() pointsReconstructed = torch.matmul(pointsReconstructed, rot_matrix.transpose(1, 0)) bestPoints = pointsReconstructed # print("best loss and angle : ", bestLoss.data[0], best_theta) val_loss.update(bestLoss.data[0]) if global_variables.opt.HR: faces_tosave = global_variables.network.mesh_HR.faces else: faces_tosave = global_variables.network.mesh.faces # create initial guess mesh = trimesh.Trimesh( vertices=(bestPoints[0].data.cpu().numpy() + translation) / scalefactor, faces=global_variables.network.mesh.faces, process=False) #START REGRESSION print("start regression...") # rotate with optimal angle rot_matrix = np.array([[np.cos(best_theta), 0, np.sin(best_theta)], [0, 1, 0], [-np.sin(best_theta), 0, np.cos(best_theta)]]) rot_matrix = Variable(torch.from_numpy(rot_matrix).float()).cuda() points2 = torch.matmul(rot_matrix, points) mesh_tmp = trimesh.Trimesh(vertices=points2[0].transpose( 1, 0).data.cpu().numpy(), faces=global_variables.network.mesh.faces) bbox = np.array([[ np.max(mesh_tmp.vertices[:, 0]), np.max(mesh_tmp.vertices[:, 1]), np.max(mesh_tmp.vertices[:, 2]) ], [ np.min(mesh_tmp.vertices[:, 0]), np.min(mesh_tmp.vertices[:, 1]), np.min(mesh_tmp.vertices[:, 2]) ]]) norma = Variable(torch.from_numpy((bbox[0] + bbox[1]) / 2).float().cuda()) norma2 = norma.unsqueeze(1).expand(3, points2.size(2)).contiguous() points2[0] = points2[0] - norma2 pointsReconstructed1 = regress(points2) # unrotate with optimal angle norma3 = norma.unsqueeze(0).expand(pointsReconstructed1.size(1), 3).contiguous() rot_matrix = np.array([[np.cos(-best_theta), 0, np.sin(-best_theta)], [0, 1, 0], [-np.sin(-best_theta), 0, np.cos(-best_theta)]]) rot_matrix = Variable(torch.from_numpy(rot_matrix).float()).cuda() pointsReconstructed1[0] = pointsReconstructed1[0] + norma3 pointsReconstructed1 = torch.matmul(pointsReconstructed1, rot_matrix.transpose(1, 0)) # create optimal reconstruction meshReg = trimesh.Trimesh( vertices=(pointsReconstructed1[0].data.cpu().numpy() + translation) / scalefactor, faces=faces_tosave, process=False) print("... Done!") return mesh, meshReg
def run(input, scalefactor): """ :param input: input mesh to reconstruct optimally. :return: final reconstruction after optimisation """ input, translation = center_V(input) if not global_variables.opt.HR: mesh_ref = global_variables.mesh_ref_LR else: mesh_ref = global_variables.mesh_ref ## Extract points and put them on GPU points = input if np.shape(points)[0] > 10000: random_sample = np.random.choice(np.shape(points)[0], size=10000) else: random_sample = np.arange(np.shape(points)[0]) points = torch.from_numpy(points.astype( np.float32)).contiguous().unsqueeze(0) points = Variable(points) points = points.transpose(2, 1).contiguous() points = points.cuda() # Get a low resolution PC to find the best reconstruction after a rotation on the Y axis points_LR = torch.from_numpy(input[random_sample].astype( np.float32)).contiguous().unsqueeze(0) points_LR = Variable(points_LR) points_LR = points_LR.transpose(2, 1).contiguous() points_LR = points_LR.cuda() theta = 0 flip_y = 1 bestLoss = 10 pointsReconstructed = global_variables.network(points_LR) dist1, dist2 = distChamfer( points_LR.transpose(2, 1).contiguous(), pointsReconstructed) loss_net = (torch.mean(dist1)) + (torch.mean(dist2)) # print("loss : ", loss_net.data[0], 0) # ---- Search best angle for best reconstruction on the Y axis--- for flip_y in [-1, 1]: for theta in np.linspace(-np.pi, np.pi, global_variables.opt.num_angles): if global_variables.opt.num_angles == 1: theta = 0 # Rotate mesh by theta and renormalise rot_y_matrix = rotation_matrix(theta, flip_y) points2 = torch.matmul(rot_y_matrix, points_LR) mesh_vert = points2[0].transpose(1, 0).detach().data bbox0 = torch.max(mesh_vert, dim=0)[0] bbox1 = torch.min(mesh_vert, dim=0)[0] norma = Variable((bbox0 + bbox1) / 2) norma2 = norma.unsqueeze(1).expand(3, points2.size(2)).contiguous() points2[0] = points2[0] - norma2 # reconstruct rotated mesh pointsReconstructed = global_variables.network(points2) dist1, dist2 = distChamfer( points2.transpose(2, 1).contiguous(), pointsReconstructed) loss_net = (torch.mean(dist1)) + (torch.mean(dist2)) if loss_net < bestLoss: bestLoss = loss_net best_theta, best_flip_y = theta, flip_y # unrotate the mesh norma3 = norma.unsqueeze(0).expand(pointsReconstructed.size(1), 3).contiguous() pointsReconstructed[0] = pointsReconstructed[0] + norma3 rot_y_matrix = rotation_matrix(-theta, flip_y) pointsReconstructed = torch.matmul( pointsReconstructed, rot_y_matrix.transpose(1, 0)) bestPoints = pointsReconstructed # print("best loss and angle : ", bestLoss.data[0], best_theta) val_loss.update(bestLoss.data[0]) #START REGRESSION print("start regression...") # rotate with optimal angle rot_y_matrix = rotation_matrix(best_theta, best_flip_y) points2 = torch.matmul(rot_y_matrix, points) mesh_vert = points2[0].transpose(1, 0).detach().data bbox0 = torch.max(mesh_vert, dim=0)[0] bbox1 = torch.min(mesh_vert, dim=0)[0] norma = Variable((bbox0 + bbox1) / 2) norma2 = norma.unsqueeze(1).expand(3, points2.size(2)).contiguous() points2[0] = points2[0] - norma2 pointsReconstructed1, final_loss = regress(points2) # unrotate with optimal angle norma3 = norma.unsqueeze(0).expand(pointsReconstructed1.size(1), 3).contiguous() rot_y_matrix = rotation_matrix(-best_theta, best_flip_y) pointsReconstructed1[0] = pointsReconstructed1[0] + norma3 pointsReconstructed1 = torch.matmul(pointsReconstructed1, rot_y_matrix.transpose(1, 0)) # create optimal reconstruction print("... Done!") final_points = (pointsReconstructed1[0].data.cpu().numpy() + translation) / scalefactor return final_points, final_loss