<commit_before>#!/usr/bin/python
import re
import sys

CONFIGLINERE=re.compile("[ \t]*([a-z0-9_-]+)[ \t]*[=][ \t]*([0-9-]+) (.*)", re.IGNORECASE)

USAGE='''USAGE:
  configtool <FILE> set <key> <val>
  configtool <FILE> get <key>
  configtool <FILE> search <regexp>
  configtool <FILE> list
  configtool <FILE> print
'''

def getfilter(key):
  key=key.encode('ascii')
  def f(k,v):
    if k==key:
      print v
    return k,v
  return f

def setfilter(key, val):
  key=key.encode('ascii')
  def f(k,v):
    if k==key:
      v=val
    return k,v
  return f

def listfilter():
  def f(k,v):
    print k
    return k,v
  return f

def printfilter():
  def f(k,v):
    print k,'=',v
    return k,v
  return f

def searchfilter(term):
  printer=printfilter()
  pat=re.compile(term)
  def f(k,v):
    m=pat.search(k)
    if m:
      printer(k,v)
    return k,v
  return f

def processConfigFile(IN, OUT, FILTERS):
  #preload entire file
  lines=[x for x in open(IN)]
  o=open(OUT, "w")
  #make output
  for line in lines:
    try:
      key,val,com = CONFIGLINERE.match(line).group(1,2,3)
      for f in FILTERS:
        key,val=f(key,val)
      o.write("%s = %s %s\n" % (key, val, com))
    except:
      o.write(line);

def getConfigVal(file, key):
  val=[]
  def tmp(k,v):
    if k==key:
      val.append(v)
    return k,v
  processConfigFile(file, "/dev/null", [tmp])
  if len(val)>0:
    return val[0]
  else:
    return None

def setConfigVal(file, key, val):
  processConfigFile(file, file, [setfilter(key, val)])

def main(argv):
  IN="/dev/stdin"
  OUT="/dev/null"
  FILTERS=[]
  #parse args
  try:
    IN=argv[1]
    i=2
    while i<len(argv):
      act=argv[i].lower()
      if act=="set":
        OUT=IN
        FILTERS.append(setfilter(argv[i+1], argv[i+2]))
        i+=3
      elif act=="get":
        FILTERS.append(getfilter(argv[i+1]))
        i+=2
      elif act=="search":
        FILTERS.append(searchfilter(argv[i+1]))
        i+=2
      elif act=="list":
        FILTERS.append(listfilter())
        i+=1
      elif act=="print":
        FILTERS.append(printfilter())
        i+=1
      else:
        raise None
  except:
    sys.stderr.write(USAGE)
  processConfigFile(IN, OUT, FILTERS)

if __name__ == "__main__":
  main(sys.argv)

<commit_msg>rewrote configtool<commit_after>#!/usr/bin/python
import re
import sys

CONFIGLINERE=re.compile("[ \t]*([a-z0-9_-]+)[ \t]*[=][ \t]*([0-9-]+)(.*)", re.IGNORECASE)
USAGE='''USAGE:
  configtool <FILE> set <key> <val>
  configtool <FILE> get <key>
  configtool <FILE> search <regexp>
  configtool <FILE> list
  configtool <FILE> print
'''

class ConfigFile:
  def __init__(self, filename = None):
    self.values = dict()
    if filename:
      self.load(filename)

  def load(self, filename):
    fd = open(filename)
    for n, line in enumerate(fd):
      try:
        key,val,com = CONFIGLINERE.match(line).group(1,2,3)
        self.add(key, val, com.strip())
      except:
        sys.stderr.write("WARNING: %s:%d -- failed to parse config line\n" % (filename,n))
    fd.close()

  def save(self, filename):
    fd = open(filename, "w")
    for k in self.values.keys():
      val, com = self.values[k]
      fd.write("%s = %d %s\n" % (k, val, com))
    fd.close()

  def __getitem__(self, k):
    return self.values[k][0]

  def __setitem__(self, k, v):
    self.values[k][0] = int(v)

  def add(self, k, v, com="# added in script"):
    self.values[k] = [int(v), com]

  def keys(self):
    return self.values.keys()

def getConfigVal(filename, key):
  try:
    return ConfigFile(filename)[key]
  except:
    return None

def setConfigVal(filename, key, val):
  cfg = ConfigFile()
  try:
    cfg.load(filename)
  except:
    sys.stderr.write("WARNING: failed to load config file "+filename+"\n")
  try:
    cfg[key]=val
  except:
    sys.stderr.write("WARNING: missing val %s in %s\n" % (key, filename))
    cfg.add(key,val)
  cfg.save(filename)

def main(argv):
  #parse args
  try:
    IN=argv[1]
    OUT="/dev/null"
    cfg = ConfigFile(IN)
    i=2
    while i<len(argv):
      act=argv[i].lower()
      if act=="set":
        OUT=IN
        cfg[argv[i+1]] = argv[i+2]
        i+=3
      elif act=="get":
        print cfg[argv[i+1]]
        i+=2
      elif act=="search":
        r = re.compile(argv[i+1], re.IGNORECASE)
        for k in filter(lambda x: r.search(x), cfg.keys()):
          print k
        i+=2
      elif act=="list":
        for k in cfg.keys():
            print k
        i+=1
      elif act=="print":
        for k in cfg.keys():
            print k,'=',cfg[k]
        i+=1
      else:
        raise None
  except Exception, e:
    print e
    sys.stderr.write(USAGE)
  cfg.save(OUT)

if __name__ == "__main__":
  main(sys.argv)

<|endoftext|>


<commit_before>"""distutils.command

Package containing implementation of all the standard Distutils
commands."""

__revision__ = "$Id$"

__all__ = ['build',
           'build_py',
           'build_ext',
           'build_clib',
           'install',
           'install_lib',
           'install_scripts',
           'install_data',
           'clean',
           'sdist',
           'bdist',
           'bdist_dumb',
          ]
<commit_msg>Harry Henry Gebel: add 'bdist_rpm' command.<commit_after>"""distutils.command

Package containing implementation of all the standard Distutils
commands."""

__revision__ = "$Id$"

__all__ = ['build',
           'build_py',
           'build_ext',
           'build_clib',
           'install',
           'install_lib',
           'install_scripts',
           'install_data',
           'clean',
           'sdist',
           'bdist',
           'bdist_dumb',
           'bdist_rpm',
          ]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2019-2020 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import os
import pathlib
import typing
import functools
import shutil
import zipfile
import torch
import yaml
from pyannote.audio.features import Pretrained as _Pretrained
from pyannote.pipeline import Pipeline as _Pipeline

dependencies = ['pyannote.audio', 'torch']

_HUB_REPO = 'https://github.com/pyannote/pyannote-audio-hub'
_ZIP_URL = f'{_HUB_REPO}/raw/master/{{kind}}s/{{name}}.zip'
_PRETRAINED_URL = f'{_HUB_REPO}/raw/master/pretrained.yml'

# path where pre-trained models and pipelines are downloaded and cached
_HUB_DIR = pathlib.Path(os.environ.get("PYANNOTE_AUDIO_HUB",
                                       "~/.pyannote/hub")).expanduser().resolve()

# download pretrained.yml if needed
_PRETRAINED_YML = _HUB_DIR / 'pretrained.yml'

if not _PRETRAINED_YML.exists():
    msg = (
        f'Downloading list of pretrained models and pipelines '
        f'to "{_PRETRAINED_YML}".'
    )
    print(msg)
    torch.hub.download_url_to_file(_PRETRAINED_URL,
                                   _PRETRAINED_YML,
                                   progress=True)

def _generic(name: str,
             duration: float = None,
             step: float = 0.25,
             batch_size: int = 32,
             device: typing.Optional[typing.Union[typing.Text, torch.device]] = None,
             pipeline: typing.Optional[bool] = None,
             force_reload: bool = False) -> typing.Union[_Pretrained, _Pipeline]:
    """Load pretrained model or pipeline

    Parameters
    ----------
    name : str
        Name of pretrained model or pipeline
    duration : float, optional
        Override audio chunks duration.
        Defaults to the one used during training.
    step : float, optional
        Ratio of audio chunk duration used for the internal sliding window.
        Defaults to 0.25 (i.e. 75% overlap between two consecutive windows).
        Reducing this value might lead to better results (at the expense of
        slower processing).
    batch_size : int, optional
        Batch size used for inference. Defaults to 32.
    device : torch.device, optional
        Device used for inference.
    pipeline : bool, optional
        Wrap pretrained model in a (not fully optimized) pipeline.
    force_reload : bool
        Whether to discard the existing cache and force a fresh download.
        Defaults to use existing cache.

    Returns
    -------
    pretrained: `Pretrained` or `Pipeline`

    Usage
    -----
    >>> sad_pipeline = torch.hub.load('pyannote/pyannote-audio', 'sad_ami')
    >>> scores = model({'audio': '/path/to/audio.wav'})
    """

    model_exists = name in _MODELS
    pipeline_exists = name in _PIPELINES

    if model_exists and pipeline_exists:

        if pipeline is None:
            msg = (
                f'Both a pretrained model and a pretrained pipeline called '
                f'"{name}" are available. Use option "pipeline=True" to '
                f'load the pipeline, and "pipeline=False" to load the model.')
            raise ValueError(msg)

        if pipeline:
            kind = 'pipeline'
            zip_url = _ZIP_URL.format(kind=kind, name=name)
            sha256 = _PIPELINES[name]
            return_pipeline = True

        else:
            kind = 'model'
            zip_url = _ZIP_URL.format(kind=kind, name=name)
            sha256 = _MODELS[name]
            return_pipeline = False

    elif pipeline_exists:

        if pipeline is None:
            pipeline = True

        if not pipeline:
            msg = (
                f'Could not find any pretrained "{name}" model. '
                f'A pretrained "{name}" pipeline does exist. '
                f'Did you mean "pipeline=True"?'
            )
            raise ValueError(msg)

        kind = 'pipeline'
        zip_url = _ZIP_URL.format(kind=kind, name=name)
        sha256 = _PIPELINES[name]
        return_pipeline = True

    elif model_exists:

        if pipeline is None:
            pipeline = False

        kind = 'model'
        zip_url = _ZIP_URL.format(kind=kind, name=name)
        sha256 = _MODELS[name]
        return_pipeline = pipeline

        if name.startswith('emb_') and return_pipeline:
            msg = (
                f'Pretrained model "{name}" has no associated pipeline. Use '
                f'"pipeline=False" or remove "pipeline" option altogether.'
            )
            raise ValueError(msg)

    else:
        msg = (
            f'Could not find any pretrained model nor pipeline called "{name}".'
        )
        raise ValueError(msg)

    if sha256 is None:
        msg = (
            f'Pretrained {kind} "{name}" is not available yet but will be '
            f'released shortly. Stay tuned...'
        )
        raise NotImplementedError(msg)

    pretrained_dir = _HUB_DIR / f'{kind}s'
    pretrained_subdir = pretrained_dir / f'{name}'
    pretrained_zip = pretrained_dir / f'{name}.zip'

    if not pretrained_subdir.exists() or force_reload:

        if pretrained_subdir.exists():
            shutil.rmtree(pretrained_subdir)

        from pyannote.audio.utils.path import mkdir_p
        mkdir_p(pretrained_zip.parent)
        try:
            msg = (
                f'Downloading pretrained {kind} "{name}" to "{pretrained_zip}".'
            )
            print(msg)
            torch.hub.download_url_to_file(zip_url,
                                           pretrained_zip,
                                           hash_prefix=sha256,
                                           progress=True)
        except RuntimeError as e:
            shutil.rmtree(pretrained_subdir)
            msg = (
                f'Failed to download pretrained {kind} "{name}".'
                f'Please try again.')
            raise RuntimeError(msg)

        # unzip downloaded file
        with zipfile.ZipFile(pretrained_zip) as z:
            z.extractall(path=pretrained_dir)

    if kind == 'model':

        params_yml, = pretrained_subdir.glob('*/*/*/*/params.yml')
        pretrained =  _Pretrained(validate_dir=params_yml.parent,
                                  duration=duration,
                                  step=step,
                                  batch_size=batch_size,
                                  device=device)

        if return_pipeline:
            if name.startswith('sad_'):
                from pyannote.audio.pipeline.speech_activity_detection import SpeechActivityDetection
                pipeline = SpeechActivityDetection(scores=pretrained)
            elif name.startswith('scd_'):
                from pyannote.audio.pipeline.speaker_change_detection import SpeakerChangeDetection
                pipeline = SpeakerChangeDetection(scores=pretrained)
            elif name.startswith('ovl_'):
                from pyannote.audio.pipeline.overlap_detection import OverlapDetection
                pipeline = OverlapDetection(scores=pretrained)
            else:
                # this should never happen
                msg = (
                    f'Pretrained model "{name}" has no associated pipeline. Use '
                    f'"pipeline=False" or remove "pipeline" option altogether.'
                )
                raise ValueError(msg)

            return pipeline.load_params(params_yml)

        return pretrained

    elif kind == 'pipeline':

        from pyannote.audio.pipeline.utils import load_pretrained_pipeline
        params_yml, *_ = pretrained_subdir.glob('*/*/params.yml')
        return load_pretrained_pipeline(params_yml.parent)

with open(_PRETRAINED_YML, 'r') as fp:
    _pretrained = yaml.load(fp, Loader=yaml.SafeLoader)

_MODELS = _pretrained['models']
for name in _MODELS:
    locals()[name] = functools.partial(_generic, name)

_PIPELINES = _pretrained['pipelines']
for name in _PIPELINES:
    locals()[name] = functools.partial(_generic, name)

_SHORTCUTS = _pretrained['shortcuts']
for shortcut, name in _SHORTCUTS.items():
    locals()[shortcut] = locals()[name]

if __name__ == '__main__':
    DOCOPT = """Create torch.hub zip file from validation directory

Usage:
  hubconf.py <validate_dir>
  hubconf.py (-h | --help)
  hubconf.py --version

Options:
  -h --help     Show this screen.
  --version     Show version.
    """
    from docopt import docopt
    from pyannote.audio.applications.base import create_zip
    arguments = docopt(DOCOPT, version='hubconf')
    validate_dir = pathlib.Path(arguments['<validate_dir>'])
    hub_zip = create_zip(validate_dir)
    print(f'Created file "{hub_zip.name}" in directory "{validate_dir}".')
<commit_msg>BREAKING: remove "main" function from hubconf<commit_after>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2019-2020 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import os
import pathlib
import typing
import functools
import shutil
import zipfile
import torch
import yaml
from pyannote.audio.features import Pretrained as _Pretrained
from pyannote.pipeline import Pipeline as _Pipeline

dependencies = ['pyannote.audio', 'torch']

_HUB_REPO = 'https://github.com/pyannote/pyannote-audio-hub'
_ZIP_URL = f'{_HUB_REPO}/raw/master/{{kind}}s/{{name}}.zip'
_PRETRAINED_URL = f'{_HUB_REPO}/raw/master/pretrained.yml'

# path where pre-trained models and pipelines are downloaded and cached
_HUB_DIR = pathlib.Path(os.environ.get("PYANNOTE_AUDIO_HUB",
                                       "~/.pyannote/hub")).expanduser().resolve()

# download pretrained.yml if needed
_PRETRAINED_YML = _HUB_DIR / 'pretrained.yml'

if not _PRETRAINED_YML.exists():
    msg = (
        f'Downloading list of pretrained models and pipelines '
        f'to "{_PRETRAINED_YML}".'
    )
    print(msg)
    torch.hub.download_url_to_file(_PRETRAINED_URL,
                                   _PRETRAINED_YML,
                                   progress=True)

def _generic(name: str,
             duration: float = None,
             step: float = 0.25,
             batch_size: int = 32,
             device: typing.Optional[typing.Union[typing.Text, torch.device]] = None,
             pipeline: typing.Optional[bool] = None,
             force_reload: bool = False) -> typing.Union[_Pretrained, _Pipeline]:
    """Load pretrained model or pipeline

    Parameters
    ----------
    name : str
        Name of pretrained model or pipeline
    duration : float, optional
        Override audio chunks duration.
        Defaults to the one used during training.
    step : float, optional
        Ratio of audio chunk duration used for the internal sliding window.
        Defaults to 0.25 (i.e. 75% overlap between two consecutive windows).
        Reducing this value might lead to better results (at the expense of
        slower processing).
    batch_size : int, optional
        Batch size used for inference. Defaults to 32.
    device : torch.device, optional
        Device used for inference.
    pipeline : bool, optional
        Wrap pretrained model in a (not fully optimized) pipeline.
    force_reload : bool
        Whether to discard the existing cache and force a fresh download.
        Defaults to use existing cache.

    Returns
    -------
    pretrained: `Pretrained` or `Pipeline`

    Usage
    -----
    >>> sad_pipeline = torch.hub.load('pyannote/pyannote-audio', 'sad_ami')
    >>> scores = model({'audio': '/path/to/audio.wav'})
    """

    model_exists = name in _MODELS
    pipeline_exists = name in _PIPELINES

    if model_exists and pipeline_exists:

        if pipeline is None:
            msg = (
                f'Both a pretrained model and a pretrained pipeline called '
                f'"{name}" are available. Use option "pipeline=True" to '
                f'load the pipeline, and "pipeline=False" to load the model.')
            raise ValueError(msg)

        if pipeline:
            kind = 'pipeline'
            zip_url = _ZIP_URL.format(kind=kind, name=name)
            sha256 = _PIPELINES[name]
            return_pipeline = True

        else:
            kind = 'model'
            zip_url = _ZIP_URL.format(kind=kind, name=name)
            sha256 = _MODELS[name]
            return_pipeline = False

    elif pipeline_exists:

        if pipeline is None:
            pipeline = True

        if not pipeline:
            msg = (
                f'Could not find any pretrained "{name}" model. '
                f'A pretrained "{name}" pipeline does exist. '
                f'Did you mean "pipeline=True"?'
            )
            raise ValueError(msg)

        kind = 'pipeline'
        zip_url = _ZIP_URL.format(kind=kind, name=name)
        sha256 = _PIPELINES[name]
        return_pipeline = True

    elif model_exists:

        if pipeline is None:
            pipeline = False

        kind = 'model'
        zip_url = _ZIP_URL.format(kind=kind, name=name)
        sha256 = _MODELS[name]
        return_pipeline = pipeline

        if name.startswith('emb_') and return_pipeline:
            msg = (
                f'Pretrained model "{name}" has no associated pipeline. Use '
                f'"pipeline=False" or remove "pipeline" option altogether.'
            )
            raise ValueError(msg)

    else:
        msg = (
            f'Could not find any pretrained model nor pipeline called "{name}".'
        )
        raise ValueError(msg)

    if sha256 is None:
        msg = (
            f'Pretrained {kind} "{name}" is not available yet but will be '
            f'released shortly. Stay tuned...'
        )
        raise NotImplementedError(msg)

    pretrained_dir = _HUB_DIR / f'{kind}s'
    pretrained_subdir = pretrained_dir / f'{name}'
    pretrained_zip = pretrained_dir / f'{name}.zip'

    if not pretrained_subdir.exists() or force_reload:

        if pretrained_subdir.exists():
            shutil.rmtree(pretrained_subdir)

        from pyannote.audio.utils.path import mkdir_p
        mkdir_p(pretrained_zip.parent)
        try:
            msg = (
                f'Downloading pretrained {kind} "{name}" to "{pretrained_zip}".'
            )
            print(msg)
            torch.hub.download_url_to_file(zip_url,
                                           pretrained_zip,
                                           hash_prefix=sha256,
                                           progress=True)
        except RuntimeError as e:
            shutil.rmtree(pretrained_subdir)
            msg = (
                f'Failed to download pretrained {kind} "{name}".'
                f'Please try again.')
            raise RuntimeError(msg)

        # unzip downloaded file
        with zipfile.ZipFile(pretrained_zip) as z:
            z.extractall(path=pretrained_dir)

    if kind == 'model':

        params_yml, = pretrained_subdir.glob('*/*/*/*/params.yml')
        pretrained =  _Pretrained(validate_dir=params_yml.parent,
                                  duration=duration,
                                  step=step,
                                  batch_size=batch_size,
                                  device=device)

        if return_pipeline:
            if name.startswith('sad_'):
                from pyannote.audio.pipeline.speech_activity_detection import SpeechActivityDetection
                pipeline = SpeechActivityDetection(scores=pretrained)
            elif name.startswith('scd_'):
                from pyannote.audio.pipeline.speaker_change_detection import SpeakerChangeDetection
                pipeline = SpeakerChangeDetection(scores=pretrained)
            elif name.startswith('ovl_'):
                from pyannote.audio.pipeline.overlap_detection import OverlapDetection
                pipeline = OverlapDetection(scores=pretrained)
            else:
                # this should never happen
                msg = (
                    f'Pretrained model "{name}" has no associated pipeline. Use '
                    f'"pipeline=False" or remove "pipeline" option altogether.'
                )
                raise ValueError(msg)

            return pipeline.load_params(params_yml)

        return pretrained

    elif kind == 'pipeline':

        from pyannote.audio.pipeline.utils import load_pretrained_pipeline
        params_yml, *_ = pretrained_subdir.glob('*/*/params.yml')
        return load_pretrained_pipeline(params_yml.parent)

with open(_PRETRAINED_YML, 'r') as fp:
    _pretrained = yaml.load(fp, Loader=yaml.SafeLoader)

_MODELS = _pretrained['models']
for name in _MODELS:
    locals()[name] = functools.partial(_generic, name)

_PIPELINES = _pretrained['pipelines']
for name in _PIPELINES:
    locals()[name] = functools.partial(_generic, name)

_SHORTCUTS = _pretrained['shortcuts']
for shortcut, name in _SHORTCUTS.items():
    locals()[shortcut] = locals()[name]
<|endoftext|>


<commit_before># This file is part of the GBI project.
# Copyright (C) 2012 Omniscale GmbH & Co. KG <http://omniscale.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import glob
import re

from flask import Blueprint, render_template, g, url_for, redirect, request, flash, current_app, abort, jsonify

from flaskext.babel import _

from geobox.lib.server_logging import send_task_logging
from geobox.lib.vectorconvert import is_valid_shapefile, ConvertError
from geobox.lib.fs import open_file_explorer
from geobox.lib.couchdb import vector_layers_metadata
from geobox.model import VectorImportTask, User
from geobox.web import forms
from geobox.model.tasks import VectorExportTask

from geobox.lib.couchdb import CouchFileBox

from ..utils import request_is_local
from ..helper import redirect_back

import logging
log = logging.getLogger(__name__)

vector = Blueprint('vector', __name__)

@vector.before_request
def only_if_enabled():
    if not current_app.config.geobox_state.config.get('app', 'vector_import'):
        abort(404)

def prepare_geojson_form(form):
    geojson_files = get_geojson_list()
    form.file_name.choices = [(name, name) for name in geojson_files]

    couch_url = 'http://%s:%s' % ('127.0.0.1', current_app.config.geobox_state.config.get('couchdb', 'port'))
    form.layers.choices = [(item['dbname'], item['title']) for item in vector_layers_metadata(couch_url)]
    form.layers.choices.insert(0, ('', _('-- select layer or add new --')))
    return form

@vector.route('/vector/import_geojson', methods=['GET', 'POST'])
def import_geojson():
    # upload geojson
    for upload_file in request.files.getlist('file_upload'):
        import_dir = current_app.config.geobox_state.user_data_path('import')
        target = os.path.join(import_dir, upload_file.filename)
        if os.path.splitext(target)[1].lower() in ('.json'):
            try:
                if not os.path.isdir(import_dir):
                    os.mkdir(import_dir)
                f = open(target, 'wb')
                f.write(upload_file.stream.read())
                f.close()
            except IOError:
                flash(_('failed to upload %(name)s', name=upload_file.filename), 'error')

            flash(_('file %(name)s uploaded', name=upload_file.filename), 'info')
        else:
            flash(_('filetype of %(name)s not allowed', name=upload_file.filename), 'error')

    form = forms.ImportGeoJSONEdit(request.form)
    form = prepare_geojson_form(form)

    if not len(request.files):
        if form.validate_on_submit():
            if (form.layers.data and form.name.data) or (not form.layers.data and not form.name.data):
                flash(_('please select new layer or current layer to import'), 'error')
                return redirect(url_for('.import_geojson'))

            title = None
            if form.layers.data:
                layer = form.layers.data
            else:
                title = form.name.data
                layer = 'local_vector_' + re.sub(r'[^a-z0-9]*', '',  title.lower())
                if layer == 'local_vector_':
                    flash(_('None of the characters used for the layer is allowed'))
                    return redirect(url_for('.import_geojson'))

            task = VectorImportTask(
                db_name=layer,
                title=title,
                file_name=form.file_name.data,
                type_ = 'geojson'
            )
            send_task_logging(current_app.config.geobox_state, task)
            g.db.add(task)
            g.db.commit()
            return redirect(url_for('tasks.list'))

        elif request.method == 'POST':
            flash(_('form error'), 'error')

    file_browser = request_is_local()
    return render_template('vector/geojson.html', form=form, file_browser=file_browser)

@vector.route('/vector/import', methods=['GET', 'POST'])
def import_vector():
    from ...lib.vectormapping import Mapping

    uploaded_shapes=[]
    upload_failures=[]
    not_allowed_files=[]

    for upload_file in request.files.getlist('file_upload'):
        import_dir = current_app.config.geobox_state.user_data_path('import')
        target = os.path.join(import_dir, upload_file.filename)
        if os.path.splitext(target)[1].lower() in ('.shp', '.shx', '.dbf'):
            try:
                if not os.path.isdir(import_dir):
                    os.mkdir(import_dir)

                f = open(target, 'wb')
                f.write(upload_file.stream.read())
                f.close()
            except IOError:
                upload_failures.append(upload_file.filename)

            uploaded_shapes.append(upload_file.filename)
        else:
            not_allowed_files.append(upload_file.filename)

    form = forms.ImportVectorEdit(request.form)
    form.srs.choices = [(srs, srs) for srs in current_app.config.geobox_state.config.get('web', 'available_srs')]

    couch_url = 'http://%s:%s' % ('127.0.0.1', current_app.config.geobox_state.config.get('couchdb', 'port'))
    form.layers.choices = [(item['dbname'], item['title']) for item in vector_layers_metadata(couch_url)]
    form.layers.choices.insert(0, ('', _('-- select layer or add new --')))

    shape_files, missing_files = get_shapefile_list()
    form.file_name.choices = [(name, name) for name in shape_files]
    form.file_name.choices.insert(0, ('', _('-- select a shapefile --')))

    if not len(request.files):
        if form.validate_on_submit():
            try:
                is_valid_shapefile(current_app.config.geobox_state.user_data_path('import', form.file_name.data), Mapping(None, None, '*'))
            except ConvertError as e:
                flash(e, 'error')
                return render_template('vector/import.html', form=form)
            except OSError:
                flash(_('invalid shapefile'), 'error')
                return render_template('vector/import.html', form=form)

            if (form.layers.data and form.name.data) or (not form.layers.data and not form.name.data):
                flash(_('please select new layer or current layer to import'), 'error')
                return redirect(url_for('.import_vector'))

            title = None
            if form.layers.data:
                layer = form.layers.data
            else:
                title = form.name.data
                layer = 'local_vector_' + re.sub(r'[^a-z0-9]*', '',  title.lower())
                if layer == 'local_vector_':
                    flash(_('None of the characters used for the layer is allowed'))
                    return redirect(url_for('.import_vector'))

            task = VectorImportTask(
                db_name=layer,
                title=title,
                file_name=form.file_name.data,
                srs=form.srs.data,
                type_ = 'shp',
            )
            send_task_logging(current_app.config.geobox_state, task)
            g.db.add(task)
            g.db.commit()
            return redirect(url_for('tasks.list'))
        elif request.method == 'POST':
            flash(_('form error'), 'error')

    for missing_file in missing_files:
        flash(_('file %(name)s missing', name=missing_file), 'error')
    for upload_failure in upload_failures:
        flash(_('failed to upload %(name)s', name=upload_failure), 'error')
    for not_allowed_file in not_allowed_files:
        flash(_('filetype of %(name)s not allowed', name=not_allowed_file), 'error')
    for uploaded_shape in uploaded_shapes:
        flash(_('file %(name)s uploaded', name=uploaded_shape), 'info')

    file_browser = request_is_local()
    return render_template('vector/import.html', form=form, file_browser=file_browser)

def get_geojson_list():
    file_list = []

    import_dir = current_app.config.geobox_state.user_data_path('import')
    files = set((glob.glob(os.path.join(import_dir, '*.json')) + glob.glob(os.path.join(import_dir, '*.JSON'))))
    for geojson_file in files:
        s_path, s_name = os.path.split(geojson_file)
        name, ext = os.path.splitext(s_name)
        file_list.append((s_name))
    return file_list

def get_shapefile_list():
    shape_import_dir = current_app.config.geobox_state.user_data_path('import')
    shape_file_list = []
    missing_file_list = []

    shp_files = set((glob.glob(os.path.join(shape_import_dir, '*.shp')) + glob.glob(os.path.join(shape_import_dir, '*.SHP'))))
    for shape_file in shp_files:
        missing = False
        s_path, s_name = os.path.split(shape_file)
        name, ext = os.path.splitext(s_name)
        if not os.path.exists(os.path.join(s_path, name +'.shx')):
            missing = True
            missing_file_list.append(name + '.shx')
        if not os.path.exists(os.path.join(s_path, name + '.dbf')):
            missing = True
            missing_file_list.append(name + '.dbf')
        if not missing:
            shape_file_list.append((s_name))
    return (shape_file_list, missing_file_list)

@vector.route('/file_browser/import_dir')
def import_file_browser():
    open_file_explorer(current_app.config.geobox_state.user_data_path('import', make_dirs=True))
    return redirect_back('.import_vector')

@vector.route('/vector/export', methods=['POST'])
def export_vector():
    return create_export_task(
        proj = request.form.get('srs', False),
        layername = request.form.get('name', False),
        export_type = request.form.get('export_type', False),
        destination = request.form.get('destination', False),
        filename = request.form.get('filename', False),
        geojson = request.form.get('geojson', '')
    )

@vector.route('/vector/export/selected', methods=['POST'])
def export_selected_geometries():
    user = User(current_app.config.geobox_state.config.get('user', 'type'))
    target_box_name = 'file_box' if user.is_consultant else 'upload_box'
    target_box = current_app.config.geobox_state.config.get('couchdb', target_box_name)

    filename = request.form.get('filename', False)
    data = request.form.get('geojson', False)

    file_obj = {'content-type': 'application/json' , 'file': data, 'filename': filename + '.json' }

    couch = CouchFileBox('http://%s:%s' % ('127.0.0.1', current_app.config.geobox_state.config.get('couchdb', 'port')), target_box)

    try:
        couch.store_file(file_obj, overwrite=True)
    except Exception, e:
        log.exception(e)
        abort(500)
    return jsonify({'status': 'success'})

def create_export_task(proj, layername, export_type, destination, filename, geojson):
    task = VectorExportTask(
       db_name=layername,
       srs=proj,
       type_=export_type,
       destination=destination,
       file_name=filename,
       geojson=geojson,
    )

    g.db.add(task)
    send_task_logging(current_app.config.geobox_state, task)
    g.db.commit()
    return redirect(url_for('tasks.list'))
<commit_msg>Don't preselect first geojson<commit_after># This file is part of the GBI project.
# Copyright (C) 2012 Omniscale GmbH & Co. KG <http://omniscale.com>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import glob
import re

from flask import Blueprint, render_template, g, url_for, redirect, request, flash, current_app, abort, jsonify

from flaskext.babel import _

from geobox.lib.server_logging import send_task_logging
from geobox.lib.vectorconvert import is_valid_shapefile, ConvertError
from geobox.lib.fs import open_file_explorer
from geobox.lib.couchdb import vector_layers_metadata
from geobox.model import VectorImportTask, User
from geobox.web import forms
from geobox.model.tasks import VectorExportTask

from geobox.lib.couchdb import CouchFileBox

from ..utils import request_is_local
from ..helper import redirect_back

import logging
log = logging.getLogger(__name__)

vector = Blueprint('vector', __name__)

@vector.before_request
def only_if_enabled():
    if not current_app.config.geobox_state.config.get('app', 'vector_import'):
        abort(404)

def prepare_geojson_form(form):
    geojson_files = get_geojson_list()
    form.file_name.choices = [(name, name) for name in geojson_files]
    form.file_name.choices.insert(0, ('', _('-- select a geojson --')))

    couch_url = 'http://%s:%s' % ('127.0.0.1', current_app.config.geobox_state.config.get('couchdb', 'port'))
    form.layers.choices = [(item['dbname'], item['title']) for item in vector_layers_metadata(couch_url)]
    form.layers.choices.insert(0, ('', _('-- select layer or add new --')))
    return form

@vector.route('/vector/import_geojson', methods=['GET', 'POST'])
def import_geojson():
    # upload geojson
    for upload_file in request.files.getlist('file_upload'):
        import_dir = current_app.config.geobox_state.user_data_path('import')
        target = os.path.join(import_dir, upload_file.filename)
        if os.path.splitext(target)[1].lower() in ('.json'):
            try:
                if not os.path.isdir(import_dir):
                    os.mkdir(import_dir)
                f = open(target, 'wb')
                f.write(upload_file.stream.read())
                f.close()
            except IOError:
                flash(_('failed to upload %(name)s', name=upload_file.filename), 'error')

            flash(_('file %(name)s uploaded', name=upload_file.filename), 'info')
        else:
            flash(_('filetype of %(name)s not allowed', name=upload_file.filename), 'error')

    form = forms.ImportGeoJSONEdit(request.form)
    form = prepare_geojson_form(form)

    if not len(request.files):
        if form.validate_on_submit():
            if (form.layers.data and form.name.data) or (not form.layers.data and not form.name.data):
                flash(_('please select new layer or current layer to import'), 'error')
                return redirect(url_for('.import_geojson'))

            title = None
            if form.layers.data:
                layer = form.layers.data
            else:
                title = form.name.data
                layer = 'local_vector_' + re.sub(r'[^a-z0-9]*', '',  title.lower())
                if layer == 'local_vector_':
                    flash(_('None of the characters used for the layer is allowed'))
                    return redirect(url_for('.import_geojson'))

            task = VectorImportTask(
                db_name=layer,
                title=title,
                file_name=form.file_name.data,
                type_ = 'geojson'
            )
            send_task_logging(current_app.config.geobox_state, task)
            g.db.add(task)
            g.db.commit()
            return redirect(url_for('tasks.list'))

        elif request.method == 'POST':
            flash(_('form error'), 'error')

    file_browser = request_is_local()
    return render_template('vector/geojson.html', form=form, file_browser=file_browser)

@vector.route('/vector/import', methods=['GET', 'POST'])
def import_vector():
    from ...lib.vectormapping import Mapping

    uploaded_shapes=[]
    upload_failures=[]
    not_allowed_files=[]

    for upload_file in request.files.getlist('file_upload'):
        import_dir = current_app.config.geobox_state.user_data_path('import')
        target = os.path.join(import_dir, upload_file.filename)
        if os.path.splitext(target)[1].lower() in ('.shp', '.shx', '.dbf'):
            try:
                if not os.path.isdir(import_dir):
                    os.mkdir(import_dir)

                f = open(target, 'wb')
                f.write(upload_file.stream.read())
                f.close()
            except IOError:
                upload_failures.append(upload_file.filename)

            uploaded_shapes.append(upload_file.filename)
        else:
            not_allowed_files.append(upload_file.filename)

    form = forms.ImportVectorEdit(request.form)
    form.srs.choices = [(srs, srs) for srs in current_app.config.geobox_state.config.get('web', 'available_srs')]

    couch_url = 'http://%s:%s' % ('127.0.0.1', current_app.config.geobox_state.config.get('couchdb', 'port'))
    form.layers.choices = [(item['dbname'], item['title']) for item in vector_layers_metadata(couch_url)]
    form.layers.choices.insert(0, ('', _('-- select layer or add new --')))

    shape_files, missing_files = get_shapefile_list()
    form.file_name.choices = [(name, name) for name in shape_files]
    form.file_name.choices.insert(0, ('', _('-- select a shapefile --')))

    if not len(request.files):
        if form.validate_on_submit():
            try:
                is_valid_shapefile(current_app.config.geobox_state.user_data_path('import', form.file_name.data), Mapping(None, None, '*'))
            except ConvertError as e:
                flash(e, 'error')
                return render_template('vector/import.html', form=form)
            except OSError:
                flash(_('invalid shapefile'), 'error')
                return render_template('vector/import.html', form=form)

            if (form.layers.data and form.name.data) or (not form.layers.data and not form.name.data):
                flash(_('please select new layer or current layer to import'), 'error')
                return redirect(url_for('.import_vector'))

            title = None
            if form.layers.data:
                layer = form.layers.data
            else:
                title = form.name.data
                layer = 'local_vector_' + re.sub(r'[^a-z0-9]*', '',  title.lower())
                if layer == 'local_vector_':
                    flash(_('None of the characters used for the layer is allowed'))
                    return redirect(url_for('.import_vector'))

            task = VectorImportTask(
                db_name=layer,
                title=title,
                file_name=form.file_name.data,
                srs=form.srs.data,
                type_ = 'shp',
            )
            send_task_logging(current_app.config.geobox_state, task)
            g.db.add(task)
            g.db.commit()
            return redirect(url_for('tasks.list'))
        elif request.method == 'POST':
            flash(_('form error'), 'error')

    for missing_file in missing_files:
        flash(_('file %(name)s missing', name=missing_file), 'error')
    for upload_failure in upload_failures:
        flash(_('failed to upload %(name)s', name=upload_failure), 'error')
    for not_allowed_file in not_allowed_files:
        flash(_('filetype of %(name)s not allowed', name=not_allowed_file), 'error')
    for uploaded_shape in uploaded_shapes:
        flash(_('file %(name)s uploaded', name=uploaded_shape), 'info')

    file_browser = request_is_local()
    return render_template('vector/import.html', form=form, file_browser=file_browser)

def get_geojson_list():
    file_list = []

    import_dir = current_app.config.geobox_state.user_data_path('import')
    files = set((glob.glob(os.path.join(import_dir, '*.json')) + glob.glob(os.path.join(import_dir, '*.JSON'))))
    for geojson_file in files:
        s_path, s_name = os.path.split(geojson_file)
        name, ext = os.path.splitext(s_name)
        file_list.append((s_name))
    return file_list

def get_shapefile_list():
    shape_import_dir = current_app.config.geobox_state.user_data_path('import')
    shape_file_list = []
    missing_file_list = []

    shp_files = set((glob.glob(os.path.join(shape_import_dir, '*.shp')) + glob.glob(os.path.join(shape_import_dir, '*.SHP'))))
    for shape_file in shp_files:
        missing = False
        s_path, s_name = os.path.split(shape_file)
        name, ext = os.path.splitext(s_name)
        if not os.path.exists(os.path.join(s_path, name +'.shx')):
            missing = True
            missing_file_list.append(name + '.shx')
        if not os.path.exists(os.path.join(s_path, name + '.dbf')):
            missing = True
            missing_file_list.append(name + '.dbf')
        if not missing:
            shape_file_list.append((s_name))
    return (shape_file_list, missing_file_list)

@vector.route('/file_browser/import_dir')
def import_file_browser():
    open_file_explorer(current_app.config.geobox_state.user_data_path('import', make_dirs=True))
    return redirect_back('.import_vector')

@vector.route('/vector/export', methods=['POST'])
def export_vector():
    return create_export_task(
        proj = request.form.get('srs', False),
        layername = request.form.get('name', False),
        export_type = request.form.get('export_type', False),
        destination = request.form.get('destination', False),
        filename = request.form.get('filename', False),
        geojson = request.form.get('geojson', '')
    )

@vector.route('/vector/export/selected', methods=['POST'])
def export_selected_geometries():
    user = User(current_app.config.geobox_state.config.get('user', 'type'))
    target_box_name = 'file_box' if user.is_consultant else 'upload_box'
    target_box = current_app.config.geobox_state.config.get('couchdb', target_box_name)

    filename = request.form.get('filename', False)
    data = request.form.get('geojson', False)

    file_obj = {'content-type': 'application/json' , 'file': data, 'filename': filename + '.json' }

    couch = CouchFileBox('http://%s:%s' % ('127.0.0.1', current_app.config.geobox_state.config.get('couchdb', 'port')), target_box)

    try:
        couch.store_file(file_obj, overwrite=True)
    except Exception, e:
        log.exception(e)
        abort(500)
    return jsonify({'status': 'success'})

def create_export_task(proj, layername, export_type, destination, filename, geojson):
    task = VectorExportTask(
       db_name=layername,
       srs=proj,
       type_=export_type,
       destination=destination,
       file_name=filename,
       geojson=geojson,
    )

    g.db.add(task)
    send_task_logging(current_app.config.geobox_state, task)
    g.db.commit()
    return redirect(url_for('tasks.list'))
<|endoftext|>


<commit_before>#!/usr/bin/env python2

import argparse
import json
import sys
import time

from tinctools import *

class TincVis:

    def __init__(self, net, rundir):
        self.net = net
        self.nodes = {}
        self.edges = {}
        self.tincinfo = TincInfo.TincInfo(net,rundir=rundir)

        self.n2id = {}
        self.id2n = {}

        self.minWeight = sys.maxint
        self.maxWeight = 1

    def __computeHash(self, source, target):
        if (int(source['id']) <= int(target['id'])):
            return "{}-{}".format(source['id'],
                                  target['id'])
        else:
            return "{}-{}".format(target['id'],
                                  source['id'])

    def prepare(self):
        self.tincinfo.parse_all()
        uniqueEdges = set()

        try:
            del(self.tincinfo.nodes['(broadcast)'])
        except:
            pass

        cnt = 0
        for n in self.tincinfo.nodes:
            node = self.nodes.setdefault(n, {'id': None,
                                             'networks': self.tincinfo.nodes[n].network,
                                             'edges': 0,
                                             'version': 0,
                                             'name': n})
            node['edges'] = 0
            if not node['id']:
                node['id'] = cnt+1
                cnt += 1

        for ed in self.tincinfo.edges:

            try:
                _hash = self.__computeHash(self.nodes[ed['from']], self.nodes[ed['to']])
            except KeyError:
                print('warning: empty edge found - ignoring...')
                continue

            if _hash not in uniqueEdges:
                e = self.edges.setdefault(ed['from'], [])
                e.append({'source': self.nodes[ed['from']]['id'],
                          'target' : self.nodes[ed['to']]['id'],
                          '_hash': _hash,
                          'weight': ed['weight']})
                self.nodes[ed['from']]['edges'] += 1
                self.nodes[ed['to']]['edges'] += 1

                uniqueEdges.add(_hash)

            self.nodes[ed['to']]['version'] = ed['options']>>24

            if self.minWeight > int(ed['weight']):
                self.minWeight = int(ed['weight'])

            if self.maxWeight < int(ed['weight']):
                self.maxWeight = int(ed['weight'])

    def getFracWeight(self, cur):
        return round(1-((cur*100.0)/self.maxWeight)/100, 4)

    def writeJSON(self, outfile=None):
        if not outfile: return

        nodes = []
        links = []

        for n in self.nodes:
            self.n2id[n] = self.nodes[n]['id']
            self.id2n[self.nodes[n]['id']] = n

            nodes.append({"name": n,
                          "index": self.nodes[n]['id'],
                          "id": self.nodes[n]['id'],
                          "edges": self.nodes[n]['edges'],
                          "version": self.nodes[n]['version'],
                          "nets": self.nodes[n]['networks'],
                          "group": 0 if self.nodes[n]['edges'] == 0 else 1})

        for en in self.edges:
            for e in self.edges[en]:
                e['frac'] = self.getFracWeight(int(e['weight']))
                e['sname'] = en
                e['tname'] = self.id2n[e['target']]

                links.append(e)

        x= json.dumps({'nodes': nodes,
                       'links': links})

        with open(outfile, "w") as ofp:
            ofp.write(x)
            ofp.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--net", required=True, help="network name")
    parser.add_argument("-o", "--outfile", required=True, help="where to store json file")
    parser.add_argument("-f", "--foreground", help="stay in foreground and dump data periodically", action="store_true", default=False)
    parser.add_argument("-t", "--timeout", help="wait between dumps", default=30)
    parser.add_argument("-r", "--rundir", help="location of pid- and socket-files", default="/var/run")

    args = parser.parse_args()
    while True:
        tv = TincVis(net=args.net, rundir=args.rundir)
        tv.prepare()
        tv.writeJSON(outfile=args.outfile)
        if not args.foreground:
            sys.exit()
        else:
            time.sleep(args.timeout)
<commit_msg>Added version checking and excpetion handling.<commit_after>#!/usr/bin/env python2

import argparse
import json
import sys
import socket
import time

import tinctools

if tinctools.__version__ == 0.2:
    import tinctools.TincInfo as TincInfo
else:
    raise ImportWarning, "tinctools version: {} not supported".format(tinctools.__version__)

class TincVis:

    def __init__(self, net, rundir):
        self.net = net
        self.nodes = {}
        self.edges = {}
        self.tincinfo = TincInfo.TincInfo(net,rundir=rundir)

        self.n2id = {}
        self.id2n = {}

        self.minWeight = sys.maxint
        self.maxWeight = 1

    def __computeHash(self, source, target):
        if (int(source['id']) <= int(target['id'])):
            return "{}-{}".format(source['id'],
                                  target['id'])
        else:
            return "{}-{}".format(target['id'],
                                  source['id'])

    def prepare(self):
        self.tincinfo.parse_all()
        uniqueEdges = set()

        try:
            del(self.tincinfo.nodes['(broadcast)'])
        except:
            pass

        cnt = 0
        for n in self.tincinfo.nodes:
            node = self.nodes.setdefault(n, {'id': None,
                                             'networks': self.tincinfo.nodes[n].network,
                                             'edges': 0,
                                             'version': 0,
                                             'name': n})
            node['edges'] = 0
            if not node['id']:
                node['id'] = cnt+1
                cnt += 1

        for ed in self.tincinfo.edges:

            try:
                _hash = self.__computeHash(self.nodes[ed['from']], self.nodes[ed['to']])
            except KeyError:
                print('warning: empty edge found - ignoring...')
                continue

            if _hash not in uniqueEdges:
                e = self.edges.setdefault(ed['from'], [])
                e.append({'source': self.nodes[ed['from']]['id'],
                          'target' : self.nodes[ed['to']]['id'],
                          '_hash': _hash,
                          'weight': ed['weight']})
                self.nodes[ed['from']]['edges'] += 1
                self.nodes[ed['to']]['edges'] += 1

                uniqueEdges.add(_hash)

            self.nodes[ed['to']]['version'] = ed['options']>>24

            if self.minWeight > int(ed['weight']):
                self.minWeight = int(ed['weight'])

            if self.maxWeight < int(ed['weight']):
                self.maxWeight = int(ed['weight'])

    def getFracWeight(self, cur):
        return round(1-((cur*100.0)/self.maxWeight)/100, 4)

    def writeJSON(self, outfile=None):
        if not outfile: return

        nodes = []
        links = []

        for n in self.nodes:
            self.n2id[n] = self.nodes[n]['id']
            self.id2n[self.nodes[n]['id']] = n

            nodes.append({"name": n,
                          "index": self.nodes[n]['id'],
                          "id": self.nodes[n]['id'],
                          "edges": self.nodes[n]['edges'],
                          "version": self.nodes[n]['version'],
                          "nets": self.nodes[n]['networks'],
                          "group": 0 if self.nodes[n]['edges'] == 0 else 1})

        for en in self.edges:
            for e in self.edges[en]:
                e['frac'] = self.getFracWeight(int(e['weight']))
                e['sname'] = en
                e['tname'] = self.id2n[e['target']]

                links.append(e)

        x= json.dumps({'nodes': nodes,
                       'links': links})

        with open(outfile, "w") as ofp:
            ofp.write(x)
            ofp.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--net", required=True, help="network name")
    parser.add_argument("-o", "--outfile", required=True, help="where to store json file")
    parser.add_argument("-f", "--foreground", help="stay in foreground and dump data periodically", action="store_true", default=False)
    parser.add_argument("-t", "--timeout", help="wait between dumps", default=30)
    parser.add_argument("-r", "--rundir", help="location of pid- and socket-files", default="/var/run")

    args = parser.parse_args()
    while True:
        try:
            tv = TincVis(net=args.net, rundir=args.rundir)
            tv.prepare()
            tv.writeJSON(outfile=args.outfile)
        except socket.error as e:
            print(e)

        if not args.foreground:
            sys.exit()
        else:
            time.sleep(args.timeout)
<|endoftext|>


<commit_before>import plugintypes
import json
import subprocess
import os
import uuid
import shutil
import re

from urllib.parse import urlparse

from tempfile import TemporaryFile

"""
Installable plugins are currently hardcoded here for early
development of the installation process. They will be centralized
online at a later date.

Plugin data will be downloaded in JSON format and inserted to a sqlite
database for local cache.
"""

CENTRAL_REPO_URL="https://github.com/datamachine/telegram-pybot-plugin-repo"
CENTRAL_REPO_DIR="telegram-pybot-plugin-repo"

PLUGINS_REPOS_DIR="plugins.repos"
PLUGINS_TRASH_DIR="plugins.trash"
GIT_BIN="/usr/bin/git"

class PluginsPlugin(plugintypes.TelegramPlugin):
    """
    Plugin to manage other plugins. Enable, disable or reload.
    """
    patterns = [
        "^!plugins$",
        "^!plugins? (enable) ([\w_.-]+)$",
        "^!plugins? (disable) ([\w_.-]+)$",
        "^!plugins? (install) (.*)$",
        "^!plugins? (uninstall) ([\w_.-]+)$",
        "^!plugins? (search) (.*)$",
        "^!plugins? (update)$",
        # "^!plugins? (enable) ([\w_.-]+) (chat)",
        # "^!plugins? (disable) ([\w_.-]+) (chat)",
        "^!plugins? (reload)$"
    ]

    usage = [
        "!plugins: list all plugins.",
        "!plugins enable [plugin]: enable plugin.",
        "!plugins disable [plugin]: disable plugin.",
        # "!plugins disable [plugin] chat: disable plugin only this chat.",
        "!plugins reload: reloads all plugins."
    ]

    def __refresh_central_repo_object(self):
        with open(os.path.join(PLUGINS_REPOS_DIR, CENTRAL_REPO_DIR, "repo.json"), "r") as f:
            self.central_repo = json.load(f)

    def activate_plugin(self):
        if not os.path.exists(PLUGINS_REPOS_DIR):
            os.makedirs(PLUGINS_REPOS_DIR)
        if PLUGINS_REPOS_DIR not in self.plugin_manager.getPluginLocator().plugins_places:
            self.plugin_manager.updatePluginPlaces([PLUGINS_REPOS_DIR])
            self.reload_plugins()
        else:
            self.__refresh_central_repo_object()

    def run(self, msg, matches):
        if matches.group(0) == "!plugins":
            return self.list_plugins()

        command = matches.group(1)

        if command == "enable":
            return self.enable_plugin(matches)

        if command == "disable":
            return self.disable_plugin(matches)

        if command == "install":
            return self.install_plugin(matches)

        if command == "uninstall":
            return self.uninstall_plugin(matches)

        if command == "search":
            return self.search_plugins(matches.group(2))

        if command == "update":
            return self.update_central_repo()

        if command == "reload":
            return self.reload_plugins()

    def enable_plugin(self, matches):
        if self.plugin_manager.activatePluginByName(matches.group(2)):
            return "Enabled plugin: {}".format(matches.group(2))
        else:
            return "Error loading plugin: {}".format(matches.group(2))

    def disable_plugin(self, matches):
        if self.plugin_manager.deactivatePluginByName(matches.group(2)):
            return "Disabled plugin: {}".format(matches.group(2))
        else:
            return "Error disabling plugin: {}".format(matches.group(2))

    def __clone_repository(self, url):
            fp = TemporaryFile(mode="r")
            args = [GIT_BIN, "clone", url]
            p = subprocess.Popen(args, cwd=PLUGINS_REPOS_DIR, stdout=fp, stderr=fp)
            code = p.wait()
            fp.seek(0)
            return (code, fp.read())

    def install_plugin(self, matches):
        plugin = matches.group(2)
        urldata = urlparse(matches.group(2))

        url = None
        if urldata.scheme in [ "http", "https" ]:
            url = plugin
        elif plugin in self.central_repo.keys():
            url = self.central_repo[plugin]["url"]

        print(url)

        if not url:
            return "Invalid plugin or url: {}".format(plugin)

        code, msg = self.__clone_repository(url)
        if code != 0:
            return msg

        self.reload_plugins()

        return "{}\nSuccessfully installed plugin: {}".format(msg, plugin)

    def uninstall_plugin(self, matches):
        plugin_name = matches.group(2)
        for plugin in self.plugin_manager.getAllPlugins():
            if plugin.name != plugin_name:
                continue

            plugin_dir = os.path.relpath(os.path.dirname(plugin.path))
            if not plugin_dir.startswith(PLUGINS_REPOS_DIR):
                return "Error uninstalling plugin: {}.\nCannot locate plugin directory.".format(plugin_name)

            while plugin_dir and os.path.dirname(plugin_dir) != PLUGINS_REPOS_DIR:
                plugin_dir = os.path.dirname(plugin_dir)

            if not plugin_dir:
                return "Error uninstalling plugin: {}.\nCannot locate plugin directory.".format(plugin_name)

            self.plugin_manager.deactivatePluginByName(plugin_name)

            old_base = os.path.basename(plugin_dir)
            new_base = "{}.{}".format(os.path.basename(plugin_dir), str(uuid.uuid4()))
            new_dir = os.path.join(PLUGINS_TRASH_DIR, new_base)
            shutil.move(plugin_dir, new_dir)

            return "Uninstalled plugin: {}".format(plugin_name)
        return "Unable to find plugin: {}".format(plugin_name)

    def search_plugins(self, query):
        prog = re.compile(query, flags=re.IGNORECASE)
        results = ""
        for plugin in self.central_repo:
            description = self.central_repo[plugin]["description"]
            if prog.search(plugin) or prog.search(description):
                results += "{} | {} | {}\n".format(plugin, self.central_repo[plugin]["version"], self.central_repo[plugin]["description"])
        return results

    def update_central_repo(self):
        central_repo_path = os.path.join(PLUGINS_REPOS_DIR, CENTRAL_REPO_DIR)
        f = TemporaryFile()
        if not os.path.exists(central_repo_path):
            args = [GIT_BIN, "clone", CENTRAL_REPO_URL]
            p = subprocess.Popen(args, cwd=PLUGINS_REPOS_DIR, stdout=f, stderr=f)
            p.wait()
        else:
            args = [GIT_BIN, "reset", "--hard"]
            p = subprocess.Popen(args, cwd=central_repo_path, stdout=f, stderr=f)
            p.wait()
            args = [GIT_BIN, "pull"]
            p = subprocess.Popen(args, cwd=central_repo_path, stdout=f, stderr=f)
            p.wait()
        self.__refresh_central_repo_object()
        f.seek(0)
        return f.read()

    def reload_plugins(self):
        self.plugin_manager.collectPlugins()
        return "Plugins reloaded"

    def list_plugins(self):
        text = ""
        for plugin in self.plugin_manager.getAllPlugins():
            text += "{0}: ({1})\n".format(plugin.name,
                                          "(Enabled)" if plugin.plugin_object.is_activated else "(Disabled)")

        return text
<commit_msg>git_bin is now a user configurable option<commit_after>import plugintypes
import json
import subprocess
import os
import uuid
import shutil
import re

from urllib.parse import urlparse

from tempfile import TemporaryFile

"""
Installable plugins are currently hardcoded here for early
development of the installation process. They will be centralized
online at a later date.

Plugin data will be downloaded in JSON format and inserted to a sqlite
database for local cache.
"""

CENTRAL_REPO_URL="https://github.com/datamachine/telegram-pybot-plugin-repo"
CENTRAL_REPO_DIR="telegram-pybot-plugin-repo"

PLUGINS_REPOS_DIR="plugins.repos"
PLUGINS_TRASH_DIR="plugins.trash"

class PluginsPlugin(plugintypes.TelegramPlugin):
    """
    Plugin to manage other plugins. Enable, disable or reload.
    """
    patterns = [
        "^!plugins$",
        "^!plugins? (enable) ([\w_.-]+)$",
        "^!plugins? (disable) ([\w_.-]+)$",
        "^!plugins? (install) (.*)$",
        "^!plugins? (uninstall) ([\w_.-]+)$",
        "^!plugins? (search) (.*)$",
        "^!plugins? (update)$",
        # "^!plugins? (enable) ([\w_.-]+) (chat)",
        # "^!plugins? (disable) ([\w_.-]+) (chat)",
        "^!plugins? (reload)$"
    ]

    usage = [
        "!plugins: list all plugins.",
        "!plugins enable [plugin]: enable plugin.",
        "!plugins disable [plugin]: disable plugin.",
        # "!plugins disable [plugin] chat: disable plugin only this chat.",
        "!plugins reload: reloads all plugins."
    ]

    @property
    def git_bin(self):
        if not self.has_option("git_bin"):
            self.write_option("git_bin", "/usr/bin/git")
        return self.read_option("git_bin")

    def __refresh_central_repo_object(self):
        with open(os.path.join(PLUGINS_REPOS_DIR, CENTRAL_REPO_DIR, "repo.json"), "r") as f:
            self.central_repo = json.load(f)

    def activate_plugin(self):
        if not os.path.exists(PLUGINS_REPOS_DIR):
            os.makedirs(PLUGINS_REPOS_DIR)
        if PLUGINS_REPOS_DIR not in self.plugin_manager.getPluginLocator().plugins_places:
            self.plugin_manager.updatePluginPlaces([PLUGINS_REPOS_DIR])
            self.reload_plugins()
        else:
            self.__refresh_central_repo_object()

    def run(self, msg, matches):
        if matches.group(0) == "!plugins":
            return self.list_plugins()

        command = matches.group(1)

        if command == "enable":
            return self.enable_plugin(matches)

        if command == "disable":
            return self.disable_plugin(matches)

        if command == "install":
            return self.install_plugin(matches)

        if command == "uninstall":
            return self.uninstall_plugin(matches)

        if command == "search":
            return self.search_plugins(matches.group(2))

        if command == "update":
            return self.update_central_repo()

        if command == "reload":
            return self.reload_plugins()

    def enable_plugin(self, matches):
        if self.plugin_manager.activatePluginByName(matches.group(2)):
            return "Enabled plugin: {}".format(matches.group(2))
        else:
            return "Error loading plugin: {}".format(matches.group(2))

    def disable_plugin(self, matches):
        if self.plugin_manager.deactivatePluginByName(matches.group(2)):
            return "Disabled plugin: {}".format(matches.group(2))
        else:
            return "Error disabling plugin: {}".format(matches.group(2))

    def __clone_repository(self, url):
            fp = TemporaryFile(mode="r")
            args = [self.git_bin, "clone", url]
            p = subprocess.Popen(args, cwd=PLUGINS_REPOS_DIR, stdout=fp, stderr=fp)
            code = p.wait()
            fp.seek(0)
            return (code, fp.read())

    def install_plugin(self, matches):
        plugin = matches.group(2)
        urldata = urlparse(matches.group(2))

        url = None
        if urldata.scheme in [ "http", "https" ]:
            url = plugin
        elif plugin in self.central_repo.keys():
            url = self.central_repo[plugin]["url"]

        print(url)

        if not url:
            return "Invalid plugin or url: {}".format(plugin)

        code, msg = self.__clone_repository(url)
        if code != 0:
            return msg

        self.reload_plugins()

        return "{}\nSuccessfully installed plugin: {}".format(msg, plugin)

    def uninstall_plugin(self, matches):
        plugin_name = matches.group(2)
        for plugin in self.plugin_manager.getAllPlugins():
            if plugin.name != plugin_name:
                continue

            plugin_dir = os.path.relpath(os.path.dirname(plugin.path))
            if not plugin_dir.startswith(PLUGINS_REPOS_DIR):
                return "Error uninstalling plugin: {}.\nCannot locate plugin directory.".format(plugin_name)

            while plugin_dir and os.path.dirname(plugin_dir) != PLUGINS_REPOS_DIR:
                plugin_dir = os.path.dirname(plugin_dir)

            if not plugin_dir:
                return "Error uninstalling plugin: {}.\nCannot locate plugin directory.".format(plugin_name)

            self.plugin_manager.deactivatePluginByName(plugin_name)

            old_base = os.path.basename(plugin_dir)
            new_base = "{}.{}".format(os.path.basename(plugin_dir), str(uuid.uuid4()))
            new_dir = os.path.join(PLUGINS_TRASH_DIR, new_base)
            shutil.move(plugin_dir, new_dir)

            return "Uninstalled plugin: {}".format(plugin_name)
        return "Unable to find plugin: {}".format(plugin_name)

    def search_plugins(self, query):
        prog = re.compile(query, flags=re.IGNORECASE)
        results = ""
        for plugin in self.central_repo:
            description = self.central_repo[plugin]["description"]
            if prog.search(plugin) or prog.search(description):
                results += "{} | {} | {}\n".format(plugin, self.central_repo[plugin]["version"], self.central_repo[plugin]["description"])
        return results

    def update_central_repo(self):
        central_repo_path = os.path.join(PLUGINS_REPOS_DIR, CENTRAL_REPO_DIR)
        f = TemporaryFile()
        if not os.path.exists(central_repo_path):
            args = [self.git_bin, "clone", CENTRAL_REPO_URL]
            p = subprocess.Popen(args, cwd=PLUGINS_REPOS_DIR, stdout=f, stderr=f)
            p.wait()
        else:
            args = [self.git_bin, "reset", "--hard"]
            p = subprocess.Popen(args, cwd=central_repo_path, stdout=f, stderr=f)
            p.wait()
            args = [self.git_bin, "pull"]
            p = subprocess.Popen(args, cwd=central_repo_path, stdout=f, stderr=f)
            p.wait()
        self.__refresh_central_repo_object()
        f.seek(0)
        return f.read()

    def reload_plugins(self):
        self.plugin_manager.collectPlugins()
        return "Plugins reloaded"

    def list_plugins(self):
        text = ""
        for plugin in self.plugin_manager.getAllPlugins():
            text += "{0}: ({1})\n".format(plugin.name,
                                          "(Enabled)" if plugin.plugin_object.is_activated else "(Disabled)")

        return text
<|endoftext|>


<commit_before>import numpy as np
import scipy as sp
import scipy.special as sps
from numpy.testing import (run_module_suite,
                           assert_,
                           assert_equal,
                           assert_raises,
                           assert_array_almost_equal)
from dipy.denoise.localpca import localpca
from dipy.sims.voxel import multi_tensor
from dipy.core.gradients import gradient_table, generate_bvecs
from dipy.core.sphere import disperse_charges, HemiSphere
from dipy.sims.voxel import multi_tensor

def rfiw_phantom(gtab, snr=None):
    """rectangle fiber immersed in water"""
    # define voxel index
    slice_ind = np.zeros((10, 10, 8))
    slice_ind[4:7, 4:7, :] = 1
    slice_ind[4:7, 7, :] = 2
    slice_ind[7, 7, :] = 3
    slice_ind[7, 4:7, :] = 4
    slice_ind[7, 3, :] = 5
    slice_ind[4:7, 3, :] = 6
    slice_ind[3, 3, :] = 7
    slice_ind[3, 4:7, :] = 8
    slice_ind[3, 7, :] = 9

    # Define tissue diffusion parameters
    # Restricted diffusion
    ADr = 0.99e-3
    RDr = 0.0
    # Hindered diffusion
    ADh = 2.26e-3
    RDh = 0.87
    # S0 value for tissue
    S1 = 50
    # Fraction between Restricted and Hindered diffusion
    fia = 0.51

    # Define water diffusion
    Dwater = 3e-3
    S2 = 100  # S0 value for water

    # Define tissue volume fraction for each voxel type (in index order)
    f = np.array([0., 1., 0.6, 0.18, 0.30, 0.15, 0.50, 0.35, 0.70, 0.42])

    # Define S0 for each voxel (in index order)
    S0 = S1 * f + S2 * (1 - f)

    # multi tensor simulations assume that each water pull as constant S0
    # since I am assuming that tissue and water voxels have different S0,
    # tissue volume fractions have to be adjusted to the measured f values when
    # constant S0 are assumed constant. Doing this correction, simulations will
    # be analogous to simulates that S0 are different for each media. (For more
    # details on this contact the phantom designer)
    f1 = f * S1 / S0

    mevals = np.array([[ADr, RDr, RDr], [ADh, RDh, RDh],
                       [Dwater, Dwater, Dwater]])
    angles = [(0, 0, 1), (0, 0, 1), (0, 0, 1)]
    DWI = np.zeros(slice_ind.shape + (gtab.bvals.size, ))
    for i in range(10):
        fractions = [f1[i] * fia * 100, f1[i] *
                     (1 - fia) * 100, (1 - f1[i]) * 100]
        sig, direction = multi_tensor(gtab, mevals, S0=S0[i], angles=angles,
                                      fractions=fractions, snr=None)
        DWI[slice_ind == i, :] = sig

    if snr is None:
        return DWI
    else:
        sigma = S2 * 1.0 / snr
        n1 = np.random.normal(0, sigma, size=DWI.shape)
        n2 = np.random.normal(0, sigma, size=DWI.shape)
        return [np.sqrt((DWI / np.sqrt(2) + n1)**2 +
                        (DWI / np.sqrt(2) + n2)**2), sigma]

def gen_gtab():
    # generate a gradient table for phantom data
    directions8 = generate_bvecs(8)
    directions30 = generate_bvecs(30)
    directions60 = generate_bvecs(60)
    # Create full dataset parameters
    # (6 b-values = 0, 8 directions for b-value 300, 30 directions for b-value
    # 1000 and 60 directions for b-value 2000)
    bvals = np.hstack((np.zeros(6),
                       300 * np.ones(8),
                       1000 * np.ones(30),
                       2000 * np.ones(60)))
    bvecs = np.vstack((np.zeros((6, 3)),
                       directions8, directions30, directions60))
    gtab = gradient_table(bvals, bvecs)
    return gtab

def test_lpca_static():
    S0 = 100 * np.ones((20, 20, 20, 20), dtype='f8')
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20), dtype=np.float64))
    assert_array_almost_equal(S0, S0ns)

def test_lpca_random_noise():
    S0 = 100 + 2 * np.random.standard_normal((22, 23, 30, 20))
    S0ns = localpca(S0, sigma=np.std(S0))

    assert_(S0ns.min() > S0.min())
    assert_(S0ns.max() < S0.max())
    assert_equal(np.round(S0ns.mean()), 100)

def test_lpca_boundary_behaviour():
    # check is first slice is getting denoised or not ?
    S0 = 100 * np.ones((20, 20, 20, 20), dtype='f8')
    S0[:, :, 0, :] = S0[:, :, 0, :] + 2 * \
        np.random.standard_normal((20, 20, 20))
    S0_first = S0[:, :, 0, :]
    S0ns = localpca(S0, sigma=np.std(S0))
    S0ns_first = S0ns[:, :, 0, :]
    rmses = np.sum(np.abs(S0ns_first - S0_first)) / \
        (100.0 * 20.0 * 20.0 * 20.0)

    # shows that S0n_first is not very close to S0_first
    assert_(rmses > 0.0001)
    assert_equal(np.round(S0ns_first.mean()), 100)

    rmses = np.sum(np.abs(S0ns_first - S0_first)) / \
        (100.0 * 20.0 * 20.0 * 20.0)

    # shows that S0n_first is not very close to S0_first
    assert_(rmses > 0.0001)
    assert_equal(np.round(S0ns_first.mean()), 100)

def test_lpca_rmse():
    S0_w_noise = 100 + 2 * np.random.standard_normal((22, 23, 30, 20))
    rmse_w_noise = np.sqrt(np.mean((S0_w_noise - 100) ** 2))
    S0_denoised = localpca(S0_w_noise, sigma=np.std(S0_w_noise))
    rmse_denoised = np.sqrt(np.mean((S0_denoised - 100) ** 2))
    # Denoising should always improve the RMSE:
    assert_(rmse_denoised < rmse_w_noise)

def test_lpca_sharpness():
    S0 = np.ones((30, 30, 30, 20), dtype=np.float64) * 100
    S0[10:20, 10:20, 10:20, :] = 50
    S0[20:30, 20:30, 20:30, :] = 0
    S0 = S0 + 20 * np.random.standard_normal((30, 30, 30, 20))
    S0ns = localpca(S0, sigma=20.0)
    # check the edge gradient
    edgs = np.abs(np.mean(S0ns[8, 10:20, 10:20] - S0ns[12, 10:20, 10:20]) - 50)
    assert_(edgs < 2)

def test_lpca_dtype():
    # If out_dtype is not specified, we retain the original precision:
    S0 = 200 * np.ones((20, 20, 20, 3), dtype=np.float64)
    S0ns = localpca(S0, sigma=1)
    assert_equal(S0.dtype, S0ns.dtype)

    S0 = 200 * np.ones((20, 20, 20, 20), dtype=np.uint16)
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20)))
    assert_equal(S0.dtype, S0ns.dtype)

    # If we set out_dtype, we get what we asked for:
    S0 = 200 * np.ones((20, 20, 20, 20), dtype=np.uint16)
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20)),
                    out_dtype=np.float32)
    assert_equal(np.float32, S0ns.dtype)

    # If we set a few entries to zero, this induces negative entries in the
    # Resulting denoised array:
    S0[5:8, 5:8, 5:8] = 0
    # But if we should always get all non-negative results:
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20)), out_dtype=np.uint16)
    assert_(np.all(S0ns >= 0))
    # And no wrap-around to crazy high values:
    assert_(np.all(S0ns <= 200))

def test_lpca_wrong():
    S0 = np.ones((20, 20))
    assert_raises(ValueError, localpca, S0, sigma=1)

def test_phantom():
    gtab = gen_gtab()
    DWI_clean = rfiw_phantom(gtab, snr=None)
    DWI, sigma = rfiw_phantom(gtab, snr=30)
    # To test without Rician correction
    temp = (DWI_clean / sigma)**2
    DWI_clean_wrc = (sigma * np.sqrt(np.pi / 2) * np.exp(-0.5 * temp) *
                     ((1 + 0.5 * temp) * sps.iv(0, 0.25 * temp) + 0.5 * temp *
                     sps.iv(1, 0.25 * temp))**2)

    DWI_den = localpca(DWI, sigma, patch_radius=3)
    rmse_den = np.sum(np.abs(DWI_clean - DWI_den)) / np.sum(np.abs(DWI_clean))
    rmse_noisy = np.sum(np.abs(DWI_clean - DWI)) / np.sum(np.abs(DWI_clean))

    rmse_den_wrc = np.sum(np.abs(DWI_clean_wrc - DWI_den)
                          ) / np.sum(np.abs(DWI_clean_wrc))
    rmse_noisy_wrc = np.sum(np.abs(DWI_clean_wrc - DWI)) / \
        np.sum(np.abs(DWI_clean_wrc))

    assert_(np.max(DWI_clean) / sigma < np.max(DWI_den) / sigma)
    assert_(np.max(DWI_den) / sigma < np.max(DWI) / sigma)
    assert_(rmse_den < rmse_noisy)
    assert_(rmse_den_wrc < rmse_noisy_wrc)

    # Check if the results of different PCA methods (eig, svd) are similar
    DWI_den_svd = localpca(DWI, sigma, pca_method='svd', patch_radius=3)
    assert_array_almost_equal(DWI_den, DWI_den_svd)

    assert_raises(ValueError, localpca, DWI, sigma, pca_method='empty')

    # Try this with a sigma volume, instead of a scalar
    sigma_vol = sigma * np.ones(DWI.shape[:-1])
    mask = np.zeros_like(DWI, dtype=bool)[..., 0]
    mask[2:-2, 2:-2, 2:-2] = True
    DWI_den = localpca(DWI, sigma_vol, mask, patch_radius=3)
    DWI_clean_masked = DWI_clean.copy()
    DWI_clean_masked[~mask] = 0
    DWI_masked = DWI.copy()
    DWI_masked[~mask] = 0
    rmse_den = np.sum(np.abs(DWI_clean_masked - DWI_den)) / np.sum(np.abs(
            DWI_clean_masked))
    rmse_noisy = np.sum(np.abs(DWI_clean_masked - DWI_masked)) / np.sum(np.abs(
            DWI_clean_masked))

    DWI_clean_wrc_masked = DWI_clean_wrc.copy()
    DWI_clean_wrc_masked[~mask] = 0
    rmse_den_wrc = np.sum(np.abs(DWI_clean_wrc_masked - DWI_den)
                          ) / np.sum(np.abs(DWI_clean_wrc_masked))
    rmse_noisy_wrc = np.sum(np.abs(DWI_clean_wrc_masked - DWI_masked)) / \
        np.sum(np.abs(DWI_clean_wrc_masked))

    assert_(np.max(DWI_clean) / sigma < np.max(DWI_den) / sigma)
    assert_(np.max(DWI_den) / sigma < np.max(DWI) / sigma)
    assert_(rmse_den < rmse_noisy)
    assert_(rmse_den_wrc < rmse_noisy_wrc)

def test_lpca_ill_conditioned():
    gtab = gen_gtab()
    DWI, sigma = rfiw_phantom(gtab, snr=30)
    assert_raises(ValueError, localpca, DWI, sigma, patch_radius=1)

def test_lpca_sigma_wrong_shape():
    gtab = gen_gtab()
    DWI, sigma = rfiw_phantom(gtab, snr=30)
    # If sigma is 3D but shape is not like DWI.shape[:-1], an error is raised:
    sigma = np.zeros((DWI.shape[0], DWI.shape[1] + 1, DWI.shape[2]))
    assert_raises(ValueError, localpca, DWI, sigma)

if __name__ == '__main__':
    run_module_suite()
<commit_msg>TESTS: test for pca eigenvalue classifier<commit_after>import numpy as np
import scipy.special as sps
from numpy.testing import (run_module_suite,
                           assert_,
                           assert_equal,
                           assert_raises,
                           assert_array_almost_equal)
from dipy.denoise.localpca import (localpca, pca_classifier)
from dipy.sims.voxel import multi_tensor
from dipy.core.gradients import gradient_table, generate_bvecs

def rfiw_phantom(gtab, snr=None):
    """rectangle fiber immersed in water"""
    # define voxel index
    slice_ind = np.zeros((10, 10, 8))
    slice_ind[4:7, 4:7, :] = 1
    slice_ind[4:7, 7, :] = 2
    slice_ind[7, 7, :] = 3
    slice_ind[7, 4:7, :] = 4
    slice_ind[7, 3, :] = 5
    slice_ind[4:7, 3, :] = 6
    slice_ind[3, 3, :] = 7
    slice_ind[3, 4:7, :] = 8
    slice_ind[3, 7, :] = 9

    # Define tissue diffusion parameters
    # Restricted diffusion
    ADr = 0.99e-3
    RDr = 0.0
    # Hindered diffusion
    ADh = 2.26e-3
    RDh = 0.87e-3
    # S0 value for tissue
    S1 = 50
    # Fraction between Restricted and Hindered diffusion
    fia = 0.51

    # Define water diffusion
    Dwater = 3e-3
    S2 = 100  # S0 value for water

    # Define tissue volume fraction for each voxel type (in index order)
    f = np.array([0., 1., 0.6, 0.18, 0.30, 0.15, 0.50, 0.35, 0.70, 0.42])

    # Define S0 for each voxel (in index order)
    S0 = S1 * f + S2 * (1 - f)

    # multi tensor simulations assume that each water pull as constant S0
    # since I am assuming that tissue and water voxels have different S0,
    # tissue volume fractions have to be adjusted to the measured f values when
    # constant S0 are assumed constant. Doing this correction, simulations will
    # be analogous to simulates that S0 are different for each media. (For more
    # details on this contact the phantom designer)
    f1 = f * S1 / S0

    mevals = np.array([[ADr, RDr, RDr], [ADh, RDh, RDh],
                       [Dwater, Dwater, Dwater]])
    angles = [(0, 0, 1), (0, 0, 1), (0, 0, 1)]
    DWI = np.zeros(slice_ind.shape + (gtab.bvals.size, ))
    for i in range(10):
        fractions = [f1[i] * fia * 100, f1[i] *
                     (1 - fia) * 100, (1 - f1[i]) * 100]
        sig, direction = multi_tensor(gtab, mevals, S0=S0[i], angles=angles,
                                      fractions=fractions, snr=None)
        DWI[slice_ind == i, :] = sig

    if snr is None:
        return DWI
    else:
        sigma = S2 * 1.0 / snr
        n1 = np.random.normal(0, sigma, size=DWI.shape)
        n2 = np.random.normal(0, sigma, size=DWI.shape)
        return [np.sqrt((DWI / np.sqrt(2) + n1)**2 +
                        (DWI / np.sqrt(2) + n2)**2), sigma]

def gen_gtab():
    # generate a gradient table for phantom data
    directions8 = generate_bvecs(8)
    directions30 = generate_bvecs(30)
    directions60 = generate_bvecs(60)
    # Create full dataset parameters
    # (6 b-values = 0, 8 directions for b-value 300, 30 directions for b-value
    # 1000 and 60 directions for b-value 2000)
    bvals = np.hstack((np.zeros(6),
                       300 * np.ones(8),
                       1000 * np.ones(30),
                       2000 * np.ones(60)))
    bvecs = np.vstack((np.zeros((6, 3)),
                       directions8, directions30, directions60))
    gtab = gradient_table(bvals, bvecs)
    return gtab

def test_lpca_static():
    S0 = 100 * np.ones((20, 20, 20, 20), dtype='f8')
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20), dtype=np.float64))
    assert_array_almost_equal(S0, S0ns)

def test_lpca_random_noise():
    S0 = 100 + 2 * np.random.standard_normal((22, 23, 30, 20))
    S0ns = localpca(S0, sigma=np.std(S0))

    assert_(S0ns.min() > S0.min())
    assert_(S0ns.max() < S0.max())
    assert_equal(np.round(S0ns.mean()), 100)

def test_lpca_boundary_behaviour():
    # check is first slice is getting denoised or not ?
    S0 = 100 * np.ones((20, 20, 20, 20), dtype='f8')
    S0[:, :, 0, :] = S0[:, :, 0, :] + 2 * \
        np.random.standard_normal((20, 20, 20))
    S0_first = S0[:, :, 0, :]
    S0ns = localpca(S0, sigma=np.std(S0))
    S0ns_first = S0ns[:, :, 0, :]
    rmses = np.sum(np.abs(S0ns_first - S0_first)) / \
        (100.0 * 20.0 * 20.0 * 20.0)

    # shows that S0n_first is not very close to S0_first
    assert_(rmses > 0.0001)
    assert_equal(np.round(S0ns_first.mean()), 100)

    rmses = np.sum(np.abs(S0ns_first - S0_first)) / \
        (100.0 * 20.0 * 20.0 * 20.0)

    # shows that S0n_first is not very close to S0_first
    assert_(rmses > 0.0001)
    assert_equal(np.round(S0ns_first.mean()), 100)

def test_lpca_rmse():
    S0_w_noise = 100 + 2 * np.random.standard_normal((22, 23, 30, 20))
    rmse_w_noise = np.sqrt(np.mean((S0_w_noise - 100) ** 2))
    S0_denoised = localpca(S0_w_noise, sigma=np.std(S0_w_noise))
    rmse_denoised = np.sqrt(np.mean((S0_denoised - 100) ** 2))
    # Denoising should always improve the RMSE:
    assert_(rmse_denoised < rmse_w_noise)

def test_lpca_sharpness():
    S0 = np.ones((30, 30, 30, 20), dtype=np.float64) * 100
    S0[10:20, 10:20, 10:20, :] = 50
    S0[20:30, 20:30, 20:30, :] = 0
    S0 = S0 + 20 * np.random.standard_normal((30, 30, 30, 20))
    S0ns = localpca(S0, sigma=20.0)
    # check the edge gradient
    edgs = np.abs(np.mean(S0ns[8, 10:20, 10:20] - S0ns[12, 10:20, 10:20]) - 50)
    assert_(edgs < 2)

def test_lpca_dtype():
    # If out_dtype is not specified, we retain the original precision:
    S0 = 200 * np.ones((20, 20, 20, 3), dtype=np.float64)
    S0ns = localpca(S0, sigma=1)
    assert_equal(S0.dtype, S0ns.dtype)

    S0 = 200 * np.ones((20, 20, 20, 20), dtype=np.uint16)
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20)))
    assert_equal(S0.dtype, S0ns.dtype)

    # If we set out_dtype, we get what we asked for:
    S0 = 200 * np.ones((20, 20, 20, 20), dtype=np.uint16)
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20)),
                    out_dtype=np.float32)
    assert_equal(np.float32, S0ns.dtype)

    # If we set a few entries to zero, this induces negative entries in the
    # Resulting denoised array:
    S0[5:8, 5:8, 5:8] = 0
    # But if we should always get all non-negative results:
    S0ns = localpca(S0, sigma=np.ones((20, 20, 20)), out_dtype=np.uint16)
    assert_(np.all(S0ns >= 0))
    # And no wrap-around to crazy high values:
    assert_(np.all(S0ns <= 200))

def test_lpca_wrong():
    S0 = np.ones((20, 20))
    assert_raises(ValueError, localpca, S0, sigma=1)

def test_phantom():
    gtab = gen_gtab()
    DWI_clean = rfiw_phantom(gtab, snr=None)
    DWI, sigma = rfiw_phantom(gtab, snr=30)
    # To test without Rician correction
    temp = (DWI_clean / sigma)**2
    DWI_clean_wrc = (sigma * np.sqrt(np.pi / 2) * np.exp(-0.5 * temp) *
                     ((1 + 0.5 * temp) * sps.iv(0, 0.25 * temp) + 0.5 * temp *
                     sps.iv(1, 0.25 * temp))**2)

    DWI_den = localpca(DWI, sigma, patch_radius=3)
    rmse_den = np.sum(np.abs(DWI_clean - DWI_den)) / np.sum(np.abs(DWI_clean))
    rmse_noisy = np.sum(np.abs(DWI_clean - DWI)) / np.sum(np.abs(DWI_clean))

    rmse_den_wrc = np.sum(np.abs(DWI_clean_wrc - DWI_den)
                          ) / np.sum(np.abs(DWI_clean_wrc))
    rmse_noisy_wrc = np.sum(np.abs(DWI_clean_wrc - DWI)) / \
        np.sum(np.abs(DWI_clean_wrc))

    assert_(np.max(DWI_clean) / sigma < np.max(DWI_den) / sigma)
    assert_(np.max(DWI_den) / sigma < np.max(DWI) / sigma)
    assert_(rmse_den < rmse_noisy)
    assert_(rmse_den_wrc < rmse_noisy_wrc)

    # Check if the results of different PCA methods (eig, svd) are similar
    DWI_den_svd = localpca(DWI, sigma, pca_method='svd', patch_radius=3)
    assert_array_almost_equal(DWI_den, DWI_den_svd)

    assert_raises(ValueError, localpca, DWI, sigma, pca_method='empty')

    # Try this with a sigma volume, instead of a scalar
    sigma_vol = sigma * np.ones(DWI.shape[:-1])
    mask = np.zeros_like(DWI, dtype=bool)[..., 0]
    mask[2:-2, 2:-2, 2:-2] = True
    DWI_den = localpca(DWI, sigma_vol, mask, patch_radius=3)
    DWI_clean_masked = DWI_clean.copy()
    DWI_clean_masked[~mask] = 0
    DWI_masked = DWI.copy()
    DWI_masked[~mask] = 0
    rmse_den = np.sum(np.abs(DWI_clean_masked - DWI_den)) / np.sum(np.abs(
            DWI_clean_masked))
    rmse_noisy = np.sum(np.abs(DWI_clean_masked - DWI_masked)) / np.sum(np.abs(
            DWI_clean_masked))

    DWI_clean_wrc_masked = DWI_clean_wrc.copy()
    DWI_clean_wrc_masked[~mask] = 0
    rmse_den_wrc = np.sum(np.abs(DWI_clean_wrc_masked - DWI_den)
                          ) / np.sum(np.abs(DWI_clean_wrc_masked))
    rmse_noisy_wrc = np.sum(np.abs(DWI_clean_wrc_masked - DWI_masked)) / \
        np.sum(np.abs(DWI_clean_wrc_masked))

    assert_(np.max(DWI_clean) / sigma < np.max(DWI_den) / sigma)
    assert_(np.max(DWI_den) / sigma < np.max(DWI) / sigma)
    assert_(rmse_den < rmse_noisy)
    assert_(rmse_den_wrc < rmse_noisy_wrc)

def test_lpca_ill_conditioned():
    gtab = gen_gtab()
    DWI, sigma = rfiw_phantom(gtab, snr=30)
    assert_raises(ValueError, localpca, DWI, sigma, patch_radius=1)

def test_lpca_sigma_wrong_shape():
    gtab = gen_gtab()
    DWI, sigma = rfiw_phantom(gtab, snr=30)
    # If sigma is 3D but shape is not like DWI.shape[:-1], an error is raised:
    sigma = np.zeros((DWI.shape[0], DWI.shape[1] + 1, DWI.shape[2]))
    assert_raises(ValueError, localpca, DWI, sigma)

def test_pca_classifier():
    # Produce small phantom with well aligned single voxels and ground truth
    # snr = 50, i.e signal std = 0.02 (Gaussian noise)
    std_gt = 0.02
    S0 = 1.0
    gtab = gen_gtab()
    ndir = gtab.bvals.size
    signal_test = np.zeros((5, 5, 5, ndir))
    mevals = np.array([[0.99e-3, 0.0, 0.0], [2.26e-3, 0.87e-3, 0.87e-3]])
    sig, direction = multi_tensor(gtab, mevals, S0=S0,
                                  angles=[(0, 0, 1), (0, 0, 1)],
                                  fractions=(50, 50), snr=None)
    signal_test[..., :] = sig
    noise = std_gt*np.random.standard_normal((5, 5, 5, ndir))
    dwi_test = signal_test + noise

    # Compute eigenvalues
    X = dwi_test.reshape(125, ndir)
    M = np.mean(X, axis=0)
    X = X - M
    [L, W] = np.linalg.eigh(np.dot(X.T, X)/125)

    # Find number of noise related eigenvalues
    c, var = pca_classifier(L, 125)
    std = np.sqrt(var)

    # Expected number of signal components is 1 because phantom only has one
    # voxel type. Therefore, expected noise components is (L.size - 1).
    # To allow some margin of error let's assess if c is higher than
    # L.size - 3. Note that algorithm is tuned to be conservative (i.e.
    # preserve the signal information as much as posible).
    assert_(c > L.size-3)

    # Let's check if noise std estimate as an error less than 5%
    std_error = abs(std - std_gt)/std_gt * 100
    assert_(std_error < 10)

if __name__ == '__main__':
    run_module_suite()
<|endoftext|>