提交测试
This commit is contained in:
137
examples/tutorial/visualize_main.py
Normal file
137
examples/tutorial/visualize_main.py
Normal file
@@ -0,0 +1,137 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import torch
|
||||
import sys
|
||||
|
||||
import kaolin
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def __normalize_vertices(vertices):
|
||||
"""
|
||||
Normalizes vertices to fit an [-1...1] bounding box,
|
||||
common during training, but not necessary for visualization.
|
||||
"""
|
||||
return kaolin.ops.pointcloud.center_points(res.vertices.unsqueeze(0), normalize=True).squeeze(0) * 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Example exporting 3D data during training as timed USDs; '
|
||||
'also demonstrates OBJ import and mesh to pointcloud conversions.')
|
||||
parser.add_argument('--test_objs', type=str, required=True,
|
||||
help='Comma separated list of several example obj files.')
|
||||
parser.add_argument('--output_dir', type=str, required=True,
|
||||
help='Output directory to write checkpoints to; must exist.')
|
||||
parser.add_argument('--iterations', type=int, default=101,
|
||||
help='How many training iterations to emulate.')
|
||||
parser.add_argument('--checkpoint_interval', type=int, default=10,
|
||||
help='Frequency with which to write out checkpoints.')
|
||||
parser.add_argument('--skip_normalization', action='store_true',
|
||||
help='If not set, will normalize bounding box of each input '
|
||||
'to be within -1..1 cube.')
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
|
||||
|
||||
if not os.path.isdir(args.output_dir):
|
||||
raise RuntimeError(
|
||||
'Output directory does not exist: --output_dir={}'.format(
|
||||
args.output_dir))
|
||||
|
||||
# Read test 3D models & setup fake training ----------------------------------
|
||||
obj_files = args.test_objs.split(',')
|
||||
logger.info('Parsing {} OBJ files: '.format(len(obj_files)))
|
||||
face_list = []
|
||||
gt_vert_list = []
|
||||
input_pt_clouds = []
|
||||
delta_list = []
|
||||
delta_pt_list = []
|
||||
# TODO: add textured example
|
||||
for f in obj_files:
|
||||
res = kaolin.io.obj.import_mesh(f)
|
||||
vertices = res.vertices if args.skip_normalization else __normalize_vertices(vertices)
|
||||
num_samples = random.randint(1000, 1500) # Vary to ensure robustness
|
||||
pts = kaolin.ops.mesh.sample_points(
|
||||
vertices.unsqueeze(0), res.faces, num_samples)[0].squeeze(0)
|
||||
|
||||
# Randomly displace vertices to emulate training
|
||||
delta = (2.0 * torch.rand(vertices.shape, dtype=vertices.dtype) - 1.0) * 0.25
|
||||
delta_pts = (2.0 * torch.rand(pts.shape, dtype=pts.dtype) - 1.0) * 0.25
|
||||
|
||||
face_list.append(res.faces)
|
||||
gt_vert_list.append(vertices)
|
||||
delta_list.append(delta)
|
||||
input_pt_clouds.append(pts)
|
||||
delta_pt_list.append(delta_pts)
|
||||
|
||||
# Emulate visualizing during training -------------------------------------
|
||||
logger.info('Emulating training run for {} iterations'.format(args.iterations))
|
||||
|
||||
# Create a Timelapse instance
|
||||
timelapse = kaolin.visualize.Timelapse(args.output_dir)
|
||||
|
||||
# Save static objects such as ground truth or inputs that do not change with iterations
|
||||
# just once.
|
||||
timelapse.add_mesh_batch(
|
||||
category='ground_truth',
|
||||
faces_list=face_list,
|
||||
vertices_list=gt_vert_list)
|
||||
timelapse.add_pointcloud_batch(
|
||||
category='input',
|
||||
pointcloud_list=input_pt_clouds)
|
||||
|
||||
for iteration in range(args.iterations):
|
||||
if iteration % args.checkpoint_interval == 0:
|
||||
# Emulate a training update
|
||||
out_pt_clouds = []
|
||||
out_vert_list = []
|
||||
out_voxels = []
|
||||
for i in range(len(gt_vert_list)):
|
||||
delta_weight = 1.0 - iteration / (args.iterations - 1)
|
||||
out_vert_list.append(gt_vert_list[i] * (1.0 + delta_list[i] * delta_weight))
|
||||
out_pt_clouds.append(input_pt_clouds[i] * (1.0 + delta_pt_list[i] * delta_weight))
|
||||
vg = kaolin.ops.conversions.trianglemeshes_to_voxelgrids(
|
||||
out_vert_list[-1].unsqueeze(0), face_list[i], 30)
|
||||
out_voxels.append(vg.squeeze(0).bool())
|
||||
|
||||
# Save model predictions to track training progress over time
|
||||
timelapse.add_mesh_batch(
|
||||
iteration=iteration,
|
||||
category='output',
|
||||
faces_list=face_list,
|
||||
vertices_list=out_vert_list)
|
||||
timelapse.add_pointcloud_batch(
|
||||
iteration=iteration,
|
||||
category='output',
|
||||
pointcloud_list=out_pt_clouds)
|
||||
timelapse.add_voxelgrid_batch(
|
||||
iteration=iteration,
|
||||
category='output',
|
||||
voxelgrid_list=out_voxels)
|
||||
|
||||
logger.info('Emulated training complete!\n'
|
||||
'You can now view created USD files by running:\n\n'
|
||||
f'kaolin-dash3d --logdir={args.output_dir}\n\n'
|
||||
'And then navigating to localhost:8080\n')
|
||||
|
||||
# TODO(mshugrina): once dash3d is also integrated, write an integration test
|
||||
# to ensure timelapse output is properly parsed by the visualizer
|
||||
Reference in New Issue
Block a user