def get_prediction_to_mesh_fn(self, edge_length_threshold=None):
cat_id = self.cat_id
with get_ffd_dataset(cat_id, self.n, edge_length_threshold) \
as ffd_dataset:
example_ids, bs, ps = zip(*self.get_ffd_data(ffd_dataset))
with get_template_mesh_dataset(cat_id, edge_length_threshold) as \
mesh_dataset:
all_faces = []
all_vertices = []
for k in example_ids:
sg = mesh_dataset[k]
all_faces.append(np.array(sg['faces']))
all_vertices.append(np.array(sg['vertices']))
def transform_predictions(probs, dp):
i = np.argmax(probs)
# print(ps[i].shape)
# print(dp[i].shape)
# print(bs[i].shape)
# exit()
vertices = np.matmul(bs[i], ps[i] + dp[i])
faces = all_faces[i]
original_vertices = all_vertices[i]
return dict(
vertices=vertices,
faces=faces,
original_vertices=original_vertices,
attrs=dict(template_id=example_ids[i]))
return transform_predictions
def get_prediction_to_top_k_mesh_fn(
self, edge_length_threshold=None, top_k=2):
cat_id = self.cat_id
with get_ffd_dataset(cat_id, self.n, edge_length_threshold) \
as ffd_dataset:
example_ids, bs, ps = zip(*self.get_ffd_data(ffd_dataset))
with get_template_mesh_dataset(cat_id, edge_length_threshold) as \
mesh_dataset:
all_faces = []
all_vertices = []
for k in example_ids:
sg = mesh_dataset[k]
all_faces.append(np.array(sg['faces']))
all_vertices.append(np.array(sg['vertices']))
def get_deformed_mesh(i, dp):
vertices = np.matmul(bs[i], ps[i] + dp[i])
faces = all_faces[i]
return dict(
vertices=vertices, faces=faces,
original_vertices=all_vertices[i])
def transform_predictions(probs, dp):
ks = probs.argsort()[-3:][::-1]
return [get_deformed_mesh(k, dp) for k in ks]
return transform_predictions
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 get_prediction_to_mesh_fn(self, edge_length_threshold=None):
cat_id = self.cat_id
if not isinstance(cat_id, (list, tuple)):
cat_id = [cat_id]
with get_ffd_dataset(self.cat_id, self.n,
edge_length_threshold=edge_length_threshold) as d:
cat_ids, example_ids, bs, ps = zip(*self.get_ffd_data(d))
with get_template_mesh_dataset(cat_id, edge_length_threshold) as \
mesh_dataset:
all_faces = []
all_vertices = []
for cat_id, example_id in zip(cat_ids, example_ids):
sg = mesh_dataset[cat_id, example_id]
all_faces.append(np.array(sg['faces']))
all_vertices.append(np.array(sg['vertices']))
def transform_predictions(probs, dp):
i = np.argmax(probs)
vertices = np.matmul(bs[i], ps[i] + dp[i])
faces = all_faces[i]
original_vertices = all_vertices[i]
return dict(
vertices=vertices,
faces=faces,
original_vertices=original_vertices,
attrs=dict(template_id=example_ids[i]))
return transform_predictions
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 get_ffd_data(self, ffd_dataset=None):
if ffd_dataset is None:
n_ffd_points = self.n_ffd_samples
ffd_dataset = get_ffd_dataset(
self.cat_id, self.n, n_samples=n_ffd_points)
with ffd_dataset:
return tuple(self._get_ffd_data(ffd_dataset))
else:
return self._get_ffd_data(ffd_dataset)
def get_segmented_mesh_fn(self, edge_length_threshold=None):
from shapenet.core.annotations.datasets import PointCloudDataset, \
SegmentationDataset
from dataset import Dataset
cat_id = self.cat_id
bs = []
ps = []
segs = []
faces = []
original_segs = []
original_seg_points = []
ffd_dataset = get_ffd_dataset(
cat_id, self.n, edge_length_threshold=edge_length_threshold)
with ffd_dataset:
example_ids, bs, ps = zip(*self.get_ffd_data(ffd_dataset))
template_mesh_ds = get_template_mesh_dataset(
cat_id, edge_length_threshold=edge_length_threshold)
seg_points_ds = PointCloudDataset(cat_id)
seg_ds = SegmentationDataset(cat_id)
ds = Dataset.zip(template_mesh_ds, seg_points_ds, seg_ds)
with ds:
for example_id in example_ids:
if example_id in ds:
template_mesh, seg_points, original_seg = ds[example_id]
v, f = (np.array(template_mesh[k])
for k in ('vertices', 'faces'))
centroids = get_centroids(v, f)
seg = original_seg[get_nn(seg_points, centroids)]
else:
f = None
seg = None
seg_points = None
original_seg = None
segs.append(seg)
original_seg_points.append(seg_points)
original_segs.append(original_seg)
faces.append(f)
def transform_predictions(probs, dp):
i = np.argmax(probs)
seg = segs[i]
if seg is None:
return None
else:
v = np.matmul(bs[i], ps[i] + dp[i])
return dict(faces=faces[i],
vertices=v,
segmentation=segs[i],
original_points=original_seg_points[i],
original_segmentation=original_segs[i])
return transform_predictions
def get_prediction_to_cloud_fn(self, n_samples=None):
from util3d.point_cloud import sample_points
with get_ffd_dataset(
self.cat_id, self.n, n_samples=self.n_ffd_samples) \
as ffd_dataset:
example_ids, bs, ps = zip(*self.get_ffd_data(ffd_dataset))
def transform_predictions(probs, dp):
i = np.argmax(probs)
b = bs[i]
if n_samples is not None:
b = sample_points(b, n_samples)
points = np.matmul(b, ps[i] + dp[i])
return dict(cloud=points, attrs=dict(template_id=example_ids[i]))
return transform_predictions
