def test_hybridcloud_pkl_interface():
hc = test_hybridcloud_load()
fname = f'{test_dir}/test_hc.pkl'
try:
hc.save2pkl(fname)
hc2 = HybridCloud()
hc2.load_from_pkl(fname)
assert hc == hc2
finally:
if os.path.isfile(fname):
os.remove(fname)
def preds2kzip(pred_folder: str, out_path: str, ssd_path: str, col_lookup: dict,
label_mappings: Optional[List[Tuple[int, int]]] = None):
pred_folder = os.path.expanduser(pred_folder)
out_path = os.path.expanduser(out_path)
if not os.path.exists(out_path):
os.makedirs(out_path)
files = glob.glob(pred_folder + '*_preds.pkl')
ssd = SuperSegmentationDataset(ssd_path)
for file in tqdm(files):
hc_voxeled = preds2hc(file)
sso_id = int(re.findall(r"/sso_(\d+).", file)[0])
sso = ssd.get_super_segmentation_object(sso_id)
verts = sso.mesh[1].reshape(-1, 3)
hc = HybridCloud(nodes=hc_voxeled.nodes, edges=hc_voxeled.edges, node_labels=hc_voxeled.node_labels,
pred_node_labels=hc_voxeled.pred_node_labels, vertices=verts)
hc.nodel2vertl()
hc.prednodel2predvertl()
if label_mappings is not None:
hc.map_labels(label_mappings)
cols = np.array([col_lookup[el] for el in hc.pred_labels.squeeze()], dtype=np.uint8)
sso.mesh2kzip(out_path + f'p_{sso_id}.k.zip', ext_color=cols)
cols = np.array([col_lookup[el] for el in hc.labels.squeeze()], dtype=np.uint8)
sso.mesh2kzip(out_path + f't_{sso_id}.k.zip', ext_color=cols)
comments = list(hc.pred_node_labels.reshape(-1))
for node in range(len(hc.nodes)):
if hc.pred_node_labels[node] != hc.node_labels[node] and hc.pred_node_labels[node] != -1:
comments[node] = 'e' + str(comments[node])
sso.save_skeleton_to_kzip(out_path + f'p_{sso_id}.k.zip', comments=comments)
comments = hc.node_labels.reshape(-1)
sso.save_skeleton_to_kzip(out_path + f't_{sso_id}.k.zip', comments=comments)
def load_obj(
data_type: str, file: str
) -> Union[HybridMesh, HybridCloud, PointCloud, CloudEnsemble]:
if data_type == 'obj':
return basics.load_pkl(file)
if data_type == 'ce':
return ensembles.ensemble_from_pkl(file)
if data_type == 'hc':
hc = HybridCloud()
return hc.load_from_pkl(file)
if data_type == 'hm':
hm = HybridMesh()
return hm.load_from_pkl(file)
else:
pc = PointCloud()
return pc.load_from_pkl(file)
def merge_hybrid(hc: HybridCloud,
clouds: List[PointCloud],
hc_name: str = None,
names: List = None,
preserve_obj_bounds: bool = False) -> HybridCloud:
""" Merges a list of PointCloud objects with a single HybridCloud using the merge_clouds method.
See merge_clouds method for more information.
Args:
hc: The single HybridCloud.
clouds: The list of PointCloud objects.
hc_name: The name of the HybridCloud.
names: The names of the clouds in clouds.
preserve_obj_bounds: see merge_clouds method.
"""
clouds.append(hc)
names.append(hc_name)
pc = merge(clouds, names, preserve_obj_bounds=preserve_obj_bounds)
return HybridCloud(nodes=hc.nodes,
edges=hc.edges,
vertices=pc.vertices,
labels=pc.labels,
pred_labels=pc.pred_labels,
features=pc.features,
types=pc.types,
encoding=pc.encoding,
obj_bounds=pc.obj_bounds,
predictions=pc.predictions,
no_pred=pc.no_pred,
node_labels=hc.node_labels,
pred_node_labels=hc.pred_node_labels)
def test_composition():
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
transform = clouds.Compose([
clouds.Normalization(10),
clouds.RandomRotate((60, 60)),
clouds.Center()
])
transform(pc)
transform(hc)
assert np.all(
np.round(np.mean(pc.vertices, axis=0)) == np.array([0, 0, 0]))
assert np.all(
np.round(np.mean(hc.vertices, axis=0)) == np.array([0, 0, 0]))
dummy = np.array([[10, 10, 10], [20, 20, 20]]) / 10
angle_range = (60, 60)
angles = np.random.uniform(angle_range[0], angle_range[1], (1, 3))[0]
rot = Rot.from_euler('xyz', angles, degrees=True)
dummy = rot.apply(dummy)
centroid = np.mean(dummy, axis=0)
dummy = dummy - centroid
assert np.all(pc.vertices == dummy)
assert np.all(hc.vertices == dummy)
assert np.all(hc.vertices == dummy)
def filter_labels(cloud: PointCloud, labels: list) -> PointCloud:
""" Returns a pointcloud which contains only those vertices which labels occuring in 'labels'. If 'cloud'
is a HybridCloud, the skeleton is taken as it is and should later be filtered with the 'filter_traverser'
method.
Args:
cloud: PointCloud which should be filtered.
labels: List of labels for which the corresponding vertices should be extracted.
Returns:
PointCloud object which contains only vertices with the filtered labels. Skeletons in case of HybridClouds are
the same.
"""
mask = np.zeros(len(cloud.labels), dtype=bool)
for label in labels:
mask = np.logical_or(mask, cloud.labels == label)
if isinstance(cloud, HybridCloud):
f_cloud = HybridCloud(cloud.nodes,
cloud.edges,
vertices=cloud.vertices[mask],
labels=cloud.labels[mask])
else:
f_cloud = PointCloud(cloud.vertices[mask], labels=cloud.labels[mask])
return f_cloud
def visualize_prediction(obj: Union[CloudEnsemble, HybridCloud],
out_path: str,
random_seed: int = 4,
save_to_image: bool = True):
""" Saves images of prediction of given file in out_path. First, the predictions get reduced onto the labels.
Labels without predictions get filtered. Second, the vertex labels are mapped onto the skeleton, where
nodes without vertices take on the label of the nearest neighbor with vertices. Third, the node labels get
mapped onto the mesh again. Then an image of the mesh is saved to file.
Args:
obj: MorphX object with predictions
out_path: Folder where image should be saved
random_seed: Defines the color coding
"""
obj.generate_pred_labels()
if isinstance(obj, CloudEnsemble):
obj = obj.hc
# reduce object to vertices where predictions exist
mask = obj.pred_labels != -1
mask = mask.reshape(-1)
red_obj = HybridCloud(nodes=obj.nodes,
edges=obj.edges,
vertices=obj.vertices[mask],
labels=obj.labels[mask],
pred_labels=obj.pred_labels[mask])
obj.set_pred_node_labels(red_obj.pred_node_labels)
obj.prednodel2predvertl()
obj.set_labels(obj.pred_labels)
visualize_clouds([obj],
capture=save_to_image,
path=out_path,
random_seed=random_seed)
def label_search(hc: HybridCloud, node_labels: np.ndarray, source: int) -> int:
""" Performs BFS on nodes starting from source until first node with label != -1 has been found.
Notes:
* Node IDs f input and output HybridCloud will not be the same!
Args:
hc: HybridCloud in which source is part of the graph
node_labels: array of node labels (cannot be accessed through hc as it doesn't exists).
source: The node for which the first neighboring node with label != -1 should be found.
Returns:
The index of the first node with label != -1
"""
if node_labels is None:
return source
g = hc.graph()
visited = [source]
neighbors = g.neighbors(source)
de = deque([i for i in neighbors])
while de:
curr = de.pop()
if node_labels[curr] != -1:
return curr
if curr not in visited:
visited.append(curr)
neighbors = g.neighbors(curr)
de.extendleft([i for i in neighbors if i not in visited])
return source
def test_bfs_vertices():
expected = [0, 1, 4, 5]
hc = HybridCloud(nodes=np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2],
[3, 3, 3], [4, 4, 4], [5, 5, 5]]),
edges=np.array([[0, 1], [1, 2], [2, 3], [0, 4], [0, 5]]),
vertices=np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0],
[1, 1, 1], [1, 1, 1], [1, 1, 1],
[1, 1, 1], [2, 2, 2], [3, 3, 3],
[3, 3, 3], [3, 3, 3], [4, 4, 4],
[4, 4, 4], [4, 4, 4], [4, 4, 4],
[4, 4, 4], [4, 4, 4], [5, 5, 5]]))
chosen = morphx.processing.objects.bfs_vertices(hc, 0, 14)
print(chosen)
assert len(chosen) == len(expected)
for item in chosen:
assert item in expected
expected = [0]
chosen = morphx.processing.objects.bfs_vertices(hc, 0, 3)
print(chosen)
assert len(chosen) == len(expected)
for item in chosen:
assert item in expected
expected = [1, 2, 3]
chosen = morphx.processing.objects.bfs_vertices(hc, 2, 10)
print(chosen)
assert len(chosen) == len(expected)
for item in chosen:
assert item in expected
def ensemble_from_pkl(path):
""" Loads an ensemble from an existing pickle file.
Args:
path: File path of pickle file.
"""
path = os.path.expanduser(path)
if not os.path.exists(path):
print(f"File with name: {path} was not found at this location.")
with open(path, 'rb') as f:
obj = pickle.load(f)
f.close()
if not isinstance(obj, dict):
raise ValueError(f"Object at {path} is no valid ensemble file.")
try:
h = HybridCloud(**obj['hybrid'])
except TypeError:
h = HybridMesh(**obj['hybrid'])
cloudlist = {}
for key in obj['clouds']:
try:
cloudlist[key] = PointCloud(**obj['clouds'][key])
except TypeError:
try:
cloudlist[key] = HybridCloud(**obj['clouds'][key])
except TypeError:
cloudlist[key] = HybridMesh(**obj['clouds'][key])
# check for empty clouds
empty = []
for key in cloudlist:
if cloudlist[key].vertices is None:
empty.append(key)
for key in empty:
cloudlist.pop(key, None)
try:
predictions = obj['predictions']
except KeyError:
predictions = None
try:
verts2node = obj['verts2node']
except KeyError:
verts2node = None
return CloudEnsemble(cloudlist,
h,
obj['no_pred'],
predictions=predictions,
verts2node=verts2node)
def extract_subset(hybrid: HybridCloud,
local_bfs: np.ndarray) -> Tuple[PointCloud, np.ndarray]:
""" Returns the mesh subset of given skeleton nodes based on a mapping dict
between skeleton and mesh.
Notes:
* Node IDs f input and output HybridCloud will not be the same!
Args:
hybrid: MorphX HybridCloud object from which the subset should be extracted.
local_bfs: Skeleton node index array which was generated by a local BFS.
Returns:
Mesh subset as HybridCloud object, vertex IDs from the parent
HybridCloud with the same ordering as vertices (enables "global" look-up).
"""
idcs = []
local_bfs = np.sort(local_bfs)
for i in local_bfs:
idcs.extend(hybrid.verts2node[i])
verts = hybrid.vertices[idcs]
feats, labels, node_labels = None, None, None
if hybrid.features is not None and len(hybrid.features) > 0:
feats = hybrid.features[idcs]
if hybrid.labels is not None and len(hybrid.labels) > 0:
labels = hybrid.labels[idcs]
if hybrid.node_labels is not None and len(hybrid.node_labels) > 0:
node_labels = hybrid.node_labels[local_bfs]
g = hybrid.graph(simple=True)
sub_g = g.subgraph(local_bfs)
relabel_dc = {n: ii for ii, n in enumerate(sub_g.nodes)}
sub_g = nx.relabel.relabel_nodes(sub_g, relabel_dc)
edges = np.array(list(sub_g.edges))
nodes = hybrid.nodes[local_bfs]
hc = HybridCloud(nodes=nodes,
edges=edges,
node_labels=node_labels,
vertices=verts,
labels=labels,
features=feats,
no_pred=hybrid.no_pred,
obj_bounds=hybrid.obj_bounds,
encoding=hybrid.encoding)
# add relabel dc in case it is needed downstream
hc.relabel_dc = relabel_dc
return hc, np.array(idcs)
def test_hybridcloud_load():
# TODO: maybe decouple from other tests.
dir_ex = os.path.abspath(f'{test_dir}/../data')
_, vertices = read_mesh_from_ply(f'{dir_ex}/example_mesh.ply')
nodes, edges = load_skeleton_nx_pkl(f'{dir_ex}/example_skel.pkl')
hc = HybridCloud(vertices=vertices, nodes=nodes, edges=edges)
assert not np.any(np.isnan(hc.vertices))
assert hc.vertices.ndim == 2, hc.vertices.shape[1] == 3
return hc
def test_normalization():
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
pc.scale(-10)
hc.scale(-10)
assert np.all(pc.vertices == np.array([[1, 1, 1], [2, 2, 2]]))
assert np.all(hc.vertices == np.array([[1, 1, 1], [2, 2, 2]]))
assert np.all(hc.nodes == np.array([[1, 1, 1], [2, 2, 2]]))
# test transformation class from processing.clouds
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
transform = clouds.Normalization(10)
transform(pc)
transform(hc)
assert np.all(pc.vertices == np.array([[1, 1, 1], [2, 2, 2]]))
assert np.all(hc.vertices == np.array([[1, 1, 1], [2, 2, 2]]))
assert np.all(hc.nodes == np.array([[1, 1, 1], [2, 2, 2]]))
def test_center():
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
relation = hc.vertices[0] - hc.nodes[1]
pc.move(np.array([1, 1, 1]))
hc.move(np.array([1, 1, 1]))
assert np.all(pc.vertices == np.array([[11, 11, 11], [21, 21, 21]]))
assert np.all(hc.vertices == np.array([[11, 11, 11], [21, 21, 21]]))
assert np.all(hc.nodes == np.array([[11, 11, 11], [21, 21, 21]]))
assert np.all(hc.vertices[0] - hc.nodes[1] == relation)
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
relation = hc.vertices[0] - hc.nodes[1]
# test transformation class from processing.clouds
transform = clouds.Center()
transform(pc)
transform(hc)
assert np.all(np.mean(pc.vertices, axis=0) == np.array([0, 0, 0]))
assert np.all(np.mean(hc.vertices, axis=0) == np.array([0, 0, 0]))
assert np.all(hc.vertices[0] - hc.nodes[1] == relation)
def test_node_labels():
vertices = [[i, i, i] for i in range(10)]
vertices += vertices
vertices += vertices
labels = [i for i in range(10)]
labels += labels
labels += labels
labels = np.array(labels)
labels[10:20] += 1
hc = HybridCloud(np.array([[i, i, i] for i in range(10)]),
np.array([[i, i + 1] for i in range(9)]),
vertices=np.array(vertices),
pred_labels=labels)
pred_node_labels = hc.pred_node_labels
expected = np.array([i for i in range(10)])
assert np.all(pred_node_labels == expected.reshape((len(expected), 1)))
node_labels = np.array([1, 2, 1, 1, 2, 2, 2, 1])
hc = HybridCloud(np.array([[i, i, i] for i in range(8)]),
np.array([[i, i + 1] for i in range(7)]),
vertices=np.array([[1, 1, 1], [2, 2, 2]]),
pred_node_labels=node_labels)
hc.node_sliding_window_bfs(2)
expected = np.array([1, 1, 1, 1, 2, 2, 2, 2])
assert np.all(hc.pred_node_labels == expected.reshape((len(expected), 1)))
def test_ensemble2pointcloud():
pc1 = PointCloud(np.array([[i, i, i] for i in range(10)]),
np.array([[i] for i in range(10)]),
obj_bounds={
'obj1': np.array([0, 5]),
'obj2': np.array([5, 10])
},
encoding={
'e1': 1,
'e2': 2
})
pc2 = PointCloud(np.array([[i, i, i] for i in range(10)]),
np.array([[i] for i in range(10)]))
hc = HybridCloud(np.array([[1, 1, 1], [2, 2, 2]]),
np.array([[0, 1]]),
vertices=np.array([[i, i, i] for i in range(10)]),
encoding={
'e1': 1,
'e3': 3
},
obj_bounds={
'obj3': np.array([0, 2]),
'obj4': np.array([2, 10])
})
ce = CloudEnsemble({'pc1': pc1, 'pc2': pc2}, hybrid=hc)
result = ensembles.ensemble2pointcloud(ce)
obj_bounds = {
'obj3': np.array([0, 2]),
'obj4': np.array([2, 10]),
'obj1': np.array([10, 15]),
'obj2': np.array([15, 20]),
'pc2': np.array([20, 30])
}
encoding = {'e1': 1, 'e2': 2, 'e3': 3}
assert np.all(result.nodes == hc.nodes)
assert np.all(result.edges == hc.edges)
assert np.all(result.vertices == np.concatenate(
(hc.vertices, pc1.vertices, pc2.vertices), axis=0))
assert np.all(result.labels == np.concatenate(
(np.ones((10, 1)) * -1, pc1.labels, pc2.labels), axis=0))
assert len(result.obj_bounds) == len(obj_bounds)
for key in obj_bounds:
assert np.all(obj_bounds[key] == result.obj_bounds[key])
assert len(result.encoding) == len(encoding)
for key in result.encoding:
assert result.encoding[key] == encoding[key]
def map_labels(cloud: PointCloud, labels: list, target) -> PointCloud:
""" Returns a PointCloud where all labels given in the labels list got mapped to the target label. E.g. if the
label array was [1,1,2,3] and the label 1 and 2 were mapped onto the target 3, the label array now is [3,3,3,3].
This method works for PointClouds and HybridClouds, not for more specific classes (HybridMesh is returned as
HybridCloud).
Args:
cloud: The PointCloud whose labels should get merged.
labels: A list of keys of the encoding dict of the PointCloud, or a list of actual labels which should get
mapped onto the target.
target: A key of the encoding dict of the PointCloud, or an actual label on which the labels should be mapped.
Returns:
A PointCloud where the labels were replaced by the target.
"""
mask = np.zeros(cloud.labels.shape, dtype=bool)
for label in labels:
if cloud.encoding is not None and label in cloud.encoding.keys():
label = cloud.encoding[label]
mask = np.logical_or(mask, cloud.labels == label)
else:
mask = np.logical_or(mask, cloud.labels == label)
if cloud.encoding is not None and target in cloud.encoding.keys():
target = cloud.encoding[target]
new_labels = cloud.labels.copy()
new_labels[mask] = target
if cloud.encoding is not None:
new_encoding = cloud.encoding.copy()
for label in labels:
new_encoding.pop(label, None)
else:
new_encoding = None
if isinstance(cloud, HybridCloud):
new_cloud = HybridCloud(cloud.nodes,
cloud.edges,
vertices=cloud.vertices,
labels=new_labels,
encoding=new_encoding)
else:
new_cloud = PointCloud(cloud.vertices,
labels=new_labels,
encoding=new_encoding)
return new_cloud
def voxel_down(ce: CloudEnsemble,
voxel_size: Union[dict, int] = 500) -> CloudEnsemble:
if type(voxel_size) is not dict:
voxel_size = dict(hc=voxel_size)
for k in ce.clouds:
voxel_size[k] = voxel_size['hc']
hc = ce.hc
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(hc.vertices)
pcd, idcs = pcd.voxel_down_sample_and_trace(voxel_size['hc'],
pcd.get_min_bound(),
pcd.get_max_bound())
idcs = np.max(idcs, axis=1)
new_labels = None
new_types = None
new_features = None
if len(hc.labels) != 0:
new_labels = hc.labels[idcs]
if len(hc.types) != 0:
new_types = hc.types[idcs]
if len(hc.features) != 0:
new_features = hc.features[idcs]
new_hc = HybridCloud(hc.nodes,
hc.edges,
vertices=np.asarray(pcd.points),
labels=new_labels,
types=new_types,
features=new_features,
encoding=hc.encoding,
node_labels=hc.node_labels,
no_pred=hc.no_pred)
new_clouds = {}
for key in ce.clouds:
pc = ce.clouds[key]
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(ce.clouds[key].vertices)
pcd, idcs = pcd.voxel_down_sample_and_trace(voxel_size[key],
pcd.get_min_bound(),
pcd.get_max_bound())
idcs = np.max(idcs, axis=1)
new_pc = PointCloud(np.asarray(pcd.points),
labels=pc.labels[idcs],
encoding=pc.encoding,
no_pred=pc.no_pred)
new_clouds[key] = new_pc
return CloudEnsemble(new_clouds, new_hc, no_pred=ce.no_pred)
def test_rotate_randomly():
angle_range = (60, 60)
angles = np.random.uniform(angle_range[0], angle_range[1], (1, 3))[0]
rot = Rot.from_euler('xyz', angles, degrees=True)
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
pc.rotate_randomly(angle_range)
hc.rotate_randomly(angle_range)
assert np.all(
pc.vertices == rot.apply(np.array([[10, 10, 10], [20, 20, 20]])))
assert np.all(
hc.vertices == rot.apply(np.array([[10, 10, 10], [20, 20, 20]])))
assert np.all(
hc.nodes == rot.apply(np.array([[10, 10, 10], [20, 20, 20]])))
pc = PointCloud(np.array([[10, 10, 10], [20, 20, 20]]))
hc = HybridCloud(np.array([[10, 10, 10], [20, 20, 20]]),
np.array([[0, 1]]),
vertices=np.array([[10, 10, 10], [20, 20, 20]]))
# test transformation class from processing.clouds
transform = clouds.RandomRotate(angle_range, apply_flip=False)
transform(pc)
transform(hc)
assert np.all(
pc.vertices == rot.apply(np.array([[10, 10, 10], [20, 20, 20]])))
assert np.all(
hc.vertices == rot.apply(np.array([[10, 10, 10], [20, 20, 20]])))
assert np.all(
hc.nodes == rot.apply(np.array([[10, 10, 10], [20, 20, 20]])))
def merge_clouds(
clouds: List[Union[PointCloud, HybridCloud]],
names: Optional[List[Union[str, int]]] = None,
ignore_hybrids: bool = False
) -> Optional[Union[PointCloud, HybridCloud]]:
""" Merges the PointCloud objects in the given list. If the names list is given, the object boundary information
is saved in the obj_bounds dict. Vertices of PointClouds without label / feature get the label /feature -1.
If no PointCloud has labels / features, then the label /feature array of the merged PointCloud is empty.
Args:
clouds: List of clouds which should get merged.
names: Names for each cloud in order to save object boundaries. This is only used if the clouds themselve have
no obj_bounds dicts.
ignore_hybrids: Flag for treating HybridClouds as sole PointClouds (ignoring nodes and edges).
Returns:
PointCloud which consists of the merged clouds.
"""
for cloud in clouds:
if cloud is None:
clouds.remove(cloud)
if names is not None:
if len(names) != len(clouds):
raise ValueError("Not enough names given.")
feats_dim = 1
if len(clouds[0].features) != 0:
feats_dim = clouds[0].features.shape[1]
for cloud in clouds:
if len(cloud.features) != 0:
if cloud.features.shape[1] != feats_dim:
raise ValueError(
"Feature dimensions of PointCloud do not fit. Clouds cannot be merged."
)
# find required size of new arrays
total_verts = 0
total_nodes = 0
total_edges = 0
for cloud in clouds:
total_verts += len(cloud.vertices)
if isinstance(
cloud, HybridCloud
) and cloud.nodes is not None and cloud.edges is not None:
total_nodes += len(cloud.nodes)
total_edges += len(cloud.edges)
# TODO: Generalize to support graph merging as well
if total_verts == 0:
return None
# reserve arrays of required size and initialize new attributes
t_verts = np.zeros((total_verts, 3))
t_labels = np.zeros((total_verts, 1))
t_features = np.zeros((total_verts, feats_dim))
t_types = np.zeros((total_verts, 1))
nodes = np.zeros((total_nodes, 3))
edges = np.zeros((total_edges, 2))
offset = 0
obj_bounds = {}
encoding = {}
no_pred = []
for ix, cloud in enumerate(clouds):
# handle hybrids
if not ignore_hybrids:
if isinstance(
cloud, HybridCloud
) and cloud.nodes is not None and cloud.edges is not None:
nodes[offset:offset + len(cloud.nodes)] = cloud.nodes
edges[offset:offset + len(cloud.edges)] = cloud.edges + offset
# handle pointclouds
t_verts[offset:offset + len(cloud.vertices)] = cloud.vertices
if len(cloud.labels) != 0:
t_labels[offset:offset + len(cloud.vertices)] = cloud.labels
else:
t_labels[offset:offset + len(cloud.vertices)] = -1
if len(cloud.features) != 0:
t_features[offset:offset + len(cloud.features)] = cloud.features
else:
t_features[offset:offset + len(cloud.features)] = -1
if len(cloud.types) != 0:
t_types[offset:offset + len(cloud.vertices)] = cloud.types
else:
t_types[offset:offset + len(cloud.vertices)] = -1
# TODO: Handle similar keys from different clouds and handle obj_bounds
# which don't span the entire vertex array
# Save object boundaries
if cloud.obj_bounds is not None:
for key in cloud.obj_bounds.keys():
obj_bounds[key] = cloud.obj_bounds[key] + offset
else:
if names is not None:
obj_bounds[names[ix]] = np.array(
[offset, offset + len(cloud.vertices)])
offset += len(cloud.vertices)
# Merge encodings
if cloud.encoding is not None:
for item in cloud.encoding:
encoding[item] = cloud.encoding[item]
# Merge no_preds
if cloud.no_pred is not None:
for item in cloud.no_pred:
no_pred.append(item)
if len(obj_bounds) == 0:
obj_bounds = None
if len(encoding) == 0:
encoding = None
if np.all(t_types == -1):
t_types = None
if np.all(t_labels == -1):
t_labels = None
if np.all(t_features == -1):
t_features = None
if np.all(nodes == 0):
nodes = None
if np.all(edges == 0):
edges = None
if ignore_hybrids:
return PointCloud(t_verts,
labels=t_labels,
features=t_features,
obj_bounds=obj_bounds,
encoding=encoding,
no_pred=no_pred,
types=t_types)
else:
return HybridCloud(nodes,
edges,
vertices=t_verts,
labels=t_labels,
features=t_features,
obj_bounds=obj_bounds,
encoding=encoding,
no_pred=no_pred,
types=t_types)
def test_sanity():
hc = HybridCloud()
assert len(hc.nodes) == 0
assert len(hc.node_labels) == 0
assert len(hc.pred_node_labels) == 0
assert len(hc.edges) == 0
def worker_split(id_queue: Queue,
chunk_queue: Queue,
ssd: SuperSegmentationDataset,
ctx: int,
base_node_dst: int,
parts: Dict[str, List[int]],
labels_itf: str,
label_mappings: List[Tuple[int, int]],
split_jitter: int = 0):
"""
Args:
id_queue: Input queue with cell ids.
chunk_queue: Output queue with cell chunks.
ssd: SuperSegmentationDataset which contains the cells to which the chunkhandler should get applied.
ctx: Context size for splitting.
base_node_dst: Distance between base nodes. Corresponds to redundancy / number of chunks per cell.
parts: Information about cell surface and organelles, Tuples like (voxel_param, feature) keyed by identifier
compatible with syconn (e.g. 'sv' or 'mi').
labels_itf: Label identifier for existing label predictions within the sso objects of the ssd dataset.
label_mappings: Tuples where label at index 0 should get mapped to label at index 1.
split_jitter: Derivation from context size during splitting.
"""
while True:
if not id_queue.empty():
ssv_id = id_queue.get()
sso = ssd.get_super_segmentation_object(ssv_id)
vert_dc = {}
label_dc = {}
encoding = {}
offset = 0
obj_bounds = {}
for ix, k in enumerate(parts):
pcd = o3d.geometry.PointCloud()
verts = sso.load_mesh(k)[1].reshape(-1, 3)
pcd.points = o3d.utility.Vector3dVector(verts)
pcd, idcs = pcd.voxel_down_sample_and_trace(
parts[k][0], pcd.get_min_bound(), pcd.get_max_bound())
idcs = np.max(idcs, axis=1)
vert_dc[k] = np.asarray(pcd.points)
obj_bounds[k] = [offset, offset + len(pcd.points)]
offset += len(pcd.points)
if k == 'sv':
# prepare mask for filtering background / unpredicted points
mask = None
if labels_itf == 'axoness':
# 0: dendrite, 1: axon, 2: soma, 3: bouton, 4: terminal, 5/6: background/unpredicted
labels_total = sso.label_dict()[labels_itf][idcs]
mask = labels_total < 5
labels_total = labels_total[mask]
elif labels_itf == 'spiness':
# 1: head, 0: neck, 2: shaft, 3: other, 4/5: background/unpredicted
labels_total = sso.label_dict()['axoness'][idcs]
spiness = sso.label_dict()['spiness'][idcs]
mask = np.logical_not(
np.logical_or(
labels_total > 4,
np.logical_and(labels_total == 0,
spiness > 3)))
labels_total = labels_total[mask]
spiness = spiness[mask]
labels_total[labels_total != 0] = 3
labels_total[labels_total == 0] = spiness[labels_total
== 0]
else:
labels_total = sso.label_dict()[labels_itf][idcs]
mask = np.ones(len(labels_total)).astype(bool)
labels = labels_total
vert_dc[k] = vert_dc[k][mask]
else:
labels = np.ones(len(vert_dc[k])) + ix + labels_total.max()
encoding[k] = ix + 1 + labels_total.max()
label_dc[k] = labels
sample_feats = np.concatenate([[parts[k][1]] * len(vert_dc[k])
for k in parts]).reshape(-1, 1)
sample_feats = label_binarize(sample_feats,
classes=np.arange(len(parts)))
sample_pts = np.concatenate([vert_dc[k] for k in parts])
sample_labels = np.concatenate([label_dc[k] for k in parts])
# mark cellular organelles to be excluded from loss calculation - see torchhandler for use of no_pred
no_pred = list(encoding.keys())
if not sso.load_skeleton():
raise ValueError(f'Couldnt find skeleton of {sso}')
nodes, edges = sso.skeleton['nodes'] * sso.scaling, sso.skeleton[
'edges']
hc = HybridCloud(nodes,
edges,
vertices=sample_pts,
features=sample_feats,
obj_bounds=obj_bounds,
no_pred=no_pred,
labels=sample_labels,
encoding=encoding)
if label_mappings is not None:
hc.map_labels(label_mappings)
_ = hc.verts2node
jitter = random.randint(0, split_jitter)
node_arrs, source_nodes = splitting.split_single(
hc, ctx + jitter, base_node_dst)
for ix, node_arr in enumerate(node_arrs):
sample, _ = objects.extract_cloud_subset(hc, node_arr)
chunk_queue.put(sample)
else:
time.sleep(0.5)
def hybridmesh2poisson(hm: HybridMesh, tech_density: int, obj_factor: float) -> PointCloud:
""" If there is a skeleton, it gets split into chunks of approximately equal size. For each chunk
the corresponding mesh piece gets extracted and gets sampled according to its area. If there
is no skeleton, the mesh is split into multiple parts, depending on its area. Each part is
then again sampled based on its area.
Args:
hm: HybridMesh which should be transformed into a HybridCloud with poisson disk sampled points.
tech_density: poisson sampling density in point/um². With tech_density = -1, the number of sampled points
equals the number of vertices in the given HybridMesh.
"""
if len(hm.nodes) == 0:
offset = 0
mesh = trimesh.Trimesh(vertices=hm.vertices, faces=hm.faces)
area = mesh.area * 1e-06
# number of chunks should be relative to area
chunk_number = round(area / 6)
if area == 0 or chunk_number == 0:
return PointCloud()
total = None
for i in tqdm(range(int(chunk_number))):
# process all faces left with last chunk
if i == chunk_number-1:
chunk_faces = hm.faces[offset:]
else:
chunk_faces = hm.faces[offset:offset + floor(len(hm.faces) // chunk_number)]
offset += floor(len(hm.faces) // chunk_number)
chunk_hm = HybridMesh(vertices=hm.vertices, faces=chunk_faces, labels=hm.labels, types=hm.types)
mesh = trimesh.Trimesh(vertices=chunk_hm.vertices, faces=chunk_hm.faces)
area = mesh.area * 1e-06
if tech_density == -1:
pc = meshes.sample_mesh_poisson_disk(chunk_hm, int(len(chunk_hm.vertices) * obj_factor))
else:
pc = meshes.sample_mesh_poisson_disk(chunk_hm, tech_density * area * obj_factor)
if total is None:
total = pc
else:
total = clouds.merge_clouds([total, pc])
result = PointCloud(vertices=total.vertices, labels=total.labels, encoding=hm.encoding, no_pred=hm.no_pred,
types=total.types)
else:
total = None
intermediate = None
context_size = 5
skel2node_mapping = True
counter = 0
chunks = graphs.bfs_iterative(hm.graph(), 0, context_size)
for chunk in tqdm(chunks):
chunk = np.array(chunk)
# At the first iteration the face2node mapping must be done
if skel2node_mapping:
print("Mapping faces to node for further processing. This might take a while...")
skel2node_mapping = False
extract = hybrids.extract_mesh_subset(hm, chunk)
if len(hm.faces) == 0:
continue
# get the mesh area in trimesh units and use it to determine how many points should be sampled
mesh = trimesh.Trimesh(vertices=extract.vertices, faces=extract.faces)
area = mesh.area * 1e-06
if area == 0:
continue
else:
if tech_density == -1:
pc = meshes.sample_mesh_poisson_disk(extract, len(extract.vertices))
else:
pc = meshes.sample_mesh_poisson_disk(extract, tech_density * area)
if intermediate is None:
intermediate = pc
else:
intermediate = clouds.merge_clouds([intermediate, pc])
# merging slows down process => hold speed constant by reducing merging operations
counter += 1
if counter % 50 == 0:
if total is None:
total = intermediate
else:
total = clouds.merge_clouds(([total, intermediate]))
intermediate = None
total = clouds.merge_clouds([total, intermediate])
result = HybridCloud(nodes=hm.nodes, edges=hm.edges, vertices=total.vertices, labels=total.labels,
encoding=hm.encoding, no_pred=hm.no_pred, types=total.types)
return result
