def vis_voxels(model_id, edge_length_threshold, filled, shuffle=False):
from mayavi import mlab
from util3d.mayavi_vis import vis_voxels
from shapenet.core import cat_desc_to_id
from template_ffd.inference.voxels import get_voxel_dataset
from template_ffd.data.voxels import get_gt_voxel_dataset
from template_ffd.model import load_params
from template_ffd.data.ids import get_example_ids
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
gt_ds = get_gt_voxel_dataset(cat_id, filled)
inf_ds = get_voxel_dataset(model_id, edge_length_threshold)
example_ids = get_example_ids(cat_id, 'eval')
if shuffle:
example_ids = list(example_ids)
example_ids.shuffle
with gt_ds:
with inf_ds:
for example_id in example_ids:
gt = gt_ds[example_id].data
inf = inf_ds[example_id].data
vis_voxels(gt, color=(0, 0, 1))
mlab.figure()
vis_voxels(inf, color=(0, 1, 0))
mlab.show()
def vis():
def vis_mesh(mesh, include_wireframe=False, **kwargs):
from util3d.mayavi_vis import vis_mesh as vm
v, f = (np.array(mesh[k]) for k in ('vertices', 'faces'))
vm(v, f, include_wireframe=include_wireframe, **kwargs)
example_ids = list(get_example_ids(cat_id, 'eval'))
random.shuffle(example_ids)
with all_ds:
for example_id in example_ids:
print(example_id)
image, gt_mesh, predictions = all_ds[example_id]
meshes = top_k_mesh_fn(*(np.array(predictions[k])
for k in ('probs', 'dp')))
plt.imshow(image)
mlab.figure()
vis_mesh(gt_mesh, color=(0, 0, 1))
for mesh in meshes:
v, f, ov = (mesh[k] for k in ('vertices', 'faces',
'original_vertices'))
mlab.figure()
vis_mesh({'vertices': v, 'faces': f}, color=(0, 1, 0))
mlab.figure()
vis_mesh({'vertices': ov, 'faces': f}, color=(1, 0, 0))
plt.show(block=False)
mlab.show()
plt.close()
def _get_lazy_evaluation_dataset_single(inf_cloud_ds, cat_id, n_samples,
eval_fn):
def sample_fn(cloud):
return sample_points(np.array(cloud), n_samples)
example_ids = get_example_ids(cat_id, 'eval')
normalization_ds = get_normalization_params_dataset(cat_id)
gt_cloud_ds = get_point_cloud_dataset(
cat_id, n_samples, example_ids=example_ids).map(sample_fn)
with inf_cloud_ds:
keys = tuple(inf_cloud_ds.keys())
normalization_ds = normalization_ds.map_keys(lambda key: key[:2])
gt_cloud_ds = gt_cloud_ds.map_keys(lambda key: key[:2])
zipped = Dataset.zip(inf_cloud_ds, gt_cloud_ds,
normalization_ds).subset(keys, check_present=False)
def map_fn(data):
inf_cloud, gt_cloud, norm_params = data
inf_cloud = normalized(inf_cloud, **norm_params)
gt_cloud = normalized(gt_cloud, **norm_params)
return eval_fn(inf_cloud, gt_cloud)
dataset = zipped.map(map_fn)
return dataset
def vis_voxels(model_id, edge_length_threshold, filled, shuffle=False):
from mayavi import mlab
from util3d.mayavi_vis import vis_voxels
from shapenet.core import cat_desc_to_id
from template_ffd.inference.voxels import get_voxel_dataset
from template_ffd.data.voxels import get_gt_voxel_dataset
from template_ffd.model import load_params
from template_ffd.data.ids import get_example_ids
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
gt_ds = get_gt_voxel_dataset(cat_id, filled)
inf_ds = get_voxel_dataset(model_id, edge_length_threshold)
example_ids = get_example_ids(cat_id, 'eval')
if shuffle:
example_ids = list(example_ids)
example_ids.shuffle
with gt_ds:
with inf_ds:
for example_id in example_ids:
gt = gt_ds[example_id].data
inf = inf_ds[example_id].data
vis_voxels(gt, color=(0, 0, 1))
mlab.figure()
vis_voxels(inf, color=(0, 1, 0))
mlab.show()
def vis():
def vis_mesh(mesh, include_wireframe=False, **kwargs):
from util3d.mayavi_vis import vis_mesh as vm
v, f = (np.array(mesh[k]) for k in ('vertices', 'faces'))
vm(v, f, include_wireframe=include_wireframe, **kwargs)
example_ids = list(get_example_ids(cat_id, 'eval'))
random.shuffle(example_ids)
with all_ds:
for example_id in example_ids:
print(example_id)
image, gt_mesh, predictions = all_ds[example_id]
meshes = top_k_mesh_fn(
*(np.array(predictions[k]) for k in ('probs', 'dp')))
plt.imshow(image)
mlab.figure()
vis_mesh(gt_mesh, color=(0, 0, 1))
for mesh in meshes:
v, f, ov = (mesh[k] for k in
('vertices', 'faces', 'original_vertices'))
mlab.figure()
vis_mesh({'vertices': v, 'faces': f}, color=(0, 1, 0))
mlab.figure()
vis_mesh({'vertices': ov, 'faces': f}, color=(1, 0, 0))
plt.show(block=False)
mlab.show()
plt.close()
def get_predictions_data(model_id):
builder = get_builder(model_id)
cat_id = builder.cat_id
mode = 'infer'
example_ids = get_example_ids(cat_id, mode)
estimator = builder.get_estimator()
predictions = estimator.predict(builder.get_predict_inputs)
return LengthedGenerator(predictions, len(example_ids))
def get_dataset(self, mode, repeat=None):
cat_id = self.cat_id
if isinstance(cat_id, (list, tuple)):
example_ids = [get_example_ids(c, mode) for c in cat_id]
else:
example_ids = get_example_ids(cat_id, mode)
render_config = self.render_config
view_index = self.view_index
n_samples = self.n_samples
n_resamples = self.params.get('n_resamples', 1024)
if repeat is None:
repeat = mode == tf.estimator.ModeKeys.TRAIN
shuffle = repeat
batch_size = self.batch_size
dataset = get_dataset(
render_config, view_index, n_samples, n_resamples, cat_id,
example_ids, shuffle=shuffle, repeat=repeat, batch_size=batch_size)
return dataset
def get_lazy_dataset(self):
from shapenet.core.meshes import get_mesh_dataset
example_ids = get_example_ids(self._cat_id, 'eval')
mesh_ds = get_mesh_dataset(self._cat_id).subset(example_ids)
def map_fn(mesh):
vertices = mesh['vertices']
offset, scale_factor = get_normalization_params(vertices)
return dict(offset=[float(o) for o in offset],
scale_factor=float(scale_factor))
return mesh_ds.map(map_fn)
def get_predictions_data(model_id, mode='infer'):
builder = get_builder(model_id)
cat_id = builder.cat_id
example_ids = get_example_ids(cat_id, mode)
n = len(example_ids)
view_index = builder.view_index
if isinstance(view_index, (list, tuple)):
n *= len(view_index)
estimator = builder.get_estimator()
predictions = estimator.predict(builder.get_predict_inputs)
return LengthedGenerator(predictions, n)
def get_lazy_dataset(self):
from shapenet.core.meshes import get_mesh_dataset
example_ids = get_example_ids(self._cat_id, 'eval')
mesh_ds = get_mesh_dataset(self._cat_id).subset(example_ids)
def map_fn(mesh):
vertices = mesh['vertices']
offset, scale_factor = get_normalization_params(vertices)
return dict(
offset=[float(o) for o in offset],
scale_factor=float(scale_factor))
return mesh_ds.map(map_fn)
def vis_clouds(model_id,
pre_sampled=True,
n_samples=1024,
edge_length_threshold=0.1,
shuffle=False):
import random
import numpy as np
from mayavi import mlab
import matplotlib.pyplot as plt
from dids import Dataset
from shapenet.core.blender_renderings.config import RenderConfig
from shapenet.core.meshes import get_mesh_dataset
from util3d.mayavi_vis import vis_point_cloud
from util3d.mayavi_vis import vis_mesh
from template_ffd.data.ids import get_example_ids
from template_ffd.inference.clouds import get_inferred_cloud_dataset
from template_ffd.model import get_builder
builder = get_builder(model_id)
cat_id = builder.cat_id
kwargs = dict(model_id=model_id, n_samples=n_samples)
if not pre_sampled:
kwargs['edge_length_threshold'] = edge_length_threshold
cloud_dataset = get_inferred_cloud_dataset(pre_sampled=pre_sampled,
**kwargs)
image_dataset = RenderConfig().get_dataset(cat_id, builder.view_index)
example_ids = get_example_ids(cat_id, 'eval')
if shuffle:
example_ids = list(example_ids)
random.shuffle(example_ids)
mesh_dataset = get_mesh_dataset(cat_id)
zipped_dataset = Dataset.zip(image_dataset, cloud_dataset, mesh_dataset)
# zipped_dataset = Dataset.zip(image_dataset, cloud_dataset)
with zipped_dataset:
for example_id in example_ids:
image, cloud, mesh = zipped_dataset[example_id]
# image, cloud = zipped_dataset[example_id]
plt.imshow(image)
vis_point_cloud(np.array(cloud),
color=(0, 1, 0),
scale_factor=0.01)
v, f = (np.array(mesh[k]) for k in ('vertices', 'faces'))
vis_mesh(v,
f,
color=(0, 0, 1),
opacity=0.1,
include_wireframe=False)
plt.show(block=False)
mlab.show()
plt.close()
def get_iou_dataset(model_id, edge_length_threshold=0.1, filled=False):
from shapenet.core import cat_desc_to_id
from template_ffd.data.ids import get_example_ids
from template_ffd.model import load_params
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
n_eval = len(get_example_ids(cat_id, 'eval'))
manager = IouAutoSavingManager(model_id=model_id,
edge_length_threshold=edge_length_threshold,
filled=filled)
with manager.get_saving_dataset() as ds:
needs_calc = len(ds) < n_eval
if needs_calc:
manager.save_all()
return manager.get_saving_dataset()
def get_dataset(self, mode):
cat_id = self.cat_id
example_ids = get_example_ids(cat_id, mode)
render_config = self.render_config
view_index = self.view_index
n_samples = self.n_samples
n_resamples = self.params.get('n_resamples', 1024)
repeat = mode == tf.estimator.ModeKeys.TRAIN
shuffle = repeat
batch_size = self.batch_size
dataset = get_dataset(
render_config, view_index, n_samples, n_resamples, cat_id,
example_ids, shuffle=shuffle, repeat=repeat, batch_size=batch_size)
return dataset
def report_emd_average(model_id, pre_sampled=True, **kwargs):
from shapenet.core import cat_desc_to_id
from template_ffd.data.ids import get_example_ids
from template_ffd.model import load_params
manager = get_emd_manager(model_id, pre_sampled, **kwargs)
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
n_eval = len(get_example_ids(cat_id, 'eval'))
with manager.get_saving_dataset('r') as ds:
if len(ds) == n_eval:
values = np.array([ds.values()])
else:
values = None
if values is None:
with manager.get_saved_dataset() as ds:
values = np.array([ds.values()])
print(np.mean(values))
def get_lazy_dataset(self):
cat_id = cat_desc_to_id(load_params(self._model_id)['cat_desc'])
example_ids = get_example_ids(cat_id, 'eval')
inferred_dataset = get_voxel_dataset(
self._model_id, self._edge_length_threshold, self._voxel_config,
filled=self._filled, example_ids=example_ids)
gt_dataset = get_gt_voxel_dataset(cat_id, filled=self._filled)
voxel_datasets = Dataset.zip(inferred_dataset, gt_dataset)
voxel_datasets = voxel_datasets.subset(example_ids)
def map_fn(v):
return intersection_over_union(v[0].data, v[1].data)
iou_dataset = voxel_datasets.map(map_fn)
return iou_dataset
def get_lazy_evaluation_dataset(inf_cloud_ds, cat_id, n_samples, eval_fn):
example_ids = get_example_ids(cat_id, 'eval')
def sample_fn(cloud):
return sample_points(np.array(cloud), n_samples)
normalization_ds = get_normalization_params_dataset(cat_id)
gt_cloud_ds = get_point_cloud_dataset(
cat_id, n_samples, example_ids=example_ids).map(sample_fn)
zipped = Dataset.zip(inf_cloud_ds, gt_cloud_ds, normalization_ds)
def map_fn(data):
inf_cloud, gt_cloud, norm_params = data
inf_cloud = normalized(inf_cloud, **norm_params)
gt_cloud = normalized(gt_cloud, **norm_params)
return eval_fn(inf_cloud, gt_cloud)
return zipped.map(map_fn)
def report_chamfer_average(model_id, pre_sampled=True, **kwargs):
import os
from shapenet.core import cat_desc_to_id
from template_ffd.data.ids import get_example_ids
from template_ffd.model import load_params
manager = get_chamfer_manager(model_id, pre_sampled, **kwargs)
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
n_eval = len(get_example_ids(cat_id, 'eval'))
values = None
if os.path.isfile(manager.path):
with manager.get_saving_dataset('r') as ds:
if len(ds) == n_eval:
values = np.array(tuple(ds.values()))
if values is None:
manager.save_all()
with manager.get_saving_dataset('r') as ds:
values = np.array(tuple(ds.values()))
print(np.mean(values))
def report_emd_average(model_id, pre_sampled=True, **kwargs):
import os
from shapenet.core import cat_desc_to_id
from template_ffd.data.ids import get_example_ids
from template_ffd.model import load_params
manager = get_emd_manager(model_id, pre_sampled, **kwargs)
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
n_eval = len(get_example_ids(cat_id, 'eval'))
values = None
if os.path.isfile(manager.path):
with manager.get_saving_dataset('r') as ds:
if len(ds) == n_eval:
values = np.array(tuple(ds.values()))
if values is None:
manager.save_all()
with manager.get_saving_dataset('r') as ds:
values = np.array(tuple(ds.values()))
print(np.mean(values))
def vis_clouds(
model_id, pre_sampled=True, n_samples=1024, edge_length_threshold=0.1,
shuffle=False):
import random
import numpy as np
from mayavi import mlab
import matplotlib.pyplot as plt
from dids import Dataset
from shapenet.core.blender_renderings.config import RenderConfig
from shapenet.core.meshes import get_mesh_dataset
from util3d.mayavi_vis import vis_point_cloud
from util3d.mayavi_vis import vis_mesh
from template_ffd.data.ids import get_example_ids
from template_ffd.inference.clouds import get_inferred_cloud_dataset
from template_ffd.model import get_builder
builder = get_builder(model_id)
cat_id = builder.cat_id
kwargs = dict(model_id=model_id, n_samples=n_samples)
if not pre_sampled:
kwargs['edge_length_threshold'] = edge_length_threshold
cloud_dataset = get_inferred_cloud_dataset(
pre_sampled=pre_sampled, **kwargs)
image_dataset = RenderConfig().get_dataset(cat_id, builder.view_index)
example_ids = get_example_ids(cat_id, 'eval')
if shuffle:
example_ids = list(example_ids)
random.shuffle(example_ids)
mesh_dataset = get_mesh_dataset(cat_id)
zipped_dataset = Dataset.zip(image_dataset, cloud_dataset, mesh_dataset)
# zipped_dataset = Dataset.zip(image_dataset, cloud_dataset)
with zipped_dataset:
for example_id in example_ids:
image, cloud, mesh = zipped_dataset[example_id]
# image, cloud = zipped_dataset[example_id]
plt.imshow(image)
vis_point_cloud(
np.array(cloud), color=(0, 1, 0), scale_factor=0.01)
v, f = (np.array(mesh[k]) for k in ('vertices', 'faces'))
vis_mesh(
v, f, color=(0, 0, 1), opacity=0.1, include_wireframe=False)
plt.show(block=False)
mlab.show()
plt.close()
def get_lazy_dataset(self):
cat_id = cat_desc_to_id(load_params(self._model_id)['cat_desc'])
example_ids = get_example_ids(cat_id, 'eval')
inferred_dataset = get_voxel_dataset(self._model_id,
self._edge_length_threshold,
self._voxel_config,
filled=self._filled,
example_ids=example_ids)
gt_dataset = get_gt_voxel_dataset(cat_id, filled=self._filled)
voxel_datasets = Dataset.zip(inferred_dataset, gt_dataset)
voxel_datasets = voxel_datasets.subset(example_ids)
def map_fn(v):
return intersection_over_union(v[0].data, v[1].data)
iou_dataset = voxel_datasets.map(map_fn)
return iou_dataset
def get_iou_dataset(
model_id, edge_length_threshold=0.1, filled=False, recalc=False):
from shapenet.core import cat_desc_to_id
from template_ffd.data.ids import get_example_ids
from template_ffd.model import load_params
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
n_eval = len(get_example_ids(cat_id, 'eval'))
manager = IouAutoSavingManager(
model_id=model_id,
edge_length_threshold=edge_length_threshold,
filled=filled
)
if not recalc:
with manager.get_saving_dataset() as ds:
recalc = len(ds) < n_eval
if recalc:
manager.save_all()
return manager.get_saving_dataset('r')
def _create_unfilled_voxel_data(
model_id, edge_length_threshold=0.1, voxel_config=None,
overwrite=False, example_ids=None):
from template_ffd.data.ids import get_example_ids
from shapenet.core import cat_desc_to_id
from template_ffd.model import load_params
import numpy as np
from progress.bar import IncrementalBar
if voxel_config is None:
voxel_config = _default_config
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
if example_ids is None:
example_ids = get_example_ids(cat_id, 'eval')
mesh_dataset = get_inferred_mesh_dataset(model_id, edge_length_threshold)
voxel_dataset = get_voxel_dataset(
model_id, edge_length_threshold, voxel_config, filled=False,
auto_save=False)
if not overwrite:
example_ids = [i for i in example_ids if i not in voxel_dataset]
if len(example_ids) == 0:
return
print('Creating %d voxels for model %s' % (len(example_ids), model_id))
kwargs = dict(
voxel_dim=voxel_config.voxel_dim,
exact=voxel_config.exact,
dc=voxel_config.dc,
aw=voxel_config.aw)
with mesh_dataset:
bar = IncrementalBar(max=len(example_ids))
for example_id in example_ids:
bar.next()
mesh = mesh_dataset[example_id]
vertices, faces = (
np.array(mesh[k]) for k in ('vertices', 'faces'))
binvox_path = voxel_dataset.path(example_id)
# x, z, y = vertices.T
# vertices = np.stack([x, y, z], axis=1)
bio.mesh_to_binvox(
vertices, faces, binvox_path, **kwargs)
bar.finish()
def test(resourceid,keyname):
img_inp,x,pt_gt,training_flag,loss,optimizer,batchno,batchnoinc,mindist,loss_nodecay,dists_forward,dists_backward,dist0=build_graph(resourceid,"test")
config=tf.ConfigProto()
config.gpu_options.allow_growth=True
config.allow_soft_placement=True
saver=tf.train.Saver()
cat_id=cat_desc_to_id('plane')
length=len(get_example_ids(cat_id,'eval'))
fout = open("%s/%s.v.txt"%(dumpdir,keyname),'wb')
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess,"%s/%s.ckpt"%(dumpdir,keyname))
values=[]
for i in xrange(0,length):
_,pred,total_loss,_,forward,backward,_=sess.run([batchnoinc,x,loss,loss_nodecay,dists_forward,dists_backward,dist0],
feed_dict={training_flag:False})
print i,total_loss
values.append(total_loss)
print np.mean(values)
def vis(cat_desc, regime='e', shuffle=True):
import matplotlib.pyplot as plt
from mayavi import mlab
from util3d.mayavi_vis import vis_point_cloud, vis_voxels
all_ds = get_ds(cat_desc, regime)
cat_id = cat_desc_to_id(cat_desc)
example_ids = list(get_example_ids(cat_id, 'eval'))
random.shuffle(example_ids)
def vis_mesh(mesh, include_wireframe=False, **kwargs):
from util3d.mayavi_vis import vis_mesh as vm
v, f = (np.array(mesh[k]) for k in ('vertices', 'faces'))
vm(v, f, include_wireframe=include_wireframe, **kwargs)
with all_ds:
for example_id in example_ids:
print(example_id)
image, gt_mesh, cloud, mesh, voxels, template_mesh = \
all_ds[example_id]
plt.imshow(image)
mlab.figure()
vis_mesh(gt_mesh, color=(0, 0, 1))
mlab.figure()
vis_mesh(mesh, color=(0, 1, 0))
mlab.figure()
vis_mesh(template_mesh, color=(1, 0, 0))
mlab.figure()
vis_point_cloud(np.array(cloud),
scale_factor=0.01,
color=(0, 1, 0))
mlab.figure()
vis_voxels(voxels.data, color=(0, 1, 0))
plt.show(block=False)
mlab.show()
plt.close()
def check_predictions(model_id):
from template_ffd.inference.predictions import get_predictions_dataset
from template_ffd.model import get_builder
from template_ffd.data.ids import get_example_ids
builder = get_builder(model_id)
cat_id = builder.cat_id
example_ids = get_example_ids(cat_id, 'eval')
missing = []
with get_predictions_dataset(model_id, 'r') as dataset:
for example_id in example_ids:
if example_id not in dataset:
missing.append(example_id)
else:
example = dataset[example_id]
if not all(k in example for k in ('probs', 'dp')):
missing.append(example_id)
if len(missing) == 0:
print('No predictions missing!')
else:
print('%d / %d predictions missing' % (len(missing), len(example_ids)))
for example_id in example_ids:
print(example_id)
def vis(cat_desc, regime='e', shuffle=True):
import matplotlib.pyplot as plt
from mayavi import mlab
from util3d.mayavi_vis import vis_point_cloud, vis_voxels
all_ds = get_ds(cat_desc, regime)
cat_id = cat_desc_to_id(cat_desc)
example_ids = list(get_example_ids(cat_id, 'eval'))
random.shuffle(example_ids)
def vis_mesh(mesh, include_wireframe=False, **kwargs):
from util3d.mayavi_vis import vis_mesh as vm
v, f = (np.array(mesh[k]) for k in ('vertices', 'faces'))
vm(v, f, include_wireframe=include_wireframe, **kwargs)
with all_ds:
for example_id in example_ids:
print(example_id)
image, gt_mesh, cloud, mesh, voxels, template_mesh = \
all_ds[example_id]
plt.imshow(image)
mlab.figure()
vis_mesh(gt_mesh, color=(0, 0, 1))
mlab.figure()
vis_mesh(mesh, color=(0, 1, 0))
mlab.figure()
vis_mesh(template_mesh, color=(1, 0, 0))
mlab.figure()
vis_point_cloud(
np.array(cloud), scale_factor=0.01, color=(0, 1, 0))
mlab.figure()
vis_voxels(voxels.data, color=(0, 1, 0))
plt.show(block=False)
mlab.show()
plt.close()
def vis_segmentations(model_id,
example_ids=None,
vis_mesh=False,
edge_length_threshold=0.02,
include_wireframe=False,
save=False):
from scipy.misc import imsave
if save and example_ids is None:
raise ValueError('Cannot save without specifying example_ids')
builder = get_builder(model_id)
cat_id = builder.cat_id
if example_ids is None:
example_ids = example_ids = get_example_ids(cat_id, 'eval')
if vis_mesh:
segmented_fn = builder.get_segmented_mesh_fn(edge_length_threshold)
else:
segmented_fn = builder.get_segmented_cloud_fn()
config = RenderConfig()
with get_predictions_dataset(model_id) as predictions:
with config.get_dataset(cat_id, builder.view_index) as image_ds:
for example_id in example_ids:
example = predictions[example_id]
probs, dp = (np.array(example[k]) for k in ('probs', 'dp'))
result = segmented_fn(probs, dp)
if result is not None:
image = image_ds[example_id]
print(example_id)
segmentation = result['segmentation']
if vis_mesh:
vertices = result['vertices']
faces = result['faces']
original_points = result['original_points']
original_seg = result['original_segmentation']
f0 = mlab.figure(bgcolor=(1, 1, 1))
vis_segmented_mesh(vertices,
segmented_cloud(
faces, segmentation),
include_wireframe=include_wireframe,
opacity=0.2)
f1 = mlab.figure(bgcolor=(1, 1, 1))
vis_clouds(
segmented_cloud(original_points, original_seg))
else:
points = result['points']
original_points = result['original_points']
f0 = mlab.figure(bgcolor=(1, 1, 1))
vis_clouds(segmented_cloud(points, segmentation))
f1 = mlab.figure(bgcolor=(1, 1, 1))
vis_clouds(
segmented_cloud(original_points, segmentation))
if save:
folder = os.path.join(_paper_dir, 'segmentations',
model_id, example_id)
if not os.path.isdir(folder):
os.makedirs(folder)
fn = 'inferred_%s.png' % ('mesh'
if vis_mesh else 'cloud')
p0 = os.path.join(folder, fn)
mlab.savefig(p0, figure=f0)
p1 = os.path.join(folder, 'annotated_cloud.png')
mlab.savefig(p1, figure=f1)
pi = os.path.join(folder, 'query_image.png')
imsave(pi, image)
mlab.close()
else:
plt.imshow(image)
plt.show(block=False)
mlab.show()
plt.close()
def vis_segmentations(
model_id, example_ids=None, vis_mesh=False,
edge_length_threshold=0.02, include_wireframe=False,
save=False):
from scipy.misc import imsave
if save and example_ids is None:
raise ValueError('Cannot save without specifying example_ids')
builder = get_builder(model_id)
cat_id = builder.cat_id
if example_ids is None:
example_ids = example_ids = get_example_ids(cat_id, 'eval')
if vis_mesh:
segmented_fn = builder.get_segmented_mesh_fn(edge_length_threshold)
else:
segmented_fn = builder.get_segmented_cloud_fn()
config = RenderConfig()
with get_predictions_dataset(model_id) as predictions:
with config.get_dataset(cat_id, builder.view_index) as image_ds:
for example_id in example_ids:
example = predictions[example_id]
probs, dp = (np.array(example[k]) for k in ('probs', 'dp'))
result = segmented_fn(probs, dp)
if result is not None:
image = image_ds[example_id]
print(example_id)
segmentation = result['segmentation']
if vis_mesh:
vertices = result['vertices']
faces = result['faces']
original_points = result['original_points']
original_seg = result['original_segmentation']
f0 = mlab.figure(bgcolor=(1, 1, 1))
vis_segmented_mesh(
vertices, segmented_cloud(faces, segmentation),
include_wireframe=include_wireframe,
opacity=0.2)
f1 = mlab.figure(bgcolor=(1, 1, 1))
vis_clouds(
segmented_cloud(original_points, original_seg))
else:
points = result['points']
original_points = result['original_points']
f0 = mlab.figure(bgcolor=(1, 1, 1))
vis_clouds(segmented_cloud(points, segmentation))
f1 = mlab.figure(bgcolor=(1, 1, 1))
vis_clouds(
segmented_cloud(original_points, segmentation))
if save:
folder = os.path.join(
_paper_dir, 'segmentations', model_id, example_id)
if not os.path.isdir(folder):
os.makedirs(folder)
fn = 'inferred_%s.png' % (
'mesh' if vis_mesh else 'cloud')
p0 = os.path.join(folder, fn)
mlab.savefig(p0, figure=f0)
p1 = os.path.join(folder, 'annotated_cloud.png')
mlab.savefig(p1, figure=f1)
pi = os.path.join(folder, 'query_image.png')
imsave(pi, image)
mlab.close()
else:
plt.imshow(image)
plt.show(block=False)
mlab.show()
plt.close()
if not os.path.isfile(manager.path):
manager.save_all()
datasets[cat_id] = manager.get_saving_dataset(
mode='r').subset(e_ids)
return BiKeyDataset(datasets).map(
lambda x: sample_points(np.array(x, dtype=np.float32), n_resamples))
if __name__ == '__main__':
from shapenet.core import cat_desc_to_id
from template_ffd.data.ids import get_example_ids
import random
cat_ids = [cat_desc_to_id(i) for i in ('plane', 'car')]
view_indices = [1, 5, 6]
mode = 'train'
example_ids = [get_example_ids(cat_id, mode) for cat_id in cat_ids]
image_dataset = get_image_dataset(cat_ids, example_ids, view_indices)
cloud_dataset = get_cloud_dataset(cat_ids, example_ids)
image_dataset.open()
cloud_dataset.open()
keys = list((tuple(k) for k in image_dataset.keys()))
random.shuffle(keys)
def vis(image, cloud):
import matplotlib.pyplot as plt
from util3d.mayavi_vis import vis_point_cloud, mlab
plt.imshow(image)
plt.show(block=False)
vis_point_cloud(
