def voxelize(pts, polys, shape=(256, 256, 256), center=(128, 128, 128), mp=True): from tvtk.api import tvtk pd = tvtk.PolyData(points=pts + center + (0, 0, 0), polys=polys) plane = tvtk.Planes(normals=[(0,0,1)], points=[(0,0,0)]) clip = tvtk.ClipPolyData(clip_function=plane, input=pd) feats = tvtk.FeatureEdges( manifold_edges=False, non_manifold_edges=False, feature_edges=False, boundary_edges=True, input=clip.output) def func(i): plane.points = [(0,0,i)] feats.update() vox = np.zeros(shape[:2][::-1], np.uint8) if feats.output.number_of_lines > 0: epts = feats.output.points.to_array() edges = feats.output.lines.to_array().reshape(-1, 3)[:,1:] for poly in trace_poly(edges): vox += rasterize(epts[poly][:,:2]+[.5, .5], shape=shape[:2][::-1]) return vox % 2 if mp: from . import mp layers = mp.map(func, range(shape[2])) else: #layers = map(func, range(shape[2])) layers = [func(x) for x in range(shape[2])] # python3 compatible return np.array(layers).T
def _getmask(pia, wm, polys, shape): from .. import polyutils mask = sparse.csr_matrix((len(wm), np.prod(shape))) from tvtk.api import tvtk measure = tvtk.MassProperties() planes = tvtk.PlaneCollection() for norm in np.vstack([-np.eye(3), np.eye(3)]): planes.append(tvtk.Plane(normal=norm)) ccs = tvtk.ClipClosedSurface(clipping_planes=planes) feats = tvtk.FeatureEdges(boundary_edges=1, non_manifold_edges=0, manifold_edges=0, feature_edges=0) feats.set_input(ccs.output) surf = polyutils.Surface(pia, polys) for i, (pts, faces) in enumerate(surf.polyhedra(wm)): if len(pts) > 0: poly = tvtk.PolyData(points=pts, polys=faces) measure.set_input(poly) measure.update() totalvol = measure.volume ccs.set_input(poly) measure.set_input(ccs.output) bmin = pts.min(0).round().astype(int) bmax = (pts.max(0).round() + 1).astype(int) vidx = np.mgrid[bmin[0]:bmax[0], bmin[1]:bmax[1], bmin[2]:bmax[2]] for vox in vidx.reshape(3, -1).T: try: idx = np.ravel_multi_index(vox[::-1], shape) for plane, m in zip(planes, [.5, .5, .5, -.5, -.5, -.5]): plane.origin = vox + m ccs.update() if ccs.output.number_of_cells > 2: measure.update() mask[i, idx] = measure.volume except ValueError: print('Voxel not in volume: (%d, %d, %d)' % tuple(vox)) mask.data[mask.indptr[i]:mask.indptr[i + 1]] /= mask[i].sum() return mask