Source code for

#    This script is part of navis (
#    Copyright (C) 2018 Philipp Schlegel
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    GNU General Public License for more details.

import os

import multiprocessing as mp
import trimesh as tm

from pathlib import Path
from typing import Union, Iterable, Optional, Dict, Any
from typing_extensions import Literal

from .. import config, utils, core
from . import base

# Set up logging
logger = config.get_logger(__name__)

[docs] def read_mesh(f: Union[str, Iterable], include_subdirs: bool = False, parallel: Union[bool, int] = 'auto', output: Union[Literal['neuron'], Literal['volume'], Literal['trimesh']] = 'neuron', errors: Union[Literal['raise'], Literal['log'], Literal['ignore']] = 'log', limit: Optional[int] = None, **kwargs) -> 'core.NeuronObject': """Create Neuron/List from mesh. This is a thin wrapper around `trimesh.load_mesh` which supports most common formats (obj, ply, stl, etc.). Parameters ---------- f : str | iterable Filename(s) or folder. If folder must include file extension (e.g. `my/dir/*.ply`). include_subdirs : bool, optional If True and ``f`` is a folder, will also search subdirectories for meshes. parallel : "auto" | bool | int, Defaults to ``auto`` which means only use parallel processing if more than 100 mesh files are imported. Spawning and joining processes causes overhead and is considerably slower for imports of small numbers of neurons. Integer will be interpreted as the number of cores (otherwise defaults to ``os.cpu_count() - 2``). output : "neuron" | "volume" | "trimesh" Determines function's output. See Returns. errors : "raise" | "log" | "ignore" If "log" or "ignore", errors will not be raised. limit : int, optional If reading from a folder you can use this parameter to read only the first ``limit`` files. Useful when wanting to get a sample from a large library of meshes. **kwargs Keyword arguments passed to :class:`navis.MeshNeuron` or :class:`navis.Volume`. You can use this to e.g. set the units on the neurons. Returns ------- navis.MeshNeuron If ``output="neuron"`` (default). navis.Volume If ``output="volume"``. trimesh.Trimesh If ``output='trimesh'``. navis.NeuronList If ``output="neuron"`` and import has multiple meshes will return NeuronList of MeshNeurons. list If ``output!="neuron"`` and import has multiple meshes will return list of Volumes or Trimesh. Examples -------- Read a single file into :class:`navis.MeshNeuron`: >>> m = navis.read_mesh('mesh.obj') # doctest: +SKIP Read all e.g. .obj files in a directory: >>> nl = navis.read_mesh('/some/directory/*.obj') # doctest: +SKIP Sample first 50 files in folder: >>> nl = navis.read_mesh('/some/directory/*.obj', limit=50) # doctest: +SKIP Read single file into :class:`navis.Volume`: >>> nl = navis.read_mesh('mesh.obj', output='volume') # doctest: +SKIP """ utils.eval_param(output, name='output', allowed_values=('neuron', 'volume', 'trimesh')) # If is directory, compile list of filenames if isinstance(f, str) and '*' in f: f, ext = f.split('*') f = Path(f).expanduser() if not f.is_dir(): raise ValueError(f'{f} does not appear to exist') if not include_subdirs: f = list(f.glob(f'*{ext}')) else: f = list(f.rglob(f'*{ext}')) if limit: f = f[:limit] if utils.is_iterable(f): # Do not use if there is only a small batch to import if isinstance(parallel, str) and parallel.lower() == 'auto': if len(f) < 100: parallel = False if parallel: # Do not swap this as ``isinstance(True, int)`` returns ``True`` if isinstance(parallel, (bool, str)): n_cores = os.cpu_count() - 2 else: n_cores = int(parallel) with mp.Pool(processes=n_cores) as pool: results = pool.imap(_worker_wrapper, [dict(f=x, output=output, errors=errors, include_subdirs=include_subdirs, parallel=False) for x in f], chunksize=1) res = list(config.tqdm(results, desc='Importing', total=len(f), disable=config.pbar_hide, leave=config.pbar_leave)) else: # If not parallel just import the good 'ole way: sequentially res = [read_mesh(x, include_subdirs=include_subdirs, output=output, errors=errors, parallel=parallel, **kwargs) for x in config.tqdm(f, desc='Importing', disable=config.pbar_hide, leave=config.pbar_leave)] if output == 'neuron': return core.NeuronList([r for r in res if r]) return res try: # Open the file fname = '.'.join(os.path.basename(f).split('.')[:-1]) mesh = tm.load_mesh(f) if output == 'trimesh': return mesh attrs = {'name': fname, 'origin': f} attrs.update(kwargs) if output == 'volume': return core.Volume(mesh.vertices, mesh.faces, **attrs) else: return core.MeshNeuron(mesh, **attrs) except BaseException as e: msg = f'Error reading file {fname}.' if errors == 'raise': raise ImportError(msg) from e elif errors == 'log': logger.error(f'{msg}: {e}') return
def _worker_wrapper(kwargs): """Helper for importing meshes using multiple processes.""" return read_mesh(**kwargs)
[docs] def write_mesh(x: Union['core.NeuronList', 'core.MeshNeuron', 'core.Volume', 'tm.Trimesh'], filepath: Optional[str] = None, filetype: str = None, ) -> None: """Export meshes (MeshNeurons, Volumes, Trimeshes) to disk. Under the hood this is using trimesh to export meshes. Parameters ---------- x : MeshNeuron | Volume | Trimesh | NeuronList If multiple objects, will generate a file for each neuron (see also ``filepath``). filepath : None | str | list, optional If ``None``, will return byte string or list of thereof. If filepath will save to this file. If path will save neuron(s) in that path using ``{}`` as filename(s). If list, input must be NeuronList and a filepath must be provided for each neuron. filetype : stl | ply | obj, optional If ``filepath`` does not include the file extension, you need to provide it as ``filetype``. Returns ------- None If filepath is not ``None``. bytes If filepath is ``None``. See Also -------- :func:`navis.read_mesh` Import neurons. :func:`navis.write_precomputed` Write meshes to Neuroglancer's precomputed format. Examples -------- Write ``MeshNeurons`` to folder: >>> import navis >>> nl = navis.example_neurons(3, kind='mesh') >>> navis.write_mesh(nl, tmp_dir, filetype='obj') Specify the filenames: >>> import navis >>> nl = navis.example_neurons(3, kind='mesh') >>> navis.write_mesh(nl, tmp_dir / '{}.obj') Write directly to zip archive: >>> import navis >>> nl = navis.example_neurons(3, kind='mesh') >>> navis.write_mesh(nl, tmp_dir / '', filetype='obj') """ ALLOWED_FILETYPES = ('stl', 'ply', 'obj') if filetype is not None: utils.eval_param(filetype, name='filetype', allowed_values=ALLOWED_FILETYPES) else: # See if we can get filetype from filepath if filepath is not None: for f in ALLOWED_FILETYPES: if str(filepath).endswith(f'.{f}'): filetype = f break if not filetype: raise ValueError('Must provide mesh type either explicitly via ' '`filetype` variable or implicitly via the ' 'file extension in `filepath`') writer = base.Writer(_write_mesh, ext=f'.{filetype}') return writer.write_any(x, filepath=filepath)
def _write_mesh(x: Union['core.MeshNeuron', 'core.Volume', 'tm.Trimesh'], filepath: Optional[str] = None) -> None: """Write single mesh to disk.""" if filepath and os.path.isdir(filepath): if isinstance(x, core.MeshNeuron): if not raise ValueError('Neuron(s) must have an ID when destination ' 'is a folder') filepath = os.path.join(filepath, f'{}') elif isinstance(x, core.Volume): filepath = os.path.join(filepath, f'{}') else: raise ValueError(f'Unable to generate filename for {type(x)}') if isinstance(x, core.MeshNeuron): mesh = x.trimesh elif isinstance(x, tm.Trimesh): mesh = x else: raise TypeError(f'Unable to write data of type "{type(x)}"') # Write to disk (will only return content if filename is None) content = mesh.export(filepath) if not filepath: return content