Source code for syconn.proc.rendering

# -*- coding: utf-8 -*-
# SyConn - Synaptic connectivity inference toolkit
#
# Copyright (c) 2016 - now
# Max Planck Institute of Neurobiology, Martinsried, Germany
# Authors: Philipp Schubert, Sven Dorkenwald, Joergen Kornfeld

from typing import Union, List, Callable, Optional, Tuple, TYPE_CHECKING, Iterable

from . import log_proc
from .meshes import MeshObject, calc_rot_matrices
from .. import global_params
from ..backend.storage import CompressedStorage
from ..handler.basics import flatten_list
from ..handler.compression import arrtolz4string
from ..handler.multiviews import generate_palette, remap_rgb_labelviews, \
    rgb2id_array, rgba2id_array, id2rgba_array_contiguous
from ..mp.mp_utils import start_multiprocess_imap

if TYPE_CHECKING:
    from ..reps.super_segmentation import SuperSegmentationObject
    from ..reps.segmentation import SegmentationObject
import time
import os
import sys
import numpy as np


[docs]def load_rendering_func(func_name: str) -> Callable: # can't load more than one platform simultaneously os.environ['PYOPENGL_PLATFORM'] = global_params.config['pyopengl_platform'] if global_params.config['pyopengl_platform'] == 'egl': try: try: import OpenGL.EGL except AttributeError: # OSMesa has been enabled before # hacky, but successfully removes all OpenGL related imports for k in list(sys.modules.keys()): if 'OpenGL' in k: del sys.modules[k] import OpenGL.EGL from ..proc import rendering_egl as rendering_module except ImportError as e: os.environ['PYOPENGL_PLATFORM'] = 'osmesa' global_params.config['pyopengl_platform'] = 'osmesa' try: import OpenGL.osmesa except AttributeError: # OSMesa has been enabled before # hacky, but successfully removes all OpenGL related imports for k in list(sys.modules.keys()): if 'OpenGL' in k: del sys.modules[k] import OpenGL.osmesa from ..proc import rendering_osmesa as rendering_module log_proc.error('EGL requirements could not be imported ({}). ' 'Switched to OSMESA platform.'.format(e)) elif global_params.config['pyopengl_platform'] == 'osmesa': try: import OpenGL.osmesa except AttributeError: # OSMesa has been enabled before # hacky, but successfully removes all OpenGL related imports for k in list(sys.modules.keys()): if 'OpenGL' in k: del sys.modules[k] import OpenGL.osmesa from ..proc import rendering_osmesa as rendering_module else: msg = 'PYOpenGL environment has to be "egl" or "osmesa".' log_proc.error(msg) raise NotImplementedError(msg) return getattr(rendering_module, func_name)
[docs]def render_mesh(mo: MeshObject, **kwargs) -> np.ndarray: """ Render super voxel raw views located at randomly chosen center of masses in vertex cloud. Args: mo: Mesh. **kwargs: Keyword arguments pass to :py:func:`~multi_view_mesh` call. Returns: View array. """ multi_view_mesh = load_rendering_func('multi_view_mesh') if "physical_scale" in kwargs.keys(): kwargs["physical_scale"] = kwargs["physical_scale"] / mo.max_dist mo_views = multi_view_mesh(mo.indices, mo.vertices, mo.normals, colors=mo.colors, **kwargs) return mo_views
[docs]def render_mesh_coords(coords: np.ndarray, ind: np.ndarray, vert: np.ndarray, **kwargs) -> np.ndarray: """ Render raw views located at given coordinates in mesh Returns ViewContainer list if dest_dir is None, else writes views to dest_path. Args: coords: Rendering locations. ind: Mesh indices/faces [N, 1] vert: Mesh vertices [M, 1] **kwargs: Keyword arguments passed to :py:func:`_render_mesh_coords`. Returns: Views at each coordinate. """ _render_mesh_coords = load_rendering_func('_render_mesh_coords') mesh = MeshObject("views", ind, vert) mesh._colors = None # this enables backwards compatibility, check why this was used return _render_mesh_coords(coords, mesh, **kwargs)
# ------------------------------------------------------------------------------ # SSO rendering code
[docs]def render_sampled_sso(sso: 'SuperSegmentationObject', ws: Optional[Tuple[int, int]] = None, verbose: bool = False, woglia: bool = True, return_rot_mat: bool = False, overwrite: bool = True, add_cellobjects: Optional[Union[bool, Iterable[str]]] = None, index_views: bool = False, return_views: bool = False, cellobjects_only: bool = False, rot_mat: Optional[np.ndarray] = None, view_key: Optional[str] = None) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray], None]: """ Renders for each SV views at sampled locations (number is dependent on SV mesh size with scaling fact) from combined mesh of all SV. Args: sso: SuperSegmentationObject ws: Window size in pixels (y, x). Default: See config.yml or custom configs in the working directory. verbose: Log additional information. woglia: Store views with "without glia" identifier, i.e. flag the views as being created after the glia separation. return_rot_mat: Return rotation matrices. add_cellobjects: Default: ('sj', 'vc', 'mi') overwrite: If True, do not skip existing views. index_views: Also render index views. return_views: Return view arrays. cellobjects_only: Only render cell objects. rot_mat: Rotation matrix array for every rendering location [N, 4, 4]. view_key: String identifier for storing view arrays. Only needed if ``return_views=False``. Returns: Depending on `return_views` and `return_rot_mat`: None; View array; View array and rotation matrices; rotation matrices. """ view_cfg = global_params.config['views'] view_props_default = view_cfg['view_properties'] if ws is None: ws = view_props_default['ws'] if add_cellobjects is None or add_cellobjects is True: add_cellobjects = view_cfg['subcell_objects'] if view_cfg['use_onthefly_views'] and 'sj' in add_cellobjects: add_cellobjects[add_cellobjects.index('sj')] = 'syn_ssv' # get coordinates for N SV's in SSO coords = sso.sample_locations(cache=False) if not overwrite: missing_sv_ixs = ~np.array(sso.view_existence( woglia=woglia, index_views=index_views, view_key=view_key), dtype=np.bool) missing_svs = np.array(sso.svs)[missing_sv_ixs] coords = np.array(coords)[missing_sv_ixs] else: missing_svs = np.array(sso.svs) if len(missing_svs) == 0: if return_views: return sso.load_views(woglia=woglia) return # len(part_views) == N + 1 part_views = np.cumsum([0] + [len(c) for c in coords]) flat_coords = np.array(flatten_list(coords)) if verbose: start = time.time() if index_views: views, rot_mat = render_sso_coords_index_views( sso, flat_coords, ws=ws, return_rot_matrices=True, verbose=verbose, rot_mat=rot_mat) else: views, rot_mat = render_sso_coords( sso, flat_coords, ws=ws, verbose=verbose, add_cellobjects=add_cellobjects, return_rot_mat=True, cellobjects_only=cellobjects_only, rot_mat=rot_mat) if verbose: dur = time.time() - start log_proc.debug("Rendering of %d views took %0.2fs. " "%0.4fs/SV" % (len(views), dur, float(dur) / len(sso.svs))) if not return_views: if cellobjects_only: log_proc.warning('`cellobjects_only=True` in `render_sampled_sso` call, views ' 'will be written to file system in serial (this is slow).') for i, so in enumerate(missing_svs): sv_views = views[part_views[i]:part_views[i + 1]] so.save_views(sv_views, woglia=woglia, cellobjects_only=cellobjects_only, index_views=index_views, enable_locking=True, view_key=view_key) else: write_sv_views_chunked(missing_svs, views, part_views, dict(woglia=woglia, index_views=index_views, view_key=view_key)) if return_views: if return_rot_mat: return views, rot_mat return views if return_rot_mat: return rot_mat
[docs]def render_sso_coords(sso: 'SuperSegmentationObject', coords: np.ndarray, add_cellobjects: Optional[Union[bool, Iterable[str]]] = None, verbose: bool = False, clahe: bool = False, ws: Optional[Tuple[int]] = None, cellobjects_only: bool = False, wire_frame: bool = False, nb_views: Optional[int] = None, comp_window: Optional[float] = None, rot_mat: Optional[np.ndarray] = None, return_rot_mat: bool = False) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ Render views of SuperSegmentationObject at given coordinates. Args: sso: SuperSegmentationObject coords: Rendering locations [N, 3]. add_cellobjects: Default: ('sj', 'vc', 'mi'). This ordering determines the channel order of the view array. verbose: Log additional information. clahe: Use clahe to enahnce view contrast. ws: Window size in pixels (y, x). Default: See config.yml or custom configs in the working directory. cellobjects_only: Only render cell objects. wire_frame: Render the mesh as a wire frame. nb_views: Number of views. Default: See config.yml or custom configs in the working directory. comp_window: Window size in nm. the clipping box during rendering will have an extent of [comp_window, comp_window / 2, comp_window]. Default: 8 um. rot_mat: Rotation matrix array for every rendering location [N, 4, 4]. return_rot_mat: Return rotation matrices, e.g. if not provided via `rot_mat`, this output can be provided for other rendering calls. Returns: Resulting views rendered at each location. Output shape: [len(coords), 4 (default: cell outline + number of cell objects), nb_views, y, x]. """ view_cfg = global_params.config['views'] view_props_default = view_cfg['view_properties'] if comp_window is None: comp_window = view_props_default['comp_window'] if ws is None: ws = view_props_default['ws'] if add_cellobjects is None or add_cellobjects is True: add_cellobjects = view_cfg['subcell_objects'] if view_cfg['use_onthefly_views'] and 'sj' in add_cellobjects: add_cellobjects[add_cellobjects.index('sj')] = 'syn_ssv' if verbose: log_proc.debug('Started "render_sso_coords" at {} locations with sub-cellular structures {} for SSO {} using ' 'PyOpenGL platform "{}".'.format(len(coords), add_cellobjects, sso, global_params.config['pyopengl_platform'])) start = time.time() if nb_views is None: nb_views = view_props_default['nb_views'] mesh = sso.mesh if verbose: log_proc.debug(f'Loaded cell mesh after {(time.time() - start):.2f} s.') if cellobjects_only: assert len(add_cellobjects) > 0, "Add cellobjects must contain at least one entry " \ "when rendering cellobjects only." raw_views = np.ones((len(coords), nb_views, ws[0], ws[1]), dtype=np.uint8) * 255 if rot_mat is None: mo = MeshObject("raw", mesh[0], mesh[1]) mo._colors = None querybox_edgelength = comp_window / mo.max_dist rot_mat = calc_rot_matrices(mo.transform_external_coords(coords), mo.vert_resh, querybox_edgelength, nb_cpus=sso.nb_cpus) else: if len(mesh[1]) == 0 or len(coords) == 0: raw_views = np.ones((len(coords), nb_views, ws[0], ws[1]), dtype=np.uint8) * 255 msg = "No mesh for SSO {} found with {} locations.".format(sso, len(coords)) log_proc.warning(msg) else: raw_views, rot_mat = render_mesh_coords( coords, mesh[0], mesh[1], clahe=clahe, verbose=verbose, return_rot_matrices=True, ws=ws, wire_frame=wire_frame, rot_matrices=rot_mat, nb_views=nb_views, comp_window=comp_window) if add_cellobjects is not False and len(add_cellobjects) > 0: if cellobjects_only: res = [] else: res = [raw_views[:, None]] for subcell_obj in add_cellobjects: mesh = sso.load_mesh(subcell_obj) if len(mesh[1]) != 0: views = render_mesh_coords( coords, mesh[0], mesh[1], clahe=clahe, verbose=verbose, rot_matrices=rot_mat, views_key=subcell_obj, ws=ws, nb_views=nb_views, wire_frame=wire_frame, comp_window=comp_window) else: views = np.ones_like(raw_views) * 255 res.append(views[:, None]) res = np.concatenate(res, axis=1) if return_rot_mat: return res, rot_mat return res if return_rot_mat: return raw_views[:, None], rot_mat return raw_views[:, None]
[docs]def render_sso_coords_index_views(sso: 'SuperSegmentationObject', coords: np.ndarray, verbose: bool = False, ws: Optional[Tuple[int, int]] = None, rot_mat: Optional[np.ndarray] = None, nb_views: Optional[int] = None, comp_window: Optional[float] = None, return_rot_matrices: bool = False ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ Uses per-face color via flattened vertices (i.e. vert[ind] -> slow!). This was added to be able to calculate the surface coverage captured by the views. Todo: * Add fast GL_POINT rendering to omit slow per-face coloring (redundant vertices) and expensive remapping from face IDs to vertex IDs. Args: sso: SuperSegmentationObject coords: Rendering locations [N, 3]. verbose: Log additional information. ws: Window size in pixels (y, x). Default: See config.yml or custom configs in the working directory. rot_mat: np.array nb_views: Number of views. Default: See config.yml or custom configs in the working directory. comp_window: Window size in nm. the clipping box during rendering will have an extent of [comp_window, comp_window / 2, comp_window]. Default: 8 um. return_rot_matrices: Return rotation matrices, e.g. if not provided via `rot_mat`, this output can be provided for other rendering calls. Returns: Resulting index views rendered at each location. Output shape: [len(coords), 1, nb_views, y, x]. """ view_props_default = global_params.config['views']['view_properties'] _render_mesh_coords = load_rendering_func('_render_mesh_coords') if comp_window is None: comp_window = view_props_default['comp_window'] if ws is None: ws = view_props_default['ws'] if verbose: log_proc.debug('Started "render_sso_coords_index_views" at {} locations for SSO {} using ' 'PyOpenGL platform "{}".'.format(len(coords), sso.id, global_params.config['pyopengl_platform'])) if nb_views is None: nb_views = view_props_default['nb_views'] ind, vert, norm = sso.mesh if len(vert) == 0 or len(coords) == 0: msg = "No mesh for SSO {} found with {} locations.".format(sso, len(coords)) log_proc.critical(msg) res = np.ones((len(coords), nb_views, ws[1], ws[0], 3), dtype=np.uint8) * 255 if not return_rot_matrices: return res else: return np.zeros((len(coords), 16), dtype=np.uint8), res try: color_array = id2rgba_array_contiguous(np.arange(len(vert) // 3)) except ValueError as e: msg = "'render_sso_coords_index_views' failed with {} when " \ "rendering SSV {}.".format(e, sso.id) log_proc.error(msg) raise ValueError(msg) if color_array.shape[1] == 3: # add alpha channel color_array = np.concatenate([color_array, np.ones((len(color_array), 1), dtype=np.uint8) * 255], axis=-1) color_array = color_array.astype(np.float32) / 255. # in init it seems color values have to be normalized, check problems with uniqueness if # they are normalized between 0 and 1.. OR check if it is possible to just switch color arrays to UINT8 -> Check # backwards compatibility with other color-dependent rendering methods # Create mesh object without redundant vertices to get same PCA rotation as for raw views if rot_mat is None: mo = MeshObject("raw", ind, vert, color=color_array, normals=norm) querybox_edgelength = comp_window / mo.max_dist rot_mat = calc_rot_matrices(mo.transform_external_coords(coords), mo.vert_resh, querybox_edgelength, nb_cpus=sso.nb_cpus) # create redundant vertices to enable per-face colors # vert = vert.reshape(-1, 3)[ind].flatten() # ind = np.arange(len(vert) // 3) # color_array = np.repeat(color_array, 3, axis=0) # 3 <- triangles mo = MeshObject("raw", ind, vert, color=color_array, normals=norm) if return_rot_matrices: ix_views, rot_mat = _render_mesh_coords( coords, mo, verbose=verbose, ws=ws, depth_map=False, rot_matrices=rot_mat, smooth_shade=False, views_key="index", nb_views=nb_views, comp_window=comp_window, return_rot_matrices=return_rot_matrices) if ix_views.shape[-1] == 3: # rgba rendering ix_views = rgb2id_array(ix_views)[:, None] else: # rgba rendering ix_views = rgba2id_array(ix_views)[:, None] return ix_views, rot_mat ix_views = _render_mesh_coords(coords, mo, verbose=verbose, ws=ws, depth_map=False, rot_matrices=rot_mat, smooth_shade=False, views_key="index", nb_views=nb_views, comp_window=comp_window, return_rot_matrices=return_rot_matrices) if ix_views.shape[-1] == 3: ix_views = rgb2id_array(ix_views)[:, None] else: ix_views = rgba2id_array(ix_views)[:, None] scnd_largest = np.partition(np.unique(ix_views.flatten()), -2)[-2] # largest value is background if scnd_largest > len(vert) // 3: log_proc.critical('Critical error during index-rendering: Maximum vertex' ' ID which was rendered is bigger than vertex array.' '{}, {}; SSV ID {}'.format(scnd_largest, len(vert) // 3, sso.id)) return ix_views
[docs]def render_sso_coords_label_views(sso: 'SuperSegmentationObject', vertex_labels: np.ndarray, coords: np.ndarray, verbose: bool = False, ws: Optional[Tuple[int, int]] = None, rot_mat: Optional[np.ndarray] = None, nb_views: Optional[int] = None, comp_window: Optional[float] = None, return_rot_matrices: bool = False ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ Render views with vertex colors corresponding to vertex labels. Args: sso: vertex_labels: np.array vertex labels [N, 1]. Ordering and length have to be the same as vertex array of SuperSegmentationObject (len(sso.mesh[1]) // 3). coords: Rendering locations [N, 3]. verbose: Log additional information. ws: Window size in pixels (y, x). Default: See config.yml or custom configs in the working directory. rot_mat: np.array nb_views: Number of views. Default: See config.yml or custom configs in the working directory. comp_window: Window size in nm. the clipping box during rendering will have an extent of [comp_window, comp_window / 2, comp_window]. Default: 8 um. return_rot_matrices: Return rotation matrices, e.g. if not provided via `rot_mat`, this output can be provided for other rendering calls. Returns: Resulting label views rendered at each location. Output shape: [len(coords), 1, nb_views, y, x]. """ view_props_default = global_params.config['views']['view_properties'] _render_mesh_coords = load_rendering_func('_render_mesh_coords') if comp_window is None: comp_window = view_props_default['comp_window'] if ws is None: ws = view_props_default['ws'] if nb_views is None: nb_views = view_props_default['nb_views'] ind, vert, _ = sso.mesh if len(vertex_labels) != len(vert) // 3: raise ValueError("Length of vertex labels and vertices have to be equal.") palette = generate_palette(len(np.unique(vertex_labels))) color_array = palette[vertex_labels].astype(np.float32) / 255 mo = MeshObject("neuron", ind, vert, color=color_array) label_views, rot_mat = _render_mesh_coords(coords, mo, depth_map=False, ws=ws, rot_matrices=rot_mat, nb_views=nb_views, smooth_shade=False, verbose=verbose, comp_window=comp_window, return_rot_matrices=True) label_views = remap_rgb_labelviews(label_views, palette)[:, None] if return_rot_matrices: return label_views, rot_mat return label_views
[docs]def get_sso_view_dc(sso: 'SuperSegmentationObject', verbose: bool = False) -> dict: """ Extracts views from sampled positions in SSO for each SV. Args: sso: SuperSegmentationObject verbose: bool Returns: Dictionary with `sso` id as key and lz4 compressed view array (see `:py:func:~syconn.handler.compression.arrtolz4string`). """ views = render_sampled_sso(sso, verbose=verbose, return_views=True) view_dc = {sso.id: arrtolz4string(views)} return view_dc
[docs]def render_sso_coords_multiprocessing(ssv: 'SuperSegmentationObject', n_jobs: int, rendering_locations: Optional[np.ndarray] = None, verbose: bool = False, render_kwargs: Optional[dict] = None, view_key: Optional[str] = None, render_indexviews: bool = True, return_views: bool = True) -> Union[None, np.ndarray]: """ Args: ssv: SuperSegmentationObject n_jobs: int number of parallel jobs running on same node of cluster rendering_locations: array of locations to be rendered if not given, rendering locations are retrieved from the SSV's SVs. Results will be stored at SV locations. verbose: bool flag to show the progress of rendering. render_kwargs: dict view_key: str render_indexviews: bool return_views: If False and rendering_locations is None, views will be saved on supervoxel level. Returns: Array of views after rendering of locations or None. """ if rendering_locations is not None and return_views is False: raise ValueError('"render_sso_coords_multiprocessing" received invalid ' 'parameters (`rendering_locations!=None` and `return_v' 'iews=False`). When using specific rendering locations, ' 'views have to be returned-') svs = None ssv.nb_cpus = n_jobs if rendering_locations is None: # use SV rendering locations svs = list(ssv.svs) if ssv._sample_locations is None and not ssv.attr_exists("sample_locations"): ssv.load_attr_dict() rendering_locations = ssv.sample_locations(cache=False) # store number of rendering locations per SV -> Ordering of rendered # views must be preserved! part_views = np.cumsum([0] + [len(c) for c in rendering_locations]) rendering_locations = np.concatenate(rendering_locations) if len(rendering_locations) == 0: log_proc.critical('No rendering locations found for SSV {}.'.format(ssv.id)) return np.ones((0, global_params.config['views']['view_properties']['nb_views'], 256, 128), dtype=np.uint8) * 255 params = np.array_split(rendering_locations, n_jobs) # TODO: refactor kwargs! render_kwargs_def = {'add_cellobjects': True, 'verbose': verbose, 'clahe': False, 'ws': None, 'cellobjects_only': False, 'wire_frame': False, 'nb_views': None, 'comp_window': None, 'rot_mat': None, 'woglia': True, 'return_rot_mat': False, 'render_indexviews': render_indexviews} if render_kwargs is not None: render_kwargs_def.update(render_kwargs) _ = ssv.mesh # cache mesh params = [[par, ssv, render_kwargs_def] for par in params] views = start_multiprocess_imap(_render_views_multiproc, params) views = np.concatenate(views) if svs is not None and return_views is False: start_writing = time.time() if render_kwargs_def['cellobjects_only']: log_proc.warning('`cellobjects_only=True` in `render_sampled_sso` call, views ' 'will be written to file system in serial (this is slow).') for i, so in enumerate(svs): so.enable_locking = True sv_views = views[part_views[i]:part_views[i + 1]] so.save_views(sv_views, woglia=render_kwargs_def['woglia'], cellobjects_only=render_kwargs_def['cellobjects_only'], index_views=render_kwargs_def["render_indexviews"], enable_locking=True) else: write_sv_views_chunked(svs, views, part_views, dict(woglia=render_kwargs_def['woglia'], index_views=render_kwargs_def["render_indexviews"], view_key=view_key)) end_writing = time.time() log_proc.debug('Writing SV renderings took {:.2f}s'.format( end_writing - start_writing)) return return views
def _render_views_multiproc(args: tuple) -> np.ndarray: coords, sso, kwargs = args render_indexviews = kwargs['render_indexviews'] del kwargs['render_indexviews'] # TODO: refactor kwargs if 'overwrite' in kwargs: del kwargs['overwrite'] if render_indexviews: if 'add_cellobjects' in kwargs: del kwargs['add_cellobjects'] if 'clahe' in kwargs: del kwargs['clahe'] if 'wire_frame' in kwargs: del kwargs['wire_frame'] if 'cellobjects_only' in kwargs: del kwargs['cellobjects_only'] if 'return_rot_mat' in kwargs: del kwargs['return_rot_mat'] if 'woglia' in kwargs: del kwargs['woglia'] views = render_sso_coords_index_views(sso, coords, **kwargs) else: if 'woglia' in kwargs: del kwargs['woglia'] views = render_sso_coords(sso, coords, **kwargs) return views
[docs]def write_sv_views_chunked(svs: List['SegmentationObject'], views: np.ndarray, part_views: np.ndarray, view_kwargs: dict, disable_locking: bool = False): """ Args: svs: SegmentationObjects views: View array. part_views: Cumulated number of views -> indices of start and end of SV views in `views` array. view_kwargs: disable_locking: Usually required as SVs are stored in chunks and rendered distributed on many processes. Returns: """ view_dc = {} for sv_ix, sv in enumerate(svs): curr_view_dest = sv.view_path(**view_kwargs) view_ixs = part_views[sv_ix], part_views[sv_ix + 1] if curr_view_dest in view_dc: view_dc[curr_view_dest][sv.id] = view_ixs else: view_dc[curr_view_dest] = {sv.id: view_ixs} for k, v in view_dc.items(): view_storage = CompressedStorage(k, read_only=False, disable_locking=disable_locking) for sv_id, sv_view_ixs in v.items(): view_storage[sv_id] = views[sv_view_ixs[0]:sv_view_ixs[1]] view_storage.push() del view_storage