def close(self):
# close writers
for key, gw in self.gif_writers.items():
if isinstance(gw, str):
continue
gw.close()
self.gif_writers[key] = join(self.folder_locs[key],
f"camera_simulation_{key}.gif")
# generate json file for the render
json_dict = self._dict()
json_file = join(self.folder_locs["base"], "info.json")
with open(json_file, mode="w") as f:
json.dump(json_dict, f)
if self.pred_results:
results_file = join(self.folder_locs["base"],
"reconstruction_results.json")
with open(results_file, mode="w") as f:
json.dump(self.pred_results, f)
if "predicted_mesh" in self.__dict__.keys():
for mesh, name in self.predicted_mesh:
mesh_path = join(self.folder_locs["base"],
f"predicted_mesh{name}.obj")
verts, faces = mesh.get_mesh_verts_faces(0)
save_obj(mesh_path, verts, faces)
if "pose_plot" in self.__dict__.keys():
for plot, name in self.pose_plot:
plot_path = join(self.folder_locs["base"],
f"pose_plot_{name}.png")
plot.savefig(plot_path, dpi=plot.dpi)
def postprocess_data(cached_pred_poses, output_dir_mesh, cfg, device,
recompute_meshes):
input_dir_img = cfg['dataset']['input_dir_img']
input_dir_mesh = cfg['dataset']['input_dir_mesh']
# postprocessing each mesh/img in dataset
refiner = MeshRefiner(cfg, device)
loss_info = {}
tqdm_out = utils.TqdmPrintEvery()
for instance_name in tqdm(cached_pred_poses, file=tqdm_out):
curr_obj_path = os.path.join(output_dir_mesh, instance_name + ".obj")
if recompute_meshes or not os.path.exists(curr_obj_path):
input_image = np.asarray(
Image.open(os.path.join(input_dir_img,
instance_name + ".png")))
with torch.no_grad():
mesh = utils.load_untextured_mesh(
os.path.join(input_dir_mesh, instance_name + ".obj"),
device)
pred_dist = cached_pred_poses[instance_name]['dist']
pred_elev = cached_pred_poses[instance_name]['elev']
pred_azim = cached_pred_poses[instance_name]['azim']
curr_refined_mesh, curr_loss_info = refiner.refine_mesh(
mesh, input_image, pred_dist, pred_elev, pred_azim)
save_obj(curr_obj_path, curr_refined_mesh.verts_packed(),
curr_refined_mesh.faces_packed())
loss_info[instance_name] = curr_loss_info
return loss_info
def _bm_load_obj(verts: torch.Tensor, faces: torch.Tensor,
decimal_places: int):
f = StringIO()
save_obj(f, verts, faces, decimal_places)
s = f.getvalue()
# Recreate stream so it's unaffected by how it was created.
return lambda: load_obj(StringIO(s))
def test_save_obj(self):
verts = torch.tensor(
[
[0.01, 0.2, 0.301],
[0.2, 0.03, 0.408],
[0.3, 0.4, 0.05],
[0.6, 0.7, 0.8],
],
dtype=torch.float32,
)
faces = torch.tensor([[0, 2, 1], [0, 1, 2], [3, 2, 1], [3, 1, 0]],
dtype=torch.int64)
obj_file = StringIO()
save_obj(obj_file, verts, faces, decimal_places=2)
expected_file = "\n".join([
"v 0.01 0.20 0.30",
"v 0.20 0.03 0.41",
"v 0.30 0.40 0.05",
"v 0.60 0.70 0.80",
"f 1 3 2",
"f 1 2 3",
"f 4 3 2",
"f 4 2 1",
])
actual_file = obj_file.getvalue()
self.assertEqual(actual_file, expected_file)
def translate_mesh_on_axis(
mesh, t: list, renderer, dist: float = 3.5, save: str = "", device: str = ""
):
translation = Transform3d(device=device).translate(t[0], t[1], t[2])
verts, faces = mesh.get_mesh_verts_faces(0)
verts = translation.transform_points(verts)
mesh = mesh.update_padded(verts.unsqueeze(0))
dist = dist
elev = 0
azim = 0
R, T = look_at_view_transform(dist=dist, elev=elev, azim=azim, device=device)
image_ref = renderer(meshes_world=mesh, R=R, T=T, device=device)
image_ref = image_ref.cpu().numpy()
plt.imshow(image_ref.squeeze())
plt.show()
if save:
verts, faces = mesh.get_mesh_verts_faces(0)
save_obj(save, verts, faces)
return mesh
def run_on_images(self, imgs, sid, mid, iid, sampled_idx):
dir1 = os.path.join(output_dir, str(sid), str(mid))
if not os.path.exists(dir1):
os.makedirs(dir1)
deprocess = imagenet_deprocess(rescale_image=False)
image_features = self.encoder(imgs)
raw_features, generated_volume = self.decoder(image_features)
generated_volume = self.merger(raw_features, generated_volume)
generated_volume = self.refiner(generated_volume)
mesh = cubify(generated_volume, 0.3)
# mesh = voxel_to_world(meshes)
save_mesh = os.path.join(dir1, "%s_%s.obj" % (iid, sampled_idx))
verts, faces = mesh.get_mesh_verts_faces(0)
save_obj(save_mesh, verts, faces)
generated_volume = generated_volume.squeeze()
img = image_to_numpy(deprocess(imgs[0][0]))
save_img = os.path.join(dir1, "%02d.png" % (iid))
# cv2.imwrite(save_img, img[:, :, ::-1])
cv2.imwrite(save_img, img)
img1 = image_to_numpy(deprocess(imgs[0][1]))
save_img1 = os.path.join(dir1, "%02d.png" % (sampled_idx))
cv2.imwrite(save_img1, img1)
# cv2.imwrite(save_img1, img1[:, :, ::-1])
get_volume_views(generated_volume, dir1, iid, sampled_idx)
def rotate_mesh_around_axis(
mesh, rot: list, renderer, dist: float = 3.5, save: str = "", device: str = ""
):
if not device:
device = torch.cuda.current_device()
rot_x = RotateAxisAngle(rot[0], "X", device=device)
rot_y = RotateAxisAngle(rot[1], "Y", device=device)
rot_z = RotateAxisAngle(rot[2], "Z", device=device)
rot = Transform3d(device=device).stack(rot_x, rot_y, rot_z)
verts, faces = mesh.get_mesh_verts_faces(0)
verts = rot_x.transform_points(verts)
verts = rot_y.transform_points(verts)
verts = rot_z.transform_points(verts)
mesh = mesh.update_padded(verts.unsqueeze(0))
dist = dist
elev = 0
azim = 0
R, T = look_at_view_transform(dist=dist, elev=elev, azim=azim, device=device)
image_ref = renderer(meshes_world=mesh, R=R, T=T, device=device)
image_ref = image_ref.cpu().numpy()[..., :3]
plt.imshow(image_ref.squeeze())
plt.show()
if save:
verts, faces = mesh.get_mesh_verts_faces(0)
save_obj(save, verts, faces)
return mesh
def visualize_predictions(
image_id,
image_file,
scores,
labels,
boxes,
mask_rles,
meshes,
metadata,
output_dir,
alpha=0.6,
dpi=200,
):
# create vis_dir
output_dir = os.path.join(output_dir, "results")
os.makedirs(output_dir, exist_ok=True)
cat_colors = metadata.thing_colors
cat_names = metadata.thing_classes
# read image
image_file = os.path.join(metadata.image_root, image_file)
image = detection_utils.read_image(image_file, format="RGB")
num_preds = len(scores)
for i in range(num_preds):
# box
box = boxes[i].view(1, 4)
# RLE to 2D mask
mask = mask_util.decode(mask_rles[i])
label = labels[i]
mask_color = np.array(cat_colors[label], dtype=np.float32)
cat_name = cat_names[label]
score = scores[i]
# plot mask overlayed
composite = image.copy()
composite = draw_mask(composite,
mask,
mask_color,
alpha=alpha,
draw_contours=False)
thickness = int(np.ceil(0.001 * image.shape[0]))
composite = draw_boxes(composite, box, thickness)
save_file = os.path.join(
output_dir, "%d_%d_%s_%.3f.png" % (image_id, i, cat_name, score))
cv2.imwrite(save_file, composite[:, :, ::-1])
save_file = os.path.join(
output_dir, "%d_%d_%s_%.3f.obj" % (image_id, i, cat_name, score))
verts, faces = meshes.get_mesh_verts_faces(i)
save_obj(save_file, verts, faces)
def test_save_obj_invalid_indices(self):
message_regex = "Faces have invalid indices"
verts = torch.FloatTensor([[0.1, 0.2, 0.3]])
faces = torch.LongTensor([[0, 1, 2]])
with self.assertWarnsRegex(UserWarning, message_regex):
save_obj(StringIO(), verts, faces)
faces = torch.LongTensor([[-1, 0, 1]])
with self.assertWarnsRegex(UserWarning, message_regex):
save_obj(StringIO(), verts, faces)
def visualize_prediction(image_id, img, mesh, output_dir):
# create vis_dir
output_dir = os.path.join(output_dir, "results_shapenet")
os.makedirs(output_dir, exist_ok=True)
save_img = os.path.join(output_dir, "%s.png" % (image_id))
cv2.imwrite(save_img, img[:, :, ::-1])
save_mesh = os.path.join(output_dir, "%s.obj" % (image_id))
verts, faces = mesh.get_mesh_verts_faces(0)
save_obj(save_mesh, verts, faces)
def _test_save_load(self, verts, faces):
f = StringIO()
save_obj(f, verts, faces)
f.seek(0)
expected_verts, expected_faces = verts, faces
if not len(expected_verts): # Always compare with a (V, 3) tensor
expected_verts = torch.zeros(size=(0, 3), dtype=torch.float32)
if not len(expected_faces): # Always compare with an (F, 3) tensor
expected_faces = torch.zeros(size=(0, 3), dtype=torch.int64)
actual_verts, actual_faces, _ = load_obj(f)
self.assertClose(expected_verts, actual_verts)
self.assertClose(expected_faces, actual_faces.verts_idx)
def test_save_obj_invalid_indices(self):
message_regex = "Faces have invalid indices"
verts = torch.FloatTensor([[0.1, 0.2, 0.3]])
faces = torch.LongTensor([[0, 1, 2]])
with self.assertWarnsRegex(UserWarning, message_regex):
with NamedTemporaryFile(mode="w", suffix=".obj") as f:
save_obj(Path(f.name), verts, faces)
faces = torch.LongTensor([[-1, 0, 1]])
with self.assertWarnsRegex(UserWarning, message_regex):
with NamedTemporaryFile(mode="w", suffix=".obj") as f:
save_obj(Path(f.name), verts, faces)
def save_merged_voxel_grids(
file_prefix, voxel_scores, voxel_size, cubify_threshold
):
"""
save merged voxel grids for debugging purpose
"""
from pytorch3d.io import save_obj
import open3d as o3d
meshes = cubify(voxel_scores, voxel_size, cubify_threshold)
for batch_idx, mesh in enumerate(meshes):
filename = "/tmp/cube_mesh_{}_{}_merged.obj" \
.format(file_prefix, batch_idx)
save_obj(filename, mesh.verts_packed(), mesh.faces_packed())
def _test_save_load(self, verts, faces):
with NamedTemporaryFile(mode="w", suffix=".obj") as f:
file_path = Path(f.name)
save_obj(file_path, verts, faces)
f.flush()
expected_verts, expected_faces = verts, faces
if not len(expected_verts): # Always compare with a (V, 3) tensor
expected_verts = torch.zeros(size=(0, 3), dtype=torch.float32)
if not len(expected_faces): # Always compare with an (F, 3) tensor
expected_faces = torch.zeros(size=(0, 3), dtype=torch.int64)
actual_verts, actual_faces, _ = load_obj(file_path)
self.assertClose(expected_verts, actual_verts)
self.assertClose(expected_faces, actual_faces.verts_idx)
def test_save_obj_invalid_shapes(self):
# Invalid vertices shape
with self.assertRaises(ValueError) as error:
verts = torch.FloatTensor([[0.1, 0.2, 0.3, 0.4]]) # (V, 4)
faces = torch.LongTensor([[0, 1, 2]])
save_obj(StringIO(), verts, faces)
expected_message = "Argument 'verts' should either be empty or of shape (num_verts, 3)."
self.assertTrue(expected_message, error.exception)
# Invalid faces shape
with self.assertRaises(ValueError) as error:
verts = torch.FloatTensor([[0.1, 0.2, 0.3]])
faces = torch.LongTensor([[0, 1, 2, 3]]) # (F, 4)
save_obj(StringIO(), verts, faces)
expected_message = "Argument 'faces' should either be empty or of shape (num_faces, 3)."
self.assertTrue(expected_message, error.exception)
def create_data(folder_path='meshes/',
nb_of_pointclouds=50,
nb_of_points=5000,
sphere_level=4,
normalize_data=True):
device = torch.device("cuda:0")
data_path = os.path.join(os.getcwd(), folder_path)
src_mesh = ico_sphere(sphere_level, device)
for filename in os.listdir(data_path):
print(f"{datetime.now()} Starting:{filename}")
file_path = os.path.join(data_path, filename)
cur_mesh = utils.load_mesh(file_path)
cur_deform_verts = deformation.get_deform_verts(
cur_mesh, nb_of_points, sphere_level)
data_verts = np.expand_dims(cur_deform_verts.detach().cpu().numpy(),
axis=0)
data_input = None
data_output = None
for _ in range(nb_of_pointclouds):
data_a = sample_points_from_meshes(
cur_mesh, nb_of_points).squeeze().cpu().numpy()
if normalize_data:
data_a = data_a - np.mean(data_a, axis=0)
data_a = data_a / np.max(data_a, axis=0)
data_a_sort_indices = np.argsort(np.linalg.norm(data_a,
axis=1))
data_a = data_a[data_a_sort_indices]
data_a = np.expand_dims(data_a, axis=0)
data_input = data_a if data_input is None else np.concatenate(
(data_input, data_a))
data_output = data_verts if data_output is None else np.concatenate(
(data_output, data_verts))
np.save(f'data/{os.path.splitext(filename)[0]}_input.npy', data_input)
np.save(f'data/{os.path.splitext(filename)[0]}_output.npy',
data_output)
deformed_mesh = src_mesh.offset_verts(cur_deform_verts)
final_verts, final_faces = deformed_mesh.get_mesh_verts_faces(0)
final_obj = os.path.join(
'deformed_meshes/',
f'{os.path.splitext(filename)[0]}_deformed.obj')
save_obj(final_obj, final_verts, final_faces)
print(
f"{datetime.now()} Finished:{filename}, Point Cloud Shape:{data_input.shape} Deform Verts Shape:{data_output.shape}"
)
def main():
if cfg.is_save_from_mat:
mat_path = os.path.join('../Data', 'modelnet40_2111instances10000_PointNet.mat')
dataset = sio.loadmat(mat_path)
trg_pc = torch.FloatTensor(dataset["data"]).cuda()
if not os.path.exists(os.path.join('../Data', 'All_class_ori_mesh')):
os.makedirs(os.path.join('../Data', 'All_class_ori_mesh'))
for i in range(trg_pc.size(0)):
curr_trg_pc = trg_pc[i].unsqueeze(0)
fout = open(os.path.join('../Data', 'All_class_ori_mesh', str(i)+'.xyz'), 'w')
for m in range(curr_trg_pc.shape[2]):
fout.write('%f %f %f \n' % (curr_trg_pc[0, 0, m], curr_trg_pc[0, 1, m], curr_trg_pc[0, 2, m]))
fout.close()
else:
#from modelnet40_with_vert import ModelNet40_vert
from modelnet_trn_test import ModelNetDataset
test_dataset = ModelNetDataset(root='/data/modelnet40_normal_resampled/', batch_size=1, npoints=cfg.npoint, split='test', normal_channel=True)
#test_dataset = ModelNet40_vert('/data/ModelNet40/', 40, phase='test', regenerate_dataset=False)
test_loader = torch.utils.data.DataLoader(test_dataset, 1, shuffle=True, drop_last=True, num_workers=8, pin_memory=True)
test_size = test_dataset.__len__()
if not os.path.exists(os.path.join('../Data', 'Ten_class_ori_mesh')):
os.makedirs(os.path.join('../Data', 'Ten_class_ori_mesh'))
for i, (vert, faces, label) in enumerate(test_loader):
if convert_from_modelnet40_1024_processed[label[0]] in label_indexes:
vert = vert.squeeze(0)
faces = faces.squeeze(0)
vert, _, _ = pc_normalize_torch(vert)
trg_mesh = Meshes(verts=[vert], faces=[faces]).cuda()
file_name = os.path.join('../Data', 'Ten_class_ori_mesh', str(i)+'_'+str(convert_from_modelnet40_1024_processed[label[0]])+'.obj')
final_verts, final_faces = trg_mesh.get_mesh_verts_faces(0)
print('Processing ['+str(i)+'/'+str(test_size)+' ] instance')
save_obj(file_name, final_verts, final_faces)
def test_save_obj_invalid_shapes(self):
# Invalid vertices shape
with self.assertRaises(ValueError) as error:
verts = torch.FloatTensor([[0.1, 0.2, 0.3, 0.4]]) # (V, 4)
faces = torch.LongTensor([[0, 1, 2]])
with NamedTemporaryFile(mode="w", suffix=".obj") as f:
save_obj(Path(f.name), verts, faces)
expected_message = (
"Argument 'verts' should either be empty or of shape (num_verts, 3)."
)
self.assertTrue(expected_message, error.exception)
# Invalid faces shape
with self.assertRaises(ValueError) as error:
verts = torch.FloatTensor([[0.1, 0.2, 0.3]])
faces = torch.LongTensor([[0, 1, 2, 3]]) # (F, 4)
with NamedTemporaryFile(mode="w", suffix=".obj") as f:
save_obj(Path(f.name), verts, faces)
expected_message = (
"Argument 'faces' should either be empty or of shape (num_faces, 3)."
)
self.assertTrue(expected_message, error.exception)
def test_save_obj(self):
verts = torch.tensor(
[[0.01, 0.2, 0.301], [0.2, 0.03, 0.408], [0.3, 0.4, 0.05],
[0.6, 0.7, 0.8]],
dtype=torch.float32,
)
faces = torch.tensor([[0, 2, 1], [0, 1, 2], [3, 2, 1], [3, 1, 0]],
dtype=torch.int64)
with NamedTemporaryFile(mode="w", suffix=".obj") as f:
save_obj(Path(f.name), verts, faces, decimal_places=2)
expected_file = "\n".join([
"v 0.01 0.20 0.30",
"v 0.20 0.03 0.41",
"v 0.30 0.40 0.05",
"v 0.60 0.70 0.80",
"f 1 3 2",
"f 1 2 3",
"f 4 3 2",
"f 4 2 1",
])
actual_file = open(Path(f.name), "r")
self.assertEqual(actual_file.read(), expected_file)
def test_save_obj_with_texture(self):
device = torch.device('cuda:0')
torch.cuda.set_device(device)
# Set paths
data_dir = Path(__file__).parent / 'data'
data_dir.mkdir(exist_ok=True)
obj_dir = Path(
__file__).resolve().parent.parent / "docs/tutorials/data"
obj_filename = obj_dir / "cow_mesh/cow.obj"
final_obj = data_dir / "cow_exported.obj"
# Load obj file
mesh = load_objs_as_meshes([obj_filename], device=device)
try:
texture_image = mesh.textures.maps_padded()
save_obj(final_obj,
mesh.verts_packed(),
mesh.faces_packed(),
verts_uvs=mesh.textures.verts_uvs_packed(),
texture_map=texture_image,
faces_uvs=mesh.textures.faces_uvs_packed())
except:
pass
t.set_description("loss = {}".format(loss))
if WANDB:
wandb.log({"Train Loss": loss})
losses.append(loss.detach())
# Plot mesh
if i % plot_period == 0 and i!=0:
plot_pointcloud(new_block1, title="iter: %d" % i)
plot_pointcloud(new_block2, title="iter: %d" % i)
plot_pointcloud(new_block3, title="iter: %d" % i)
plt.imshow(image.squeeze().permute(1, 2, 0).detach().cpu().numpy())
plt.show()
plt.plot(range(len(losses)), np.log10(losses))
plt.show()
save_obj("model1.obj", new_block1.verts_packed(), new_block1.faces_packed())
save_obj("model2.obj", new_block2.verts_packed(), new_block2.faces_packed())
save_obj("model3.obj", new_block3.verts_packed(), new_block3.faces_packed())
if i % save_period == 0 and i !=0:
print("saving ...")
if WANDB:
torch.save({
'epoch': epoch,
'iteration':i,
'gcn_state_dict': [gcn1.state_dict(), gcn2.state_dict(), gcn3.state_dict()],
'fe_state_dic': fe.state_dict(),
'cnn_state_dic': image_model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss}, os.path.join(wandb.run.dir, 'dic_checkpoint'))
def bm_save_simple_obj_with_init(V: int, F: int):
verts_list = torch.tensor(V * [[0.11, 0.22, 0.33]]).view(-1, 3)
faces_list = torch.tensor(F * [[1, 2, 3]]).view(-1, 3)
return lambda: save_obj(
StringIO(), verts_list, faces_list, decimal_places=2
)
def _bm_save_obj(verts: torch.Tensor, faces: torch.Tensor,
decimal_places: int):
return lambda: save_obj(StringIO(), verts, faces, decimal_places)
try:
if save_result or (opt.isTrain and not opt.no_html):
estimated_obj_path = os.path.join(
visualizer.img_dir, 'epoch%.3d_%s.obj' %
(total_iters // opt.batch_size, 'estimated_mesh'))
true_obj_path = os.path.join(
visualizer.img_dir, 'epoch%.3d_%s.obj' %
(total_iters // opt.batch_size, 'true_mesh'))
save_obj(estimated_obj_path,
model.estimated_mesh[
model.verbose_batch_ind].verts_packed(),
torch.from_numpy(
model.flamelayer.faces.astype(np.int32)),
verts_uvs=model.estimated_mesh[
model.verbose_batch_ind].textures.
verts_uvs_packed(),
texture_map=model.estimated_texture_map[
None, model.verbose_batch_ind],
faces_uvs=model.estimated_mesh[
model.verbose_batch_ind].textures.
faces_uvs_packed())
save_obj(
true_obj_path,
model.true_mesh[
model.verbose_batch_ind].verts_packed(),
torch.from_numpy(
model.flamelayer.faces.astype(np.int32)),
verts_uvs=model.true_mesh[model.verbose_batch_ind].
textures.verts_uvs_packed(),
texture_map=model.true_mesh.textures.maps_padded()[
def test_cube_mesh_render(self):
"""
End-End test of rendering a cube mesh with texture
from decreasing camera distances. The camera starts
outside the cube and enters the inside of the cube.
"""
device = torch.device("cuda:0")
mesh = self.load_cube_mesh_with_texture(device)
raster_settings = RasterizationSettings(
image_size=512,
blur_radius=1e-8,
faces_per_pixel=5,
z_clip_value=1e-2,
perspective_correct=True,
bin_size=0,
)
# Only ambient, no diffuse or specular
lights = PointLights(
device=device,
ambient_color=((1.0, 1.0, 1.0), ),
diffuse_color=((0.0, 0.0, 0.0), ),
specular_color=((0.0, 0.0, 0.0), ),
location=[[0.0, 0.0, -3.0]],
)
renderer = MeshRenderer(
rasterizer=MeshRasterizer(raster_settings=raster_settings),
shader=SoftPhongShader(device=device, lights=lights),
)
# Render the cube by decreasing the distance from the camera until
# the camera enters the cube. Check the output looks correct.
images_list = []
dists = np.linspace(0.1, 2.5, 20)[::-1]
for d in dists:
R, T = look_at_view_transform(d, 0, 0)
T[0, 1] -= 0.1 # move down in the y axis
cameras = FoVPerspectiveCameras(device=device, R=R, T=T, fov=90)
images = renderer(mesh, cameras=cameras)
rgb = images[0, ..., :3].cpu().detach()
filename = "DEBUG_cube_dist=%.1f.jpg" % d
im = (rgb.numpy() * 255).astype(np.uint8)
images_list.append(im)
# Check one of the images where the camera is inside the mesh
if d == 0.5:
filename = "test_render_mesh_clipped_cam_dist=0.5.jpg"
image_ref = load_rgb_image(filename, DATA_DIR)
self.assertClose(rgb, image_ref, atol=0.05)
# Save a gif of the output - this should show
# the camera moving inside the cube.
if DEBUG:
gif_filename = ("room_original.gif"
if raster_settings.z_clip_value is None else
"room_clipped.gif")
imageio.mimsave(DATA_DIR / gif_filename, images_list, fps=2)
save_obj(
f=DATA_DIR / "cube.obj",
verts=mesh.verts_packed().cpu(),
faces=mesh.faces_packed().cpu(),
)
def save_mesh():
save_obj(obj_file, verts_list, faces_list, decimal_places=2)
#laplacian_losses.append(loss_laplacian)
losses.append(loss)
# Plot mesh
if i % plot_period == 0 and i!=0:
plot_pointcloud(new_block1, title="iter: %d" % i)
plot_pointcloud(new_block2, title="iter: %d" % i)
plot_pointcloud(new_block3, title="iter: %d" % i)
# Optimization step
loss.backward()
torch.nn.utils.clip_grad_norm_(gcn1.parameters(), 1)
torch.nn.utils.clip_grad_norm_(gcn2.parameters(), 1)
torch.nn.utils.clip_grad_norm_(gcn3.parameters(), 1)
optimizer.step()
# for name, param in gcn1.named_parameters():
# print(name, torch.isfinite(param.grad).all())
# for name, param in gcn2.named_parameters():
# print(name, torch.isfinite(param.grad).all())
# for name, param in gcn3.named_parameters():
# print(name, torch.isfinite(param.grad).all())
save_obj("model_3d_to_3d_1.obj", new_block1.verts_packed(), new_block1.faces_packed())
save_obj("model_3d_to_3d_2.obj", new_block2.verts_packed(), new_block2.faces_packed())
save_obj("model_3d_to_3d_3.obj", new_block3.verts_packed(), new_block3.faces_packed())
plt.plot(range(len(losses)), np.log(losses))
plt.show()
# Print the losses
t.set_description('total_loss = %.6f' % loss)
# Save the losses for plotting
chamfer_losses.append(loss_chamfer)
edge_losses.append(loss_edge)
normal_losses.append(loss_normal)
laplacian_losses.append(loss_laplacian)
losses.append(loss)
# Plot mesh
if i % plot_period == 0:
plot_pointcloud(new_src_mesh, title="iter: %d" % i)
plt.show()
# Optimization step
loss.backward()
optimizer.step()
plt.plot(range(len(losses)), losses)
plt.show()
# Fetch the verts and faces of the final predicted mesh
final_verts, final_faces = new_src_mesh.get_mesh_verts_faces(0)
# Scale normalize back to the original target size
final_verts = final_verts * scale + center
# Store the predicted mesh using save_obj
final_obj = os.path.join('../', '3dTo3d/final_model.obj')
save_obj(final_obj, final_verts, final_faces)
print('Error! No existing model!')
exit(-1)
net.eval()
src_pc, _, _ = load_obj("%s/pc_4096.obj" % opt.src_shape)
tar_pc, _, _ = load_obj("%s/pc_4096.obj" % opt.tar_shape)
key_pts, _, _ = load_obj("%s/key_point_50.obj" % opt.src_shape)
_, src_faces, _ = load_obj("%s/manifold.obj" % opt.src_shape)
src_faces = src_faces.verts_idx
w_pc = torch.from_numpy(np.load("%s/w_pc_4096.npy" % opt.src_shape))
w_mesh = torch.from_numpy(np.load("%s/w_mesh.npy" % opt.src_shape))
src_pc = src_pc.unsqueeze(0).cuda()
tar_pc = tar_pc.unsqueeze(0).cuda()
key_pts = key_pts.unsqueeze(0).cuda()
w_pc = w_pc.unsqueeze(0).cuda()
_, _, _, basis, _, _, _, coef_range = net(src_pc, tar_pc, key_pts, w_pc)
os.makedirs(opt.save_dir, exist_ok=True)
for i in range(opt.num_basis):
l = coef_range[0][i][0] * 2
r = coef_range[0][i][1] * 2
for j in range(4):
scale = (r - l) / 3 * j + l
off = basis[0][i].reshape(50, 3) * scale
def_key_pts = key_pts[0] + off
def_ver = torch.matmul(w_mesh.cuda(), def_key_pts)
save_obj(("%s/%d-%d.obj") % (opt.save_dir, i, j), def_ver.cpu(),
src_faces)
opt = parse_args()
net = network.model(opt.num_basis).cuda()
net = torch.nn.DataParallel(net)
try:
checkpoint = torch.load(opt.checkpoint)
net.load_state_dict(checkpoint['model_state_dict'])
print('Use pretrain model')
except:
print('Error! No existing model!')
exit(-1)
net.eval()
src_pc, _, _ = load_obj("%s/pc_4096.obj" % opt.src_shape)
tar_pc, _, _ = load_obj("%s/pc_4096.obj" % opt.tar_shape)
key_pts, _, _ = load_obj("%s/key_point_50.obj" % opt.src_shape)
_, src_faces, _ = load_obj("%s/manifold.obj" % opt.src_shape)
src_faces = src_faces.verts_idx
w_pc = torch.from_numpy(np.load("%s/w_pc_4096.npy" % opt.src_shape))
w_mesh = torch.from_numpy(np.load("%s/w_mesh.npy" % opt.src_shape))
src_pc = src_pc.unsqueeze(0).cuda()
tar_pc = tar_pc.unsqueeze(0).cuda()
key_pts = key_pts.unsqueeze(0).cuda()
w_pc = w_pc.unsqueeze(0).cuda()
def_key_pts, _, _, _, _, _, _, _ = net(src_pc, tar_pc, key_pts, w_pc)
def_ver = torch.matmul(w_mesh.cuda(), def_key_pts[0])
save_obj(opt.save_name, def_ver.cpu(), src_faces)
