def get_data(model_id, example_ids=None):
edge_length_threshold = 0.02
builder = get_builder(model_id)
cat_id = builder.cat_id
with get_ffd_dataset(cat_id, edge_length_threshold=0.02) as ffd_ds:
template_ids, bs, ps = zip(*builder.get_ffd_data(ffd_ds))
with get_template_mesh_dataset(cat_id, edge_length_threshold) as mesh_ds:
faces = [np.array(mesh_ds[e]['faces']) for e in template_ids]
predictions_ds = get_predictions_dataset(model_id)
mesh_ds = get_mesh_dataset(cat_id)
image_ds = RenderConfig().get_dataset(cat_id, builder.view_index)
zipped = Dataset.zip(predictions_ds, mesh_ds, image_ds)
with zipped:
if example_ids is None:
example_ids = list(predictions_ds.keys())
random.shuffle(example_ids)
for example_id in example_ids:
print(example_id)
pred, mesh, image = zipped[example_id]
i = np.argmax(pred['probs'])
dp = np.array(pred['dp'][i])
b = bs[i]
p = ps[i]
yield example_id, b, p, dp, faces[i], mesh, image
def print_template_scores(model_id, by_weight=False):
builder = get_builder(model_id)
template_ids = builder.template_ids
n = len(template_ids)
counts = np.zeros((n,), dtype=np.int32)
totals = np.zeros((n,), dtype=np.float32)
dataset = get_predictions_dataset(model_id)
with dataset:
for example_id in dataset:
probs = np.array(dataset[example_id]['probs'])
counts[np.argmax(probs)] += 1
totals += probs
if by_weight:
zipped = list(zip(template_ids, range(n), totals))
zipped.sort(key=lambda x: x[2], reverse=True)
for rank, (k, i, p) in enumerate(zipped):
print(rank, i, p, k)
print([z[1] for z in zipped])
else:
zipped = list(zip(template_ids, range(n), counts))
zipped.sort(key=lambda x: x[2], reverse=True)
for rank, (k, i, p) in enumerate(zipped):
print(rank, i, p, k)
print([z[1] for z in zipped])
def print_template_scores(model_id, by_weight=False):
builder = get_builder(model_id)
template_ids = builder.template_ids
n = len(template_ids)
counts = np.zeros((n, ), dtype=np.int32)
totals = np.zeros((n, ), dtype=np.float32)
dataset = get_predictions_dataset(model_id)
with dataset:
for example_id in dataset:
probs = np.array(dataset[example_id]['probs'])
counts[np.argmax(probs)] += 1
totals += probs
if by_weight:
zipped = list(zip(template_ids, range(n), totals))
zipped.sort(key=lambda x: x[2], reverse=True)
for rank, (k, i, p) in enumerate(zipped):
print(rank, i, p, k)
print([z[1] for z in zipped])
else:
zipped = list(zip(template_ids, range(n), counts))
zipped.sort(key=lambda x: x[2], reverse=True)
for rank, (k, i, p) in enumerate(zipped):
print(rank, i, p, k)
print([z[1] for z in zipped])
def get_data(model_id, example_ids=None):
edge_length_threshold = 0.02
builder = get_builder(model_id)
cat_id = builder.cat_id
with get_ffd_dataset(cat_id, edge_length_threshold=0.02) as ffd_ds:
template_ids, bs, ps = zip(*builder.get_ffd_data(ffd_ds))
with get_template_mesh_dataset(cat_id, edge_length_threshold) as mesh_ds:
faces = [np.array(mesh_ds[e]['faces']) for e in template_ids]
predictions_ds = get_predictions_dataset(model_id)
mesh_ds = get_mesh_dataset(cat_id)
image_ds = RenderConfig().get_dataset(cat_id, builder.view_index)
zipped = Dataset.zip(predictions_ds, mesh_ds, image_ds)
with zipped:
if example_ids is None:
example_ids = list(predictions_ds.keys())
random.shuffle(example_ids)
for example_id in example_ids:
print(example_id)
pred, mesh, image = zipped[example_id]
i = np.argmax(pred['probs'])
dp = np.array(pred['dp'][i])
b = bs[i]
p = ps[i]
yield example_id, b, p, dp, faces[i], mesh, image
def clear_results(model_id, actual=False):
import os
import shutil
from template_ffd.inference.predictions import get_predictions_dataset
from template_ffd.inference.meshes import InferredMeshManager
from template_ffd.inference.clouds import get_cloud_manager
from template_ffd.inference.voxels import get_voxel_dataset
from template_ffd.eval.chamfer import \
get_chamfer_manager, get_template_chamfer_manager
from template_ffd.eval.iou import IouAutoSavingManager
def maybe_remove(path):
if os.path.isfile(path):
print('Removing file %s' % path)
if actual:
os.remove(path)
elif os.path.isdir(path):
print('Removing subdir %s' % path)
if actual:
shutil.rmtree(path)
predictions_dataset = get_predictions_dataset(model_id)
maybe_remove(predictions_dataset.path)
maybe_remove(get_cloud_manager(model_id, pre_sampled=True).path)
maybe_remove(get_chamfer_manager(model_id, pre_sampled=True).path)
for elt in (None, 0.1, 0.05, 0.02, 0.01):
maybe_remove(InferredMeshManager(model_id, elt).path)
maybe_remove(
get_cloud_manager(model_id,
pre_sampled=False,
edge_length_threshold=elt).path)
for filled in (True, False):
dataset = get_voxel_dataset(model_id,
elt,
filled=filled,
auto_save=False)
maybe_remove(dataset.root_dir)
maybe_remove(IouAutoSavingManager(model_id, elt, filled).path)
maybe_remove(
get_chamfer_manager(model_id,
pre_sampled=False,
edge_length_threshold=elt).path)
maybe_remove(get_template_chamfer_manager(model_id).path)
if not actual:
print('NOTE: this was a dry run. Files not actually removed')
print('Use -a for actual run.')
def get_lazy_dataset(self):
from shapenet.core.point_clouds import get_point_cloud_dataset
from util3d.point_cloud import sample_points
from template_ffd.model import get_builder
from template_ffd.inference.predictions import get_predictions_dataset
builder = get_builder(self._model_id)
cat_id = builder.cat_id
template_ids = builder.template_ids
clouds = []
def sample_fn(cloud):
return sample_points(np.array(cloud), self._n_samples)
gt_clouds = get_point_cloud_dataset(
cat_id, builder.n_samples).map(sample_fn)
with gt_clouds:
for example_id in template_ids:
clouds.append(np.array(gt_clouds[example_id]))
predictions = get_predictions_dataset(self._model_id)
inf_cloud_ds = predictions.map(lambda i: clouds[i].copy())
return _get_lazy_emd_dataset(inf_cloud_ds, cat_id, self._n_samples)
def get_lazy_dataset(self):
from shapenet.core.point_clouds import get_point_cloud_dataset
from util3d.point_cloud import sample_points
from template_ffd.model import get_builder
from template_ffd.inference.predictions import get_predictions_dataset
builder = get_builder(self._model_id)
cat_id = builder.cat_id
template_ids = builder.template_ids
clouds = []
def sample_fn(cloud):
return sample_points(np.array(cloud), self._n_samples)
gt_clouds = get_point_cloud_dataset(
cat_id, builder.n_samples).map(sample_fn)
with gt_clouds:
for example_id in template_ids:
clouds.append(np.array(gt_clouds[example_id]))
predictions = get_predictions_dataset(self._model_id)
inf_cloud_ds = predictions.map(lambda i: clouds[i].copy())
return _get_lazy_emd_dataset(inf_cloud_ds, cat_id, self._n_samples)
def get_hist_data(model_id, n_bins, mode):
from shapenet.core import cat_desc_to_id
from template_ffd.templates.ids import get_template_ids
from template_ffd.model import load_params
from template_ffd.inference.predictions import get_predictions_dataset
cat_id = cat_desc_to_id(load_params(model_id)['cat_desc'])
n_templates = len(get_template_ids(cat_id))
counts = np.zeros((n_bins, ), dtype=np.int32)
argmax_counts = np.zeros((n_templates, ), dtype=np.int32)
with get_predictions_dataset(model_id) as dataset:
for example_id in dataset:
probs = np.array(dataset[example_id]['probs'])
counts[int(np.max(probs) * n_bins)] += 1
# prob_indices = np.array(0.999*probs * n_bins, dtype=np.int32)
# for pi in prob_indices:
# counts[pi] += 1
argmax_counts[np.argmax(probs)] += 1
counts = counts / np.sum(counts)
argmax_counts = argmax_counts / np.sum(argmax_counts)
return counts, argmax_counts
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)
regime = 'e'
cat_desc = 'chair'
view_index = 5
edge_length_threshold = 0.02
shuffle = True
k = 3
cat_id = cat_desc_to_id(cat_desc)
model_id = '%s_%s' % (regime, cat_desc)
builder = get_builder(model_id)
image_ds = RenderConfig().get_dataset(cat_id, view_index)
gt_mesh_ds = get_mesh_dataset(cat_id)
predictions_ds = get_predictions_dataset(model_id)
top_k_mesh_fn = builder.get_prediction_to_top_k_mesh_fn(
edge_length_threshold, k)
all_ds = Dataset.zip(image_ds, gt_mesh_ds, predictions_ds)
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'))
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()
regime = 'e'
cat_desc = 'chair'
view_index = 5
edge_length_threshold = 0.02
shuffle = True
k = 3
cat_id = cat_desc_to_id(cat_desc)
model_id = '%s_%s' % (regime, cat_desc)
builder = get_builder(model_id)
image_ds = RenderConfig().get_dataset(cat_id, view_index)
gt_mesh_ds = get_mesh_dataset(cat_id)
predictions_ds = get_predictions_dataset(model_id)
top_k_mesh_fn = builder.get_prediction_to_top_k_mesh_fn(
edge_length_threshold, k)
all_ds = Dataset.zip(image_ds, gt_mesh_ds, predictions_ds)
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)
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()