def create_archive(cat_desc, shape=(192, 256), n_images=8, example_ids=None):
import os
from shapenet.core import cat_desc_to_id
from shapenet.core import get_example_ids
from shapenet.core.blender_renderings.config import RenderConfig
from progress.bar import IncrementalBar
import zipfile
cat_id = cat_desc_to_id(cat_desc)
if example_ids is None or len(example_ids) == 0:
example_ids = get_example_ids(cat_id)
config = RenderConfig(shape=shape, n_images=n_images)
zip_path = config.get_zip_path(cat_id)
with zipfile.ZipFile(zip_path, mode='a', allowZip64=True) as zf:
bar = IncrementalBar(max=len(example_ids))
for example_id in example_ids:
example_dir = config.get_example_dir(cat_id, example_id)
if not os.path.isdir(example_dir):
print('No directory at %s' % example_dir)
else:
for fn in os.listdir(example_dir):
src = os.path.join(example_dir, fn)
dst = os.path.join(cat_id, example_id, fn)
zf.write(src, dst)
bar.next()
bar.finish()
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 vis(cat, n_images, view_index=5, example_ids=None):
import matplotlib.pyplot as plt
from shapenet.core import cat_desc_to_id, get_example_ids
from shapenet.core.blender_renderings.config import RenderConfig
cat_id = cat_desc_to_id(cat)
config = RenderConfig(n_images=n_images)
dataset = config.get_dataset(cat_id, view_index)
if example_ids is not None and len(example_ids) > 0:
dataset = dataset.subset(example_ids)
else:
example_ids = get_example_ids(cat_id)
with dataset:
for example_id in example_ids:
plt.imshow(dataset[example_id])
plt.title(example_id)
plt.show()
def get_ds(cat_desc, regime='e'):
view_index = 5
edge_length_threshold = 0.02
cat_id = cat_desc_to_id(cat_desc)
model_id = '%s_%s' % (regime, cat_desc)
image_ds = RenderConfig().get_dataset(cat_id, view_index)
cloud_ds = get_cloud_manager(
model_id, pre_sampled=True, n_samples=n_samples).get_lazy_dataset()
mesh_ds = get_inferred_mesh_dataset(
model_id, edge_length_threshold=edge_length_threshold)
gt_mesh_ds = get_mesh_dataset(cat_id)
voxel_ds = get_voxel_dataset(
model_id, edge_length_threshold=edge_length_threshold, filled=False)
selected_template_ds = get_selected_template_idx_dataset(model_id)
template_meshes = []
with gt_mesh_ds:
for template_id in get_template_ids(cat_id):
mesh = gt_mesh_ds[template_id]
template_meshes.append(
{k: np.array(mesh[k]) for k in ('vertices', 'faces')})
template_mesh_ds = selected_template_ds.map(lambda i: template_meshes[i])
return Dataset.zip(
image_ds, gt_mesh_ds, cloud_ds, mesh_ds, voxel_ds, template_mesh_ds)
def get_image_dataset(cat_ids, example_ids, view_indices, render_config=None):
from shapenet.image import with_background
from dids.core import BiKeyDataset
if render_config is None:
from shapenet.core.blender_renderings.config import RenderConfig
render_config = RenderConfig()
if isinstance(cat_ids, str):
cat_ids = [cat_ids]
example_ids = [example_ids]
if isinstance(view_indices, int):
view_indices = [view_indices]
datasets = {
c: render_config.get_multi_view_dataset(
c, view_indices=view_indices, example_ids=eid)
for c, eid in zip(cat_ids, example_ids)}
dataset = BiKeyDataset(datasets).map(
lambda image: with_background(image, 255))
dataset = dataset.map_keys(
lambda key: (key[0], (key[1], key[2])),
lambda key: (key[0],) + key[1])
return dataset
def check_zip(cat_desc, shape, n_images):
import zipfile
from shapenet.core.blender_renderings.config import RenderConfig
from shapenet.core import cat_desc_to_id, get_example_ids
cat_id = cat_desc_to_id(cat_desc)
config = RenderConfig(shape=shape, n_images=n_images)
rendered_ids = set()
with zipfile.ZipFile(config.get_zip_path(cat_id)) as zf:
for name in zf.namelist():
rendered_ids.add(name.split('/')[1])
not_rendered_count = 0
example_ids = get_example_ids(cat_id)
for example_id in example_ids:
if example_id not in rendered_ids:
print(example_id)
not_rendered_count += 1
if not_rendered_count > 0:
print('%d / %d not rendered' % (not_rendered_count, len(example_ids)))
else:
print('All %d %ss rendered!' % (len(example_ids), cat_desc))
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_inputs(self, mode, repeat=None):
from shapenet.core.blender_renderings.config import RenderConfig
from ..data.ids import get_example_ids
render_config = RenderConfig()
view_index = self.params.get('view_index',
range(render_config.n_images))
cat_ids = self.cat_ids
example_ids = tuple(
get_example_ids(cat_id, mode) for cat_id in cat_ids)
if repeat is None:
repeat = mode == tf.estimator.ModeKeys.TRAIN
dataset = get_tf_dataset(render_config,
view_index,
cat_ids,
example_ids,
batch_size=self.batch_size,
shuffle=True,
repeat=repeat)
return dataset.make_one_shot_iterator().get_next()
def render_config(self):
from shapenet.core.blender_renderings.config import RenderConfig
return RenderConfig(**self.params.get('render_params', {}))
file_index = get_file_index(zip_file)
for example_id in example_ids:
bar.next()
render_example(
config, cat_id, example_id, zip_file,
overwrite, call_kwargs, blender_path=blender_path,
verbose=verbose, file_index=file_index)
bar.finish()
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('cat', type=str)
parser.add_argument('--shape', type=int, nargs=2, default=[192, 256])
parser.add_argument('--scale', type=float, default=None)
parser.add_argument('--blender_path', type=str, default='blender')
parser.add_argument('-n', '--n_images', type=int, default=8)
parser.add_argument('-d', '--debug', action='store_true')
parser.add_argument('-r', '--reverse', action='store_true')
parser.add_argument('-o', '--overwrite', action='store_true')
parser.add_argument('-i', '--example_ids', nargs='*')
parser.add_argument('-f', '--fixed_meshes', action='store_true')
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()
config = RenderConfig(args.shape, args.n_images, args.scale)
cat_id = cat_id = cat_desc_to_id(args.cat)
render_cat(config, cat_id, args.overwrite, args.reverse, args.debug,
args.example_ids, args.fixed_meshes, args.blender_path,
args.verbose)
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()
from template_ffd.data.ids import get_example_ids
from template_ffd.model import get_builder
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)
import os
import matplotlib.pyplot as plt
from shapenet.image import with_background
from shapenet.core.blender_renderings.config import RenderConfig
from shapenet.core import cat_desc_to_id, get_example_ids
cat_desc = 'plane'
view_index = 5
config = RenderConfig()
view_angle = config.view_angle(view_index)
cat_id = cat_desc_to_id(cat_desc)
example_ids = get_example_ids(cat_id)
path = config.get_zip_path(cat_id)
if not os.path.isfile(path):
raise IOError('No renderings at %s' % path)
with config.get_dataset(cat_id, view_index) as ds:
ds = ds.map(lambda image: with_background(image, 255))
for example_id in ds:
image = ds[example_id]
plt.imshow(image)
plt.show()
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()
@property
def output_shape(self):
return self._config.shape + (3, )
@property
def output_type(self):
return tf.uint8
def get_renderings_dataset(render_config, view_index, cat_id, example_ids):
manager = RenderingsManager(render_config, view_index, cat_id)
return base_dataset(example_ids).map(manager.map_tf)
if __name__ == '__main__':
import matplotlib.pyplot as plt
from shapenet.core import cat_desc_to_id, get_example_ids
from shapenet.core.blender_renderings.config import RenderConfig
cat_desc = 'plane'
view_index = 5
config = RenderConfig()
cat_id = cat_desc_to_id(cat_desc)
example_ids = get_example_ids(cat_id)
dataset = get_renderings_dataset(config, view_index, cat_id, example_ids)
image_tf = dataset.make_one_shot_iterator().get_next()
with tf.train.MonitoredSession() as sess:
while not sess.should_stop():
image = sess.run(image_tf)
plt.imshow(image)
plt.show()