def save_simple_points_ply(out_filename, points, text=False):
assert points.shape[1] == 3
assert points.ndim == 2
dirname = path.dirname(out_filename)
if not path.isdir(dirname):
os.makedirs(dirname)
log.info('mkdir -p {}'.format(dirname))
xyz_nxyz = points
vertex = np.core.records.fromarrays(xyz_nxyz.T, names='x, y, z',
formats='f4, f4, f4')
el = plyfile.PlyElement.describe(vertex, 'vertex')
ply = plyfile.PlyData([el], text=text)
ply.write(out_filename)
# Replace \r\n with \n.
with open(out_filename, 'rb') as f:
content = f.read()
beginning = content[:128]
rest = content[128:]
beginning = beginning.replace(b'ply\r\n', b'ply\n')
with open(out_filename, 'wb') as f:
f.write(beginning + rest)
return ply
def run_command(command):
log.debug(subprocess.list2cmdline(command))
start_time = time.time()
p = subprocess.Popen(command,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
return_code = p.wait()
log.debug('return code: {}'.format(return_code))
elapsed = time.time() - start_time
stdout, stderr = p.communicate()
stdout = stdout.decode('utf-8')
stderr = stderr.decode('utf-8')
if p.returncode != 0:
exec_summary = textwrap.dedent(
textwrap.dedent("""
Command: {}
Return code: {}
stdout:
{}
stderr:
{}
""").format(subprocess.list2cmdline(command), p.returncode, stdout,
stderr))
raise RuntimeError(exec_summary)
log.info('{0} took {1:.3f} seconds.'.format(path.basename(command[0]),
elapsed))
return p, stdout, stderr
def ensure_dir_exists(dirname, log_mkdir=True):
dirname = path.realpath(path.expanduser(dirname))
if not path.isdir(dirname):
# `exist_ok` in case of race condition.
os.makedirs(dirname, exist_ok=True)
if log_mkdir:
log.info('mkdir -p {}'.format(dirname))
return dirname
def main():
log.info('Getting all filenames.')
syndirs = sorted(glob.glob('/data/render_for_cnn/data/syn_images_cropped_bkg_overlaid/*'))
random.seed(42)
filenames = []
for syndir in syndirs:
modeldirs = sorted(glob.glob(path.join(syndir, '*')))
if is_subset:
modeldirs = modeldirs[:10]
for modeldir in modeldirs:
renderings = sorted(glob.glob(path.join(modeldir, '*')))
if is_subset:
renderings = renderings[:7]
filenames.extend(renderings)
log.info('{} files'.format(len(filenames)))
data_base_dir = '/data/mvshape'
start_time = time.time()
random.seed(42)
log.info('Processing rgb images.')
for i, filename in enumerate(filenames):
m = re.search(r'syn_images_cropped_bkg_overlaid/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a', filename)
synset = m.group(1)
model_name = m.group(2)
v = m.group(3)
image_num = int(v)
vc_rendering_name = '{}_{:04d}'.format(model_name, image_num)
out_filename = path.join(data_base_dir, 'out/shapenetcore/single_rgb_128/{}.png'.format(vc_rendering_name))
assert path.isfile(filename)
io_utils.ensure_dir_exists(path.dirname(out_filename))
img = io_utils.read_jpg(filename)
assert img.shape[0] == img.shape[1]
assert img.shape[2] == 3
resize_method = {0: 'bilinear',
1: 'bicubic',
2: 'lanczos'}[random.randint(0, 2)]
resized_img = scipy.misc.imresize(img, (128, 128), interp=resize_method)
assert resized_img.dtype == np.uint8
# io_utils.save_array_compressed(out_filename, resized_img)
scipy.misc.imsave(out_filename, resized_img)
if i % 100 == 0:
t_elapsed = (time.time() - start_time)
t_remaining = (t_elapsed / (i + 1) * (len(filenames) - i))
log.info('Creating examples in db. {} of {}. elapsed: {:.1f} min, remaining: {:.1f} min'.format(i, len(filenames), t_elapsed / 60, t_remaining / 60))
t_elapsed = (time.time() - start_time)
log.info('total elapsed: {:.1f} min'.format(t_elapsed / 60))
def print_status(models, num_all_examples, io_block_times: list,
accum_scalars):
print('')
log.info(' {:08d} of {:08d}. io block: {:.5f}s \n{}'.format(
models['metadata']['current_epoch_step'],
int(num_all_examples / models['metadata']['batch_size']),
np.mean(io_block_times),
pprint.pformat(
torch_utils.recursive_dict_of_lists_mean(accum_scalars)),
))
def save_model(model, save_dir):
current_epoch = model['metadata']['current_epoch']
current_epoch_step = model['metadata']['current_epoch_step']
global_step = model['metadata']['global_step']
filename = path.join(
save_dir,
'models0_{:05}_{:07}_{:08}.pth'.format(current_epoch,
current_epoch_step,
global_step))
io_utils.ensure_dir_exists(path.dirname(filename))
with open(filename, 'wb') as f:
log.info('Saving.. {}'.format(filename))
torch.save(model, f)
log.info('Saved.')
def save_points_ply(out_filename, points, normals, values, confidence, text=False, scale_normals_by_confidence=True):
assert points.shape == normals.shape
assert points.shape[1] == 3
assert points.ndim == 2
dirname = path.dirname(out_filename)
if not path.isdir(dirname):
os.makedirs(dirname)
log.info('mkdir -p {}'.format(dirname))
if values.ndim == 1:
values = values[:, None]
if confidence.ndim == 1:
confidence = confidence[:, None]
# Replace 0 confidence values with a small number to avoid zero-length normals.
confidence = confidence.copy()
confidence[np.isclose(confidence, 0.0)] = 1e-3
if scale_normals_by_confidence:
normals = normals.copy() * confidence
xyz_nxyz = np.concatenate([points, normals, values, confidence], axis=1).astype(np.float32)
vertex = np.core.records.fromarrays(xyz_nxyz.T, names='x, y, z, nx, ny, nz, value, confidence',
formats='f4, f4, f4, f4, f4, f4, f4, f4')
el = plyfile.PlyElement.describe(vertex, 'vertex')
ply = plyfile.PlyData([el], text=text)
ply.write(out_filename)
# Replace \r\n with \n.
with open(out_filename, 'rb') as f:
content = f.read()
beginning = content[:128]
rest = content[128:]
beginning = beginning.replace(b'ply\r\n', b'ply\n')
with open(out_filename, 'wb') as f:
f.write(beginning + rest)
return ply
def load_model(filename, use_cpu=True):
assert filename.endswith('.pth'), filename
parts = path.basename(filename.split('.')[0]).split('_')
current_epoch = int(parts[1])
current_epoch_step = int(parts[2])
global_step = int(parts[3])
model = eval_utils.load_torch_model(filename, use_cpu=use_cpu)
log.info('Loading model {}'.format(
(current_epoch, current_epoch_step, global_step)))
assert 'metadata' in model
# if 'metadata' in model:
# assert model['metadata']['global_step'] == global_step
# assert model['metadata']['current_epoch'] == current_epoch
# assert model['metadata']['current_epoch_step'] == current_epoch_step
# else:
# model['metadata'] = {
# 'global_step': global_step,
# 'current_epoch': current_epoch,
# 'current_epoch_step': current_epoch_step,
# }
return model
def data_queue_worker(data_loader: dataset.ExampleLoader2,
data_queue: queue.Queue, models: dict,
discard_uneven_last_batch: bool):
"""
:param data_loader: Data source object. Implements next() and reset().
:param data_queue:
:param models: Only needed to access the metadata. Read-only.
:param discard_uneven_last_batch: Should be False for evaluation mode. Often true for training.
:return:
"""
batch_size = models['metadata']['batch_size']
while True:
next_batch = data_loader.next()
if next_batch is None or (discard_uneven_last_batch and
(len(next_batch[0]) != batch_size)):
log.info('Resetting data_loader')
data_loader.reset()
data_queue.put(None, block=True)
break
if models['metadata']['name'] == 'shapenetcore_rgb_mv6':
batch_data_np = prepare_data_rgb_mv(next_batch)
elif models['metadata']['name'] == 'shapenetcore_rgb_voxels':
batch_data_np = prepare_data_rgb_voxels(next_batch)
elif models['metadata']['name'] == 'shrec12_depth_mv6':
batch_data_np = prepare_data_depth_mv(next_batch)
elif models['metadata']['name'] == 'shrec12_voxels':
batch_data_np = prepare_data_depth_voxels(next_batch)
else:
log.error('ERROR: Not implemented')
sys.exit(1)
# move to gpu
batch_data = torch_utils.recursive_numpy_to_torch(batch_data_np,
cuda=True)
data_queue.put(batch_data, block=True)
log.info('End of data queue')
def open_locked(filename, mode='r+', sleep_seconds=0.2, verbose=True, **kwargs):
"""
If `filename` does not exist, `mode` cannot be "r" or "r+".
Parent directories are created if they do not exist.
"""
dirname = path.dirname(filename)
if 'r' in mode and not path.isfile(filename):
raise FileNotFoundError('{} does not exist. `mode` cannot be {}.'.format(filename, mode))
if not path.isdir(dirname):
log.info('mkdir -p {}'.format(dirname))
os.makedirs(dirname)
with open(filename, mode, **kwargs) as fd:
start_time = time.time()
try:
i = 0
while True:
try:
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
break
except IOError as e:
if e.errno != errno.EAGAIN:
raise e
else:
time.sleep(sleep_seconds)
if verbose and i > 0 and i % 10 == 0:
log.warning('Waiting for a locked file {}. Elapsed: {:.3g} seconds.'.format(filename,
time.time() - start_time))
i += 1
yield fd
fd.flush()
finally:
fcntl.flock(fd, fcntl.LOCK_UN)
def main():
# syn_images_dir = '/data/render_for_cnn/data/syn_images/'
# out_dir = '/tmp/hi'
syn_images_dir = '/data/syn_images/'
out_dir = '/data/syn_images_final/'
assert path.isdir(syn_images_dir)
bkg_dir = '/home/daeyun/git/RenderForCNN/datasets/sun2012pascalformat/JPEGImages'
bkg_file_list = '/home/daeyun/git/RenderForCNN/datasets/sun2012pascalformat/filelist.txt'
syn_images_filenames = glob.glob(path.join(syn_images_dir, '**/*.png'), recursive=True)
syn_images_filenames = sorted(syn_images_filenames)
with open(bkg_file_list, 'r') as f:
content = f.read()
fnames = content.strip().split()
background_filenames = [path.join(bkg_dir, name) for name in fnames]
assert path.isfile(syn_images_filenames[0])
assert path.isfile(background_filenames[0])
log.info('{} object images found'.format(len(syn_images_filenames)))
log.info('{} background images found'.format(len(background_filenames)))
start_time = time.time()
all_params = [
{
'object_image_filename': filename,
'bkg_filenames': background_filenames,
'out_image_filename': path.join(out_dir, '/'.join(filename.split('/')[-3:]))
} for filename in syn_images_filenames
]
pool = mp.Pool(processes=os.cpu_count() * 2)
for i, out_filename in enumerate(pool.imap(render_for_cnn_utils.blend_and_resize, all_params, chunksize=200)):
if i % 2000 == 0:
elapsed = time.time() - start_time
remaining = elapsed / (i + 1) * (len(all_params) - (i + 1))
log.info('{} of {}. Elapsed: {} min. Remaining: {} min'.format(i, len(all_params), int(elapsed / 60), int(remaining / 60)))
def truncate_images(syn_images_cropped_exact_dir,
out_dir,
ignore_overwrite=True):
log.info(
'looking for png files in {}'.format(syn_images_cropped_exact_dir))
assert os.path.isdir(syn_images_cropped_exact_dir)
filenames = sorted(
list(glob.glob(path.join(syn_images_cropped_exact_dir, '*/*/*.png'))))
assert len(filenames) > 0
log.info('found {} images'.format(len(filenames)))
# synset filtering
new_filenames = []
for filename in filenames:
m = re.search(
r'syn_images_cropped/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a',
filename)
synset = m.group(1)
if synset not in ['02858304', '02876657', '02924116', '04379243']:
continue
new_filenames.append(filename)
filenames = new_filenames
assert len(filenames) > 100
print('Number of images to resize:', len(filenames))
start_time = time.time()
params = [{
'filename': filename,
'out_dir': out_dir,
'ignore_overwrite': ignore_overwrite,
} for filename in filenames]
pool = mp.Pool()
for i, out_filename in enumerate(
pool.imap(_truncate_images_worker, params, chunksize=200)):
if i % 2000 == 0:
elapsed = time.time() - start_time
remaining = elapsed / (i + 1) * (len(filenames) - (i + 1))
log.info('{} of {}. Elapsed: {} min. Remaining: {} min'.format(
i, len(filenames), int(elapsed / 60), int(remaining / 60)))
def main():
syn_images_dir = '/data/mvshape/shapenetcore/single_rgb_128/'
shapenetcore_dir = '/data/shapenetcore/ShapeNetCore.v1/'
log.info('Getting all filenames.')
syndirs = sorted(glob.glob(path.join(syn_images_dir, '*')))
filenames = []
for syndir in syndirs:
modeldirs = sorted(glob.glob(path.join(syndir, '*')))
if is_subset:
modeldirs = modeldirs[:10]
for modeldir in modeldirs:
renderings = sorted(glob.glob(path.join(modeldir, '*.png')))
if is_subset:
renderings = renderings[:7]
filenames.extend(renderings)
# random.seed(42)
# if not is_subset:
# random.shuffle(filenames)
# filenames = filenames[:1000000]
random.seed(42)
log.info('{} files'.format(len(filenames)))
# TODO
target_dir = '/data/mvshape/database'
if is_subset:
sqlite_file_path = join(target_dir, 'shapenetcore_subset.sqlite')
else:
sqlite_file_path = join(target_dir, 'shapenetcore.sqlite')
output_cam_distance_from_origin = 2
log.info('Setting up output directory.')
# set up debugging directory.
if path.isfile(sqlite_file_path):
os.remove(sqlite_file_path)
io_utils.ensure_dir_exists(target_dir)
# used for making sure there is no duplicate.
duplicate_name_check_set = set()
log.info('Checking for duplicates. And making sure params.txt exists.')
for i, filename in enumerate(filenames):
m = re.search(r'single_rgb_128/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a',
filename)
synset = m.group(1)
model_name = m.group(2)
v = m.group(3)
image_num = int(v)
vc_rendering_name = '{}_{}_{:04d}'.format(synset, model_name,
image_num)
if vc_rendering_name in duplicate_name_check_set:
print('duplicate found: ', (filename, vc_rendering_name))
duplicate_name_check_set.add(vc_rendering_name)
params_filename = join(syn_images_dir, synset, model_name,
'params.txt')
assert path.isfile(params_filename)
# Create the database
dbm.init(sqlite_file_path)
with dbm.db.transaction() as txn:
log.info('Creating common objects.')
make_dataset('shapenetcore')
make_rendering_type('rgb')
make_rendering_type('depth')
make_rendering_type('normal')
make_rendering_type('voxels')
make_tag('novelview')
make_tag('novelmodel')
make_tag('novelclass')
make_tag('perspective_input')
make_tag('orthographic_input')
make_tag('perspective_output')
make_tag('orthographic_output')
make_tag('viewer_centered')
make_tag('object_centered')
make_tag('real_world')
make_split('train')
make_split('test')
make_split('validation')
# Quote from http://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v1/README.txt
# "The OBJ files have been pre-aligned so that the up direction is the +Y axis, and the front is the +X axis. In addition each model is normalized to fit within a unit cube centered at the origin."
oc_output_cam = camera.OrthographicCamera.from_Rt(
transforms.lookat_matrix(cam_xyz=(0, 0,
output_cam_distance_from_origin),
obj_xyz=(0, 0, 0),
up=(0, 1, 0)),
wh=(128, 128),
is_world_to_cam=True)
db_oc_output_cam = get_db_camera(oc_output_cam, fov=None)
# Prepare all category objects.
log.info('Preparing categories.')
synset_db_category_map = {}
for synset, synset_name in synset_name_pairs:
db_category_i, _ = dbm.Category.get_or_create(name=synset_name)
synset_db_category_map[synset] = db_category_i
txn.commit()
# Prepare all mesh model objects.
# ---------------------------------------------
db_object_map = {}
# model_name -> {rendering_type_name -> rendering}
db_object_centered_renderings = {}
log.info('Preparing mesh model objects.')
start_time = time.time()
count = 0
for i, filename in enumerate(filenames):
m = re.search(
r'single_rgb_128/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a',
filename)
synset = m.group(1)
model_name = m.group(2)
if model_name not in db_object_map:
mesh_filename = join(shapenetcore_dir, synset, model_name,
'model.obj')
assert path.isfile(mesh_filename)
mesh_filename_suffix = join(
'/mesh/shapenetcore/v1',
'/'.join(mesh_filename.split('/')[-3:]))
db_category = synset_db_category_map[synset]
# Must be unique.
db_object = dbm.Object.create(
name=model_name,
category=db_category,
dataset=datasets['shapenetcore'],
num_vertices=0, # Not needed for now. Easy to fill in later.
num_faces=0,
mesh_filename=mesh_filename_suffix,
)
db_object_map[model_name] = db_object
oc_rendering_name = '{}_{}'.format(synset, model_name)
assert model_name not in db_object_centered_renderings
db_object_centered_renderings[model_name] = {
'output_rgb':
dbm.ObjectRendering.create(
type=rendering_types['rgb'],
camera=db_oc_output_cam,
object=db_object,
# JPG
filename='/shapenetcore/mv20_rgb_128/{}.bin'.format(
oc_rendering_name),
resolution=128,
num_channels=3,
set_size=20,
is_normalized=False,
),
'output_depth':
dbm.ObjectRendering.create(
type=rendering_types['depth'],
camera=db_oc_output_cam,
object=db_object,
# Since there is only one gt rendering per model, their id is the same as the model name.
filename='/shapenetcore/mv20_depth_128/{}.bin'.format(
oc_rendering_name),
resolution=128,
num_channels=1,
set_size=20,
is_normalized=False,
),
'output_normal':
dbm.ObjectRendering.create(
type=rendering_types['normal'],
camera=db_oc_output_cam,
object=db_object,
filename='/shapenetcore/mv20_normal_128/{}.bin'.format(
oc_rendering_name),
resolution=128,
num_channels=3,
set_size=20,
is_normalized=False,
),
'output_voxels':
dbm.ObjectRendering.create(
type=rendering_types['voxels'],
camera=db_oc_output_cam,
object=db_object,
filename='/shapenetcore/voxels_32/{}.bin'.format(
oc_rendering_name),
resolution=32,
num_channels=1,
set_size=1,
is_normalized=False,
)
}
if count % 5000 == 0:
txn.commit()
t_elapsed = (time.time() - start_time)
t_remaining = (t_elapsed / (i + 1) * (len(filenames) - i))
log.info(
'Creating mesh objects in db. {} of {}. elapsed: {:.1f} min, remaining: {:.1f} min'
.format(i, len(filenames), t_elapsed / 60,
t_remaining / 60))
count += 1
txn.commit()
t_elapsed = time.time() - start_time
log.info('created {} mesh objects in db. elapsed: {:.1f} min'.format(
count, t_elapsed / 60))
start_time = time.time()
log.info('Processing rgb images.')
for i, filename in enumerate(filenames):
m = re.search(
r'single_rgb_128/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a',
filename)
synset = m.group(1)
model_name = m.group(2)
v = m.group(3)
image_num = int(v)
params_filename = join(syn_images_dir, synset, model_name,
'params.txt')
assert path.isfile(params_filename)
lines = render_for_cnn_utils.read_params_file(params_filename)
Rt = render_for_cnn_utils.get_Rt_from_RenderForCNN_parameters(
lines[image_num])
# Input and output cameras
# -------------------
input_cam = camera.OrthographicCamera.from_Rt(Rt,
wh=(128, 128),
is_world_to_cam=True)
# 49.1343 degrees is the default fov in blender.
db_input_cam = get_db_camera(input_cam, fov=49.1343)
input_cam_depth_xyz = (input_cam.pos /
la.norm(input_cam.pos)) * 1.5
input_cam_depth_Rt = transforms.lookat_matrix(
cam_xyz=input_cam_depth_xyz,
obj_xyz=(0, 0, 0),
up=input_cam.up_vector)
input_cam_depth = camera.OrthographicCamera.from_Rt(
input_cam_depth_Rt, wh=(128, 128), is_world_to_cam=True)
db_input_cam_depth = get_db_camera(input_cam_depth, fov=49.1343)
output_cam_xyz = (input_cam.pos / la.norm(
input_cam.pos)) * output_cam_distance_from_origin
output_Rt = transforms.lookat_matrix(cam_xyz=output_cam_xyz,
obj_xyz=(0, 0, 0),
up=input_cam.up_vector)
vc_output_cam = camera.OrthographicCamera.from_Rt(
output_Rt, wh=(128, 128), is_world_to_cam=True)
db_vc_output_cam = get_db_camera(vc_output_cam, fov=None)
# ---
db_object = db_object_map[model_name]
vc_rendering_name = '{}_{}_{:04d}'.format(synset, model_name,
image_num)
# Viewer centered renderings.
# --------------------------------
# Input rgb image:
db_object_rendering_input_rgb = dbm.ObjectRendering.create(
type=rendering_types['rgb'],
camera=db_input_cam,
object=db_object,
# This should already exist.
filename='/shapenetcore/single_rgb_128/{}.png'.format(
vc_rendering_name),
resolution=128,
num_channels=1,
set_size=1,
is_normalized=False, # False for rgb.
)
db_object_rendering_input_depth = dbm.ObjectRendering.create(
type=rendering_types['depth'],
camera=db_input_cam_depth,
object=db_object,
filename='/shapenetcore/single_depth_128/{}.bin'.format(
vc_rendering_name),
resolution=128,
num_channels=1,
set_size=1,
is_normalized=True,
)
db_object_rendering_vc_output_rgb = dbm.ObjectRendering.create(
type=rendering_types['rgb'],
camera=db_vc_output_cam,
object=db_object,
filename='/shapenetcore/mv20_rgb_128/{}.bin'.format(
vc_rendering_name),
resolution=128,
num_channels=3,
set_size=20,
is_normalized=False,
)
db_object_rendering_vc_output_depth = dbm.ObjectRendering.create(
type=rendering_types['depth'],
camera=db_vc_output_cam,
object=db_object,
filename='/shapenetcore/mv20_depth_128/{}.bin'.format(
vc_rendering_name),
resolution=128,
num_channels=1,
set_size=20,
is_normalized=False,
)
db_object_rendering_vc_output_normal = dbm.ObjectRendering.create(
type=rendering_types['normal'],
camera=db_vc_output_cam,
object=db_object,
filename='/shapenetcore/mv20_normal_128/{}.bin'.format(
vc_rendering_name),
resolution=128,
num_channels=3,
set_size=20,
is_normalized=False,
)
db_object_rendering_vc_output_voxels = dbm.ObjectRendering.create(
type=rendering_types['voxels'],
camera=db_vc_output_cam,
object=db_object,
filename='/shapenetcore/voxels_32/{}.bin'.format(
vc_rendering_name),
resolution=32,
num_channels=1,
set_size=1,
is_normalized=False,
)
# Examples
# ----------------
# A row in the `Example` table is just an id for many-to-many references.
# View centered
example_viewer_centered = dbm.Example.create()
dbm.ExampleObjectRendering.create(
example=example_viewer_centered,
rendering=db_object_rendering_input_rgb)
dbm.ExampleObjectRendering.create(
example=example_viewer_centered,
rendering=db_object_rendering_input_depth)
dbm.ExampleObjectRendering.create(
example=example_viewer_centered,
rendering=db_object_rendering_vc_output_depth)
dbm.ExampleObjectRendering.create(
example=example_viewer_centered,
rendering=db_object_rendering_vc_output_normal)
dbm.ExampleObjectRendering.create(
example=example_viewer_centered,
rendering=db_object_rendering_vc_output_rgb)
dbm.ExampleObjectRendering.create(
example=example_viewer_centered,
rendering=db_object_rendering_vc_output_voxels)
dbm.ExampleDataset.create(example=example_viewer_centered,
dataset=datasets['shapenetcore'])
dbm.ExampleSplit.create(example=example_viewer_centered,
split=splits['train'])
dbm.ExampleTag.create(example=example_viewer_centered,
tag=tags['real_world'])
dbm.ExampleTag.create(example=example_viewer_centered,
tag=tags['viewer_centered'])
dbm.ExampleTag.create(example=example_viewer_centered,
tag=tags['perspective_input'])
dbm.ExampleTag.create(example=example_viewer_centered,
tag=tags['orthographic_output'])
dbm.ExampleTag.create(example=example_viewer_centered,
tag=tags['novelmodel'])
# Object centered
example_object_centered = dbm.Example.create()
dbm.ExampleObjectRendering.create(
example=example_object_centered,
rendering=db_object_rendering_input_rgb)
dbm.ExampleObjectRendering.create(
example=example_object_centered,
rendering=db_object_rendering_input_depth)
dbm.ExampleObjectRendering.create(
example=example_object_centered,
rendering=db_object_centered_renderings[model_name]
['output_depth'])
dbm.ExampleObjectRendering.create(
example=example_object_centered,
rendering=db_object_centered_renderings[model_name]
['output_normal'])
dbm.ExampleObjectRendering.create(
example=example_object_centered,
rendering=db_object_centered_renderings[model_name]
['output_rgb'])
dbm.ExampleObjectRendering.create(
example=example_object_centered,
rendering=db_object_centered_renderings[model_name]
['output_voxels'])
dbm.ExampleDataset.create(example=example_object_centered,
dataset=datasets['shapenetcore'])
dbm.ExampleSplit.create(example=example_object_centered,
split=splits['train'])
dbm.ExampleTag.create(example=example_object_centered,
tag=tags['real_world'])
dbm.ExampleTag.create(example=example_object_centered,
tag=tags['object_centered'])
dbm.ExampleTag.create(example=example_object_centered,
tag=tags['perspective_input'])
dbm.ExampleTag.create(example=example_object_centered,
tag=tags['orthographic_output'])
dbm.ExampleTag.create(example=example_object_centered,
tag=tags['novelmodel'])
if i % 5000 == 0:
txn.commit()
t_elapsed = (time.time() - start_time)
t_remaining = (t_elapsed / (i + 1) * (len(filenames) - i))
log.info(
'Creating examples in db. {} of {}. elapsed: {:.1f} min, remaining: {:.1f} min'
.format(i, len(filenames), t_elapsed / 60,
t_remaining / 60))
txn.commit()
dbm.db.commit()
t_elapsed = (time.time() - start_time)
log.info('total elapsed: {:.1f} min'.format(t_elapsed / 60))
def main():
experiment_name = args.experiment.strip()
save_name = args.save_name
model_name = args.model_name
resume_filename = path.expanduser(args.resume)
print_every = args.print_every
# sanity checks
assert experiment_name in ('opo', 'vpo'), experiment_name
assert save_name, save_name
assert model_name in encoderdecoder.model_names, model_name
assert print_every > 0
learning_rate = get_learning_rate(model_name)
log.info(args)
if resume_filename:
assert resume_filename.endswith('.pth')
assert path.isfile(resume_filename)
assert '/' + experiment_name in resume_filename # sanity check. doesnt necessarily have to be true
# e.g. '/data/mvshape/out/pytorch/shapenetcore_rgb_mv6/vpo/0/models0_00000_0017500_00017500.pth'
models = load_model(resume_filename, model_name)
else:
models = build_model(model_name)
# models['metadata']['batch_size']
batch_size = models['metadata']['batch_size']
# output of model is
# ((b, 1, 128, 128), (b, 1, 128, 128))
# TODO: this should probably be in gpu
helper_torch_modules = encoderdecoder.build_helper_torch_modules()
all_params = []
torch_utils.recursive_module_apply(
models, lambda m: all_params.extend(m.parameters()))
assert len(all_params) > 0
optimizer = optim.Adam(all_params, lr=learning_rate)
log.info('Model is ready. Loading training metadata.')
# Initialize dataset.
#####################################
loader = select_data_loader(model_name=model_name,
experiment_name=experiment_name,
batch_size=batch_size)
log.info('Data is ready.')
save_dir = path.join(
dataset.base,
'out/pytorch/{}/{}/{}/'.format(model_name, experiment_name, save_name))
# Main loop
while True:
encoderdecoder.train_epoch(
models=models,
data_loader=loader,
optimizer=optimizer,
helper_torch_modules=helper_torch_modules,
save_every=2500,
print_every=print_every,
save_dir=save_dir,
)
def train_epoch(models,
data_loader: dataset.ExampleLoader2,
optimizer: torch.optim.Optimizer,
helper_torch_modules: dict,
save_every=2500,
print_every=50,
save_dir=None):
log.info('--- Started training. \n{}'.format(
pprint.pformat(models['metadata'])))
torch_utils.recursive_train_setter(models, is_training=True)
gc.collect()
batch_size = models['metadata']['batch_size']
num_all_examples = data_loader.total_num_examples
assert batch_size == data_loader.batch_size, (batch_size,
data_loader.batch_size)
# sanity check. not really necessary. just making sure .reset() was called.
assert data_loader.current_batch_index == 0, data_loader.current_batch_index
current_num_examples = 0
all_scalar_outputs = {}
data_queue = queue.Queue(maxsize=1)
thread = threading.Thread(target=data_queue_worker,
args=[data_loader, data_queue, models, True],
daemon=True)
thread.start()
io_block_times = []
model_name = models['metadata']['name']
models['metadata']['current_epoch'] += 1
log.info('\n### New epoch: {}\n'.format(
models['metadata']['current_epoch']))
models['metadata']['current_epoch_step'] = 0
while True:
start_time = time.time()
batch_data = data_queue.get(block=True)
io_block_times.append(time.time() - start_time)
if models['metadata']['current_epoch_step'] > 0:
if models['metadata']['current_epoch_step'] % print_every == 0:
print_status(models,
num_all_examples,
io_block_times,
accum_scalars=all_scalar_outputs)
if save_dir is not None and models['metadata'][
'current_epoch_step'] % save_every == 0:
# save_model(models, save_dir=path.join(dataset.base, 'out/pytorch/{}/{}/{}/'.format(model_name, experiment_name, save_name)))
save_model(models, save_dir=save_dir)
# Check if end of epoch.
if batch_data is None:
print_status(models,
num_all_examples,
io_block_times,
accum_scalars=all_scalar_outputs)
if save_dir is not None:
save_model(models, save_dir=save_dir)
break
assert batch_data['in_image'].size(0) == batch_size
optimizer.zero_grad()
model_output = get_output_from_model(
models,
batch_data=batch_data,
helper_torch_modules=helper_torch_modules,
include_output_tensors=False)
total_loss = model_output['loss']['total']
# modifies `all_scalar_outputs` in-place
torch_utils.recursive_scalar_dict_merge(
{
'loss': model_output['loss'],
'metric': model_output['metric'],
}, all_scalar_outputs)
# calling `.backward()` seems to trigger garbage collection. without it, it runs out of memory during testing.
# this doesn't update any parameters, so it's safe to run when `is_training` is false too.
total_loss.backward()
optimizer.step()
# running `to_python_scalar` triggers a sync point.
# total_loss_sum += (to_python_scalar(total_loss) * batch_data['in_image'].size(0)) / this_dataset_total_num_examples
current_num_examples += batch_size
models['metadata']['current_epoch_step'] += 1
models['metadata']['global_step'] += 1
print('.', end='', flush=True)
log.info('\nTrained {} examples\n\n'.format(num_all_examples))
return torch_utils.recursive_dict_of_lists_mean(all_scalar_outputs)
def select_data_loader(model_name, experiment_name, batch_size):
if model_name == 'shapenetcore_rgb_mv6':
if experiment_name == 'vpo':
# loader = dataset.ExampleLoader2(path.join(dataset.base, 'out/splits/shapenetcore_examples_{}/subset_examples.cbor'.format(experiment_name)),
# tensors_to_read=('input_rgb', 'target_depth'), shuffle=True, batch_size=batch_size)
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shapenetcore_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_rgb',
'target_depth'),
shuffle=True,
batch_size=batch_size)
elif experiment_name == 'opo':
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shapenetcore_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_rgb',
'target_depth'),
shuffle=True,
batch_size=batch_size)
else:
raise NotImplementedError()
elif model_name == 'shapenetcore_rgb_voxels':
if experiment_name == 'vpo':
# TODO#########################
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shapenetcore_examples_{}/subset_examples.cbor'.
format(experiment_name)),
tensors_to_read=('input_rgb',
'target_voxels'),
shuffle=True,
batch_size=batch_size)
elif experiment_name == 'opo':
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shapenetcore_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_rgb',
'target_voxels'),
shuffle=True,
batch_size=batch_size)
else:
raise NotImplementedError()
elif model_name == 'shrec12_depth_mv6':
if experiment_name == 'vpo':
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shrec12_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_depth',
'target_depth'),
shuffle=True,
batch_size=batch_size)
elif experiment_name == 'opo':
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shrec12_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_depth',
'target_depth'),
shuffle=True,
batch_size=batch_size)
else:
raise NotImplementedError()
elif model_name == 'shrec12_voxels':
if experiment_name == 'vpo':
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shrec12_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_depth',
'target_voxels'),
shuffle=True,
batch_size=batch_size)
elif experiment_name == 'opo':
loader = dataset.ExampleLoader2(path.join(
dataset.base,
'out/splits/shrec12_examples_{}/all_examples.cbor'.format(
experiment_name)),
tensors_to_read=('input_depth',
'target_voxels'),
shuffle=True,
batch_size=batch_size)
else:
raise NotImplementedError()
else:
raise NotImplementedError()
total_num_examples = loader.total_num_examples
log.info('Dataset size: {}'.format(total_num_examples))
return loader
l.depth_and_normal_map.restype = ctypes.c_int
l.depth_and_normal_map.argtypes = [
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ctypes.c_size_t,
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ctypes.c_int,
ctypes.c_int,
ctypes.c_char_p,
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
]
lib = l
log.info('Loaded shared library {}'.format(lib._name))
def depth_and_normal_map(vertices, eye, center, up, wh):
tri = vertices.astype(np.float32)
eye = np.array(eye, np.float32)
center = np.array(center, np.float32)
up = np.array(up, np.float32)
rendered_depth = np.zeros([wh[0], wh[1]], dtype=np.float32)
rendered_normal = np.zeros([wh[0], wh[1], 3], dtype=np.float32)
resources_dir = path.realpath(path.join(proj_root, 'resources'))
lib.depth_and_normal_map(tri, tri.size, eye, center, up, 128, 128,
ctypes.c_char_p(resources_dir.encode('ascii')),
rendered_depth, rendered_normal)
