text:
# -*- encoding: utf-8 -*-
'''Display functions'''

from collections import defaultdict

import numpy as np
from scipy.signal import spectrogram

from matplotlib.patches import Rectangle
from matplotlib.ticker import FuncFormatter, MultipleLocator
from matplotlib.ticker import Formatter
from matplotlib.colors import LinearSegmentedColormap, LogNorm, ColorConverter
from matplotlib.collections import BrokenBarHCollection

from .melody import freq_to_voicing
from .util import midi_to_hz, hz_to_midi

def __expand_limits(ax, limits, which='x'):
    '''Helper function to expand axis limits'''

    if which == 'x':
        getter, setter = ax.get_xlim, ax.set_xlim
    elif which == 'y':
        getter, setter = ax.get_ylim, ax.set_ylim
    else:
        raise ValueError('invalid axis: {}'.format(which))

    old_lims = getter()
    new_lims = list(limits)

    # infinite limits occur on new axis objects with no data
    if np.isfinite(old_lims[0]):
        new_lims[0] = min(old_lims[0], limits[0])

    if np.isfinite(old_lims[1]):
        new_lims[1] = max(old_lims[1], limits[1])

    setter(new_lims)

def __get_axes(ax=None, fig=None):
    '''Get or construct the target axes object for a new plot.

    Parameters
    ----------
    ax : matplotlib.pyplot.axes, optional
        If provided, return this axes object directly.

    fig : matplotlib.figure.Figure, optional
        The figure to query for axes.

        By default, uses the current figure `plt.gcf()`.

    Returns
    -------
    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the segmentation.
        If none is provided, a new set of axes is created.

    new_axes : bool
        If `True`, the axis object was newly constructed.
        If `False`, the axis object already existed.

    '''

    new_axes = False

    if ax is not None:
        return ax, new_axes

    if fig is None:
        import matplotlib.pyplot as plt
        fig = plt.gcf()

    if not fig.get_axes():
        new_axes = True

    return fig.gca(), new_axes

def segments(intervals, labels, base=None, height=None, text=False,
             text_kw=None, ax=None, **kwargs):
    '''Plot a segmentation as a set of disjoint rectangles.

    Parameters
    ----------
    intervals : np.ndarray, shape=(n, 2)
        segment intervals, in the format returned by
        :func:`mir_eval.io.load_intervals` or
        :func:`mir_eval.io.load_labeled_intervals`.

    labels : list, shape=(n,)
        reference segment labels, in the format returned by
        :func:`mir_eval.io.load_labeled_intervals`.

    base : number
        The vertical position of the base of the rectangles.
        By default, this will be the bottom of the plot.

    height : number
        The height of the rectangles.
        By default, this will be the top of the plot (minus ``base``).

    text : bool
        If true, each segment's label is displayed in its
        upper-left corner

    text_kw : dict
        If ``text == True``, the properties of the text
        object can be specified here.
        See ``matplotlib.pyplot.Text`` for valid parameters

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the segmentation.
        If none is provided, a new set of axes is created.

    kwargs
        Additional keyword arguments to pass to
        ``matplotlib.patches.Rectangle``.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''
    if text_kw is None:
        text_kw = dict()
    text_kw.setdefault('va', 'top')
    text_kw.setdefault('clip_on', True)
    text_kw.setdefault('bbox', dict(boxstyle='round', facecolor='white'))

    # Make sure we have a numpy array
    intervals = np.atleast_2d(intervals)

    seg_def_style = dict(linewidth=1)

    ax, new_axes = __get_axes(ax=ax)

    if new_axes:
        ax.set_ylim([0, 1])

    # Infer height
    if base is None:
        base = ax.get_ylim()[0]

    if height is None:
        height = ax.get_ylim()[1]

    cycler = ax._get_patches_for_fill.prop_cycler

    seg_map = dict()

    for lab in labels:
        if lab in seg_map:
            continue

        style = next(cycler)
        seg_map[lab] = seg_def_style.copy()
        seg_map[lab].update(style)
        # Swap color -> facecolor here so we preserve edgecolor on rects
        seg_map[lab]['facecolor'] = seg_map[lab].pop('color')
        seg_map[lab].update(kwargs)
        seg_map[lab]['label'] = lab

    for ival, lab in zip(intervals, labels):
        rect = Rectangle((ival[0], base), ival[1] - ival[0], height,
                         **seg_map[lab])
        ax.add_patch(rect)
        seg_map[lab].pop('label', None)

        if text:
            ann = ax.annotate(lab,
                              xy=(ival[0], height), xycoords='data',
                              xytext=(8, -10), textcoords='offset points',
                              **text_kw)
            ann.set_clip_path(rect)

    if new_axes:
        ax.set_yticks([])

    # Only expand if we have data
    if intervals.size:
        __expand_limits(ax, [intervals.min(), intervals.max()], which='x')

    return ax

def labeled_intervals(intervals, labels, label_set=None,
                      base=None, height=None, extend_labels=True,
                      ax=None, tick=True, **kwargs):
    '''Plot labeled intervals with each label on its own row.

    Parameters
    ----------
    intervals : np.ndarray, shape=(n, 2)
        segment intervals, in the format returned by
        :func:`mir_eval.io.load_intervals` or
        :func:`mir_eval.io.load_labeled_intervals`.

    labels : list, shape=(n,)
        reference segment labels, in the format returned by
        :func:`mir_eval.io.load_labeled_intervals`.

    label_set : list
        An (ordered) list of labels to determine the plotting order.
        If not provided, the labels will be inferred from
        ``ax.get_yticklabels()``.
        If no ``yticklabels`` exist, then the sorted set of unique values
        in ``labels`` is taken as the label set.

    base : np.ndarray, shape=(n,), optional
        Vertical positions of each label.
        By default, labels are positioned at integers
        ``np.arange(len(labels))``.

    height : scalar or np.ndarray, shape=(n,), optional
        Height for each label.
        If scalar, the same value is applied to all labels.
        By default, each label has ``height=1``.

    extend_labels : bool
        If ``False``, only values of ``labels`` that also exist in
        ``label_set`` will be shown.

        If ``True``, all labels are shown, with those in `labels` but
        not in `label_set` appended to the top of the plot.
        A horizontal line is drawn to indicate the separation between
        values in or out of ``label_set``.

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the intervals.
        If none is provided, a new set of axes is created.

    tick : bool
        If ``True``, sets tick positions and labels on the y-axis.

    kwargs
        Additional keyword arguments to pass to
        `matplotlib.collection.BrokenBarHCollection`.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''

    # Get the axes handle
    ax, _ = __get_axes(ax=ax)

    # Make sure we have a numpy array
    intervals = np.atleast_2d(intervals)

    if label_set is None:
        # If we have non-empty pre-existing tick labels, use them
        label_set = [_.get_text() for _ in ax.get_yticklabels()]
        # If none of the label strings have content, treat it as empty
        if not any(label_set):
            label_set = []
    else:
        label_set = list(label_set)

    # Put additional labels at the end, in order
    if extend_labels:
        ticks = label_set + sorted(set(labels) - set(label_set))
    elif label_set:
        ticks = label_set
    else:
        ticks = sorted(set(labels))

    style = dict(linewidth=1)

    style.update(next(ax._get_patches_for_fill.prop_cycler))
    # Swap color -> facecolor here so we preserve edgecolor on rects
    style['facecolor'] = style.pop('color')
    style.update(kwargs)

    if base is None:
        base = np.arange(len(ticks))

    if height is None:
        height = 1

    if np.isscalar(height):
        height = height * np.ones_like(base)

    seg_y = dict()
    for ybase, yheight, lab in zip(base, height, ticks):
        seg_y[lab] = (ybase, yheight)

    xvals = defaultdict(list)
    for ival, lab in zip(intervals, labels):
        if lab not in seg_y:
            continue
        xvals[lab].append((ival[0], ival[1] - ival[0]))

    for lab in seg_y:
        ax.add_collection(BrokenBarHCollection(xvals[lab], seg_y[lab],
                                               **style))
        # Pop the label after the first time we see it, so we only get
        # one legend entry
        style.pop('label', None)

    # Draw a line separating the new labels from pre-existing labels
    if label_set != ticks:
        ax.axhline(len(label_set), color='k', alpha=0.5)

    if tick:
        ax.grid(True, axis='y')
        ax.set_yticks([])
        ax.set_yticks(base)
        ax.set_yticklabels(ticks, va='bottom')
        ax.yaxis.set_major_formatter(IntervalFormatter(base, ticks))

    if base.size:
        __expand_limits(ax, [base.min(), (base + height).max()], which='y')
    if intervals.size:
        __expand_limits(ax, [intervals.min(), intervals.max()], which='x')

    return ax

class IntervalFormatter(Formatter):
    '''Ticker formatter for labeled interval plots.

    Parameters
    ----------
    base : array-like of int
        The base positions of each label

    ticks : array-like of string
        The labels for the ticks
    '''
    def __init__(self, base, ticks):

        self._map = {int(k): v for k, v in zip(base, ticks)}

    def __call__(self, x, pos=None):

        return self._map.get(int(x), '')

def hierarchy(intervals_hier, labels_hier, levels=None, ax=None, **kwargs):
    '''Plot a hierarchical segmentation

    Parameters
    ----------
    intervals_hier : list of np.ndarray
        A list of segmentation intervals.  Each element should be
        an n-by-2 array of segment intervals, in the format returned by
        :func:`mir_eval.io.load_intervals` or
        :func:`mir_eval.io.load_labeled_intervals`.
        Segmentations should be ordered by increasing specificity.

    labels_hier : list of list-like
        A list of segmentation labels.  Each element should
        be a list of labels for the corresponding element in
        `intervals_hier`.

    levels : list of string
        Each element ``levels[i]`` is a label for the ```i`` th segmentation.
        This is used in the legend to denote the levels in a segment hierarchy.

    kwargs
        Additional keyword arguments to `labeled_intervals`.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''

    # This will break if a segment label exists in multiple levels
    if levels is None:
        levels = list(range(len(intervals_hier)))

    # Get the axes handle
    ax, _ = __get_axes(ax=ax)

    # Count the pre-existing patches
    n_patches = len(ax.patches)

    for ints, labs, key in zip(intervals_hier[::-1],
                               labels_hier[::-1],
                               levels[::-1]):
        labeled_intervals(ints, labs, label=key, ax=ax, **kwargs)

    # Reverse the patch ordering for anything we've added.
    # This way, intervals are listed in the legend from top to bottom
    ax.patches[n_patches:] = ax.patches[n_patches:][::-1]
    return ax

def events(times, labels=None, base=None, height=None, ax=None, text_kw=None,
           **kwargs):
    '''Plot event times as a set of vertical lines

    Parameters
    ----------
    times : np.ndarray, shape=(n,)
        event times, in the format returned by
        :func:`mir_eval.io.load_events` or
        :func:`mir_eval.io.load_labeled_events`.

    labels : list, shape=(n,), optional
        event labels, in the format returned by
        :func:`mir_eval.io.load_labeled_events`.

    base : number
        The vertical position of the base of the line.
        By default, this will be the bottom of the plot.

    height : number
        The height of the lines.
        By default, this will be the top of the plot (minus `base`).

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the segmentation.
        If none is provided, a new set of axes is created.

    text_kw : dict
        If `labels` is provided, the properties of the text
        objects can be specified here.
        See `matplotlib.pyplot.Text` for valid parameters

    kwargs
        Additional keyword arguments to pass to
        `matplotlib.pyplot.vlines`.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''
    if text_kw is None:
        text_kw = dict()
    text_kw.setdefault('va', 'top')
    text_kw.setdefault('clip_on', True)
    text_kw.setdefault('bbox', dict(boxstyle='round', facecolor='white'))

    # make sure we have an array for times
    times = np.asarray(times)

    # Get the axes handle
    ax, new_axes = __get_axes(ax=ax)

    # If we have fresh axes, set the limits

    if new_axes:
        # Infer base and height
        if base is None:
            base = 0
        if height is None:
            height = 1

        ax.set_ylim([base, height])
    else:
        if base is None:
            base = ax.get_ylim()[0]

        if height is None:
            height = ax.get_ylim()[1]

    cycler = ax._get_patches_for_fill.prop_cycler

    style = next(cycler).copy()
    style.update(kwargs)
    # If the user provided 'colors', don't override it with 'color'
    if 'colors' in style:
        style.pop('color', None)

    lines = ax.vlines(times, base, base + height, **style)

    if labels:
        for path, lab in zip(lines.get_paths(), labels):
            ax.annotate(lab,
                        xy=(path.vertices[0][0], height),
                        xycoords='data',
                        xytext=(8, -10), textcoords='offset points',
                        **text_kw)

    if new_axes:
        ax.set_yticks([])

    __expand_limits(ax, [base, base + height], which='y')

    if times.size:
        __expand_limits(ax, [times.min(), times.max()], which='x')

    return ax

def pitch(times, frequencies, midi=False, unvoiced=False, ax=None, **kwargs):
    '''Visualize pitch contours

    Parameters
    ----------
    times : np.ndarray, shape=(n,)
        Sample times of frequencies

    frequencies : np.ndarray, shape=(n,)
        frequencies (in Hz) of the pitch contours.
        Voicing is indicated by sign (positive for voiced,
        non-positive for non-voiced).

    midi : bool
        If `True`, plot on a MIDI-numbered vertical axis.
        Otherwise, plot on a linear frequency axis.

    unvoiced : bool
        If `True`, unvoiced pitch contours are plotted and indicated
        by transparency.

        Otherwise, unvoiced pitch contours are omitted from the display.

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the pitch contours.
        If none is provided, a new set of axes is created.

    kwargs
        Additional keyword arguments to `matplotlib.pyplot.plot`.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''

    ax, _ = __get_axes(ax=ax)

    times = np.asarray(times)

    # First, segment into contiguously voiced contours
    frequencies, voicings = freq_to_voicing(np.asarray(frequencies,
                                                       dtype=np.float))

    # Here are all the change-points
    v_changes = 1 + np.flatnonzero(voicings[1:] != voicings[:-1])
    v_changes = np.unique(np.concatenate([[0], v_changes, [len(voicings)]]))

    # Set up arrays of slices for voiced and unvoiced regions
    v_slices, u_slices = [], []
    for start, end in zip(v_changes, v_changes[1:]):
        idx = slice(start, end)
        # A region is voiced if its starting sample is voiced
        # It's unvoiced if none of the samples in the region are voiced.
        if voicings[start]:
            v_slices.append(idx)
        elif frequencies[idx].all():
            u_slices.append(idx)

    # Now we just need to plot the contour
    style = dict()
    style.update(next(ax._get_lines.prop_cycler))
    style.update(kwargs)

    if midi:
        idx = frequencies > 0
        frequencies[idx] = hz_to_midi(frequencies[idx])

        # Tick at integer midi notes
        ax.yaxis.set_minor_locator(MultipleLocator(1))

    for idx in v_slices:
        ax.plot(times[idx], frequencies[idx], **style)
        style.pop('label', None)

    # Plot the unvoiced portions
    if unvoiced:
        style['alpha'] = style.get('alpha', 1.0) * 0.5
        for idx in u_slices:
            ax.plot(times[idx], frequencies[idx], **style)

    return ax

def multipitch(times, frequencies, midi=False, unvoiced=False, ax=None,
               **kwargs):
    '''Visualize multiple f0 measurements

    Parameters
    ----------
    times : np.ndarray, shape=(n,)
        Sample times of frequencies

    frequencies : list of np.ndarray
        frequencies (in Hz) of the pitch measurements.
        Voicing is indicated by sign (positive for voiced,
        non-positive for non-voiced).

        `times` and `frequencies` should be in the format produced by
        :func:`mir_eval.io.load_ragged_time_series`

    midi : bool
        If `True`, plot on a MIDI-numbered vertical axis.
        Otherwise, plot on a linear frequency axis.

    unvoiced : bool
        If `True`, unvoiced pitches are plotted and indicated
        by transparency.

        Otherwise, unvoiced pitches are omitted from the display.

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the pitch contours.
        If none is provided, a new set of axes is created.

    kwargs
        Additional keyword arguments to `plt.scatter`.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''

    # Get the axes handle
    ax, _ = __get_axes(ax=ax)

    # Set up a style for the plot
    style_voiced = dict()
    style_voiced.update(next(ax._get_lines.prop_cycler))
    style_voiced.update(kwargs)

    style_unvoiced = style_voiced.copy()
    style_unvoiced.pop('label', None)
    style_unvoiced['alpha'] = style_unvoiced.get('alpha', 1.0) * 0.5

    # We'll collect all times and frequencies first, then plot them
    voiced_times = []
    voiced_freqs = []

    unvoiced_times = []
    unvoiced_freqs = []

    for t, freqs in zip(times, frequencies):
        if not len(freqs):
            continue

        freqs, voicings = freq_to_voicing(np.asarray(freqs, dtype=np.float))

        # Discard all 0-frequency measurements
        idx = freqs > 0
        freqs = freqs[idx]
        voicings = voicings[idx]

        if midi:
            freqs = hz_to_midi(freqs)

        n_voiced = sum(voicings)
        voiced_times.extend([t] * n_voiced)
        voiced_freqs.extend(freqs[voicings])
        unvoiced_times.extend([t] * (len(freqs) - n_voiced))
        unvoiced_freqs.extend(freqs[~voicings])

    # Plot the voiced frequencies
    ax.scatter(voiced_times, voiced_freqs, **style_voiced)

    # Plot the unvoiced frequencies
    if unvoiced:
        ax.scatter(unvoiced_times, unvoiced_freqs, **style_unvoiced)

    # Tick at integer midi notes
    if midi:
        ax.yaxis.set_minor_locator(MultipleLocator(1))
    return ax

def piano_roll(intervals, pitches=None, midi=None, ax=None, **kwargs):
    '''Plot a quantized piano roll as intervals

    Parameters
    ----------
    intervals : np.ndarray, shape=(n, 2)
        timing intervals for notes

    pitches : np.ndarray, shape=(n,), optional
        pitches of notes (in Hz).

    midi : np.ndarray, shape=(n,), optional
        pitches of notes (in MIDI numbers).

        At least one of ``pitches`` or ``midi`` must be provided.

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the intervals.
        If none is provided, a new set of axes is created.

    kwargs
        Additional keyword arguments to :func:`labeled_intervals`.

    Returns
    -------
    ax : matplotlib.pyplot.axes._subplots.AxesSubplot
        A handle to the (possibly constructed) plot axes
    '''

    if midi is None:
        if pitches is None:
            raise ValueError('At least one of `midi` or `pitches` '
                             'must be provided.')

        midi = hz_to_midi(pitches)

    scale = np.arange(128)
    ax = labeled_intervals(intervals, np.round(midi).astype(int),
                           label_set=scale,
                           tick=False,
                           ax=ax,
                           **kwargs)

    # Minor tick at each semitone
    ax.yaxis.set_minor_locator(MultipleLocator(1))

    ax.axis('auto')
    return ax

def separation(sources, fs=22050, labels=None, alpha=0.75, ax=None, **kwargs):
    '''Source-separation visualization

    Parameters
    ----------
    sources : np.ndarray, shape=(nsrc, nsampl)
        A list of waveform buffers corresponding to each source

    fs : number > 0
        The sampling rate

    labels : list of strings
        An optional list of descriptors corresponding to each source

    alpha : float in [0, 1]
        Maximum alpha (opacity) of spectrogram values.

    ax : matplotlib.pyplot.axes
        An axis handle on which to draw the spectrograms.
        If none is provided, a new set of axes is created.

    kwargs
        Additional keyword arguments to ``scipy.signal.spectrogram``

    Returns
    -------
    ax
        The axis handle for this plot
    '''

    # Get the axes handle
    ax, new_axes = __get_axes(ax=ax)

    # Make sure we have at least two dimensions
    sources = np.atleast_2d(sources)

    if labels is None:
        labels = ['Source {:d}'.format(_) for _ in range(len(sources))]

    kwargs.setdefault('scaling', 'spectrum')

    # The cumulative spectrogram across sources
    # is used to establish the reference power
    # for each individual source
    cumspec = None
    specs = []
    for i, src in enumerate(sources):
        freqs, times, spec = spectrogram(src, fs=fs, **kwargs)

        specs.append(spec)
        if cumspec is None:
            cumspec = spec.copy()
        else:
            cumspec += spec

    ref_max = cumspec.max()
    ref_min = ref_max * 1e-6

    color_conv = ColorConverter()

    for i, spec in enumerate(specs):

        # For each source, grab a new color from the cycler
        # Then construct a colormap that interpolates from
        # [transparent white -> new color]
        color = next(ax._get_lines.prop_cycler)['color']
        color = color_conv.to_rgba(color, alpha=alpha)
        cmap = LinearSegmentedColormap.from_list(labels[i],
                                                 [(1.0, 1.0, 1.0, 0.0),
                                                  color])

        ax.pcolormesh(times, freqs, spec,
                      cmap=cmap,
                      norm=LogNorm(vmin=ref_min, vmax=ref_max),
                      shading='gouraud',
                      label=labels[i])

        # Attach a 0x0 rect to the axis with the corresponding label
        # This way, it will show up in the legend
        ax.add_patch(Rectangle((0, 0), 0, 0, color=color, label=labels[i]))

    if new_axes:
        ax.axis('tight')

    return ax

def __ticker_midi_note(x, pos):
    '''A ticker function for midi notes.

    Inputs x are interpreted as midi numbers, and converted
    to [NOTE][OCTAVE]+[cents].
    '''

    NOTES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']

    cents = float(np.mod(x, 1.0))
    if cents >= 0.5:
        cents = cents - 1.0
        x = x + 0.5

    idx = int(x % 12)

    octave = int(x / 12) - 1

    if cents == 0:
        return '{:s}{:2d}'.format(NOTES[idx], octave)
    return '{:s}{:2d}{:+02d}'.format(NOTES[idx], octave, int(cents * 100))

def __ticker_midi_hz(x, pos):
    '''A ticker function for midi pitches.

    Inputs x are interpreted as midi numbers, and converted
    to Hz.
    '''

    return '{:g}'.format(midi_to_hz(x))

def ticker_notes(ax=None):
    '''Set the y-axis of the given axes to MIDI notes

    Parameters
    ----------
    ax : matplotlib.pyplot.axes
        The axes handle to apply the ticker.
        By default, uses the current axes handle.

    '''
    ax, _ = __get_axes(ax=ax)

    ax.yaxis.set_major_formatter(FMT_MIDI_NOTE)
    # Get the tick labels and reset the vertical alignment
    for tick in ax.yaxis.get_ticklabels():
        tick.set_verticalalignment('baseline')

def ticker_pitch(ax=None):
    '''Set the y-axis of the given axes to MIDI frequencies

    Parameters
    ----------
    ax : matplotlib.pyplot.axes
        The axes handle to apply the ticker.
        By default, uses the current axes handle.
    '''
    ax, _ = __get_axes(ax=ax)

    ax.yaxis.set_major_formatter(FMT_MIDI_HZ)

# Instantiate ticker objects; we don't need more than one of each
FMT_MIDI_NOTE = FuncFormatter(__ticker_midi_note)
FMT_MIDI_HZ = FuncFormatter(__ticker_midi_hz)

reason: import numpy,from scipy


text:
import sys
from Visualize import Visualizer
from mots_common.io import load_sequences, load_seqmap, load_txt
import pycocotools.mask as rletools
import glob
import os
import cv2
import colorsys
import numpy as np

def apply_mask(image, mask, color, alpha=0.5):
	"""
	 Apply the given mask to the image.
	"""
	for c in range(3):
		image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - alpha) + alpha * color[c],image[:, :, c])
	return image

class MOTSVisualizer(Visualizer):

	def load(self, FilePath):
		return load_txt(FilePath)

	def drawResults(self, im = None, t = 0):
		self.draw_boxes = False

		for obj in self.resFile[t]:

			color = self.colors[obj.track_id % len(self.colors)]

			color = tuple([int(c*255) for c in color])

			if obj.class_id == 1:
				category_name = "Car"
			elif obj.class_id == 2:
				category_name = "Ped"
			else:
				category_name = "Ignore"
				color = (0.7*255, 0.7*255, 0.7*255)

			if obj.class_id == 1 or obj.class_id == 2:  # Don't show boxes or ids for ignore regions
				x, y, w, h = rletools.toBbox(obj.mask)

				pt1=(int(x),int(y))
				pt2=(int(x+w),int(y+h))

				category_name += ":" + str(obj.track_id)
				cv2.putText(im, category_name, (int(x + 0.5 * w), int( y + 0.5 * h)), cv2.FONT_HERSHEY_TRIPLEX,self.imScale,color,thickness =2)
				if self.draw_boxes:
					cv2.rectangle(im,pt1,pt2,color,2)

			binary_mask = rletools.decode(obj.mask)

			im = apply_mask(im, binary_mask, color)
		return im

if __name__ == "__main__":
	visualizer = MOTSVisualizer(
	seqName = "MOTS20-11",
	FilePath ="data/MOTS/train/MOTS20-11/gt/gt.txt",
	image_dir = "data/MOTS/train/MOTS20-11/img1",
	mode = "gt",
	output_dir  = "vid")

	visualizer.generateVideo(
	        displayTime = True,
	        displayName = "seg",
	        showOccluder = True,
	        fps = 25 )

ext: py

max_stars_repo_name: collector-m/UniTrack

max_issues_repo_name: collector-m/UniTrack

max_forks_repo_name: collector-m/UniTrack

reason: import numpy


text:
import time
import sys
import shutil
import numpy as np
import os
import warnings
warnings.filterwarnings(action='ignore',category=UserWarning,module='gensim')
from PyQt5.QtWidgets import QFileDialog
from PyQt5.uic import loadUi
from PyQt5.QtWidgets import QMainWindow, QApplication
import jieba.posseg as pseg
import codecs#文件读取时，编码转化
from gensim import corpora, models, similarities#运用Gensim建立词典，生成BOW语料
#运行tfidf模型计算词权重，采用LsiModel进行降维，最后运用Gensim提供的MatrixSimilarity类来计算两文档的相似性【基于余弦的距离的计算】
import docx2txt#读取docx文档
import xlwt
import easygui
from QCandyUi.CandyWindow import colorful
@colorful('blueGreen')
class MainWindow(QMainWindow):
    def __init__(self):
        QMainWindow.__init__(self)
        loadUi("check.ui", self)
        global FilesName
        global Directory
        global Name
        global threshold

        #botton 3
        FilesName =self.pushButton_3.clicked.connect(self._openFile_)
        #botton 2
        Directory=str(self.pushButton_2.clicked.connect(self._openDir_))
        #botton 1

        self.pushButton.clicked.connect(lambda:self._check_(FilesName,Directory))

    def _split_(self,file_path):
        (filepath, tempfilename) = os.path.split(file_path)
        (filename, extension) = os.path.splitext(tempfilename)
        return filename

    def _check_(self,FilesName,Directory):
        #报表名前后无空格
        global Name
        Name=str(self.lineEdit_2.text()).strip()+'.xls'
        global threshold
        threshold=70 #default

        s=self.lineEdit.text().strip()
        if s!='':
            threshold=int(s)

        style = xlwt.easyxf('pattern: pattern solid, fore_colour red')
        stop_words = 'cn_stopwords.txt'#停用词路径
        stopwords = codecs.open(stop_words,'r',encoding='utf8').readlines()#读取停用词
        stopwords = [ w.strip() for w in stopwords ] # 去除首尾空格
        stop_flag = ['x', 'c', 'u','d', 'p', 't', 'uj', 'm', 'f', 'r']# 停用词性

        #对一篇文章分词、去停用词
        def tokenization(file):

            result = []
            text = file
            words = pseg.cut(text)#分词
            #去除 flag在stop_flag，word 在stopword，所对应的词语（即除二者意外其他加入了result中）
            for word, flag in words:
                if flag not in stop_flag and word not in stopwords:
                    result.append(word)
            return result

            #对目录下的所有文本进行预处理，构建字典
        start =time.clock()
        filenames=[]
        for i in FilesName:
            filenames.append(docx2txt.process(i))

        corpus = []
        for each in filenames:
            corpus.append(tokenization(each))

        dictionary = corpora.Dictionary(corpus)#为每个出现在词库中的词语分配一个独一无二的整数ID，这个操作收集单词计数和其他统计信息

        #建立词袋模型
        #  生成词向量
        doc_vectors = [dictionary.doc2bow(text) for text in corpus] # 语料库

        # 生成TF-IDF模型
        tfidf = models.TfidfModel(doc_vectors)# TF-IDF模型对语料库建模
        tfidf_vectors = tfidf[doc_vectors] # 每个词的TF-IDF值

        #把稀疏的高级矩阵变成一个计算起来会比较轻松的小矩阵，也把一些没有用的燥音给过滤掉了，这个模型可以被后来的语料查询与分类所调用
        lsi = models.LsiModel(tfidf_vectors, id2word=dictionary, num_topics=300)
        #num_topics 主题数
        #将Tf-Idf语料转化为一个潜在2-D空间（2-D是因为我们设置了num_topics=2）

        lsi.print_topics(300)# 建立的两个主题模型内容

        # 将文章投影到主题空间中
        lsi_vector = lsi[tfidf_vectors]
        count1=0
        count=1
        workbook=xlwt.Workbook(encoding="utf-8")
        worksheet=workbook.add_sheet(Name)
        worksheet.write(0,0,'查找文档名')
        worksheet.write(0,1,'对比文档名')
        worksheet.write(0,2,'重复率')
        for i in FilesName:

            text = docx2txt.process(i)#打开文档
            query = tokenization(text)#分词，去词
            query_bow = dictionary.doc2bow(query)#语义库

            query_lsi = lsi[query_bow]

            #相似矩阵计算相似度
            index = similarities.MatrixSimilarity(lsi_vector)#和第个文档相似
            sims = index[query_lsi]

            sort_sims = sorted(enumerate(sims), key=lambda item: -item[1])

            worksheet.write(count,0,self._split_(i))
            worksheet.write(count,1,self._split_(FilesName[sort_sims[1][0]]))
            if np.float32(sort_sims[1][1]).item()>(float(threshold)*0.01):
                 worksheet.write(count,2,np.float32(sort_sims[1][1]).item(),style)
                 count1=count1+1
            else:
                 worksheet.write(count,2,np.float32(sort_sims[1][1]).item())

            count=count+1

        workbook.save(Name)

        aa=os.getcwd()
        file_path=os.path.join(aa,Name)
        shutil.move(file_path,Directory)

        end = time.clock()
        print('Running time: %s Seconds'%(end-start))
        easygui.msgbox("完成啦！", title="文本查重",ok_button="知道啦")

    #button3 选择文件，返回list类型的文件路径

    def _openFile_(self):
        global FilesName

        FilesName,filetype = QFileDialog.getOpenFileNames(None, "请选择要添加的文件", "C:\\Users\\14866\\Desktop\\NLP\\Untitled Folder", "Text Files (*.docx);;Text Files (*.doc);;All Files (*)")
        return FilesName

    def _openDir_(self):
        global Directory
        Directory = QFileDialog.getExistingDirectory(self, "选择文件夹", "/")
        return Directory

app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())

max_stars_repo_name: aleeyanger/textual-replicas-detection

max_issues_repo_name: aleeyanger/textual-replicas-detection

max_forks_repo_name: aleeyanger/textual-replicas-detection


text:
import numpy as np
from numba import jit, int32, float32, double, cfunc
from numba.experimental import jitclass

spec = [
    ('x', double[:]), ('dq', double[:]), ('u', double[:]),
    ('m', double), ('Iz', double),
    ('lf', double), ('lr', double),
    ('Bf', double), ('Cf', double), ('Df', double),
    ('Br', double), ('Cr', double), ('Dr', double),
    ('Cr0', double), ('Cr2', double),
    ('Cm1', double), ('Cm2', double),
    ('iterNum', int32), ('sim_dt', double), ('control_dt', double),
    ('car_shape', double[:,:]),
]

@jitclass(spec)
class VehicleSimModel(object):
    def __init__(self,  m=0.041, Iz=27.8E-6,
                        lf=0.029, lr=0.033,
                        Bf=2.579, Cf=1.2, Df=0.192,
                        Br=3.3852, Cr=1.2691, Dr=0.1737,
                        Cm1=0.287, Cm2=0.0545,
                        Cr0= 0.0518,
                        Cr2=0.00035, scale=1.0, control_dt = 10.0, sim_dt=1.0):

        self.x = np.asfortranarray(np.zeros(3, dtype=np.float64))
        self.dq = np.zeros(3, dtype=np.float64)
        self.u = np.zeros(2, dtype=np.float64)

        self.m = m
        self.Iz= Iz

        self.lf = lf
        self.lr = lr

        self.Bf = Bf
        self.Cf = Cf
        self.Df = Df
        self.Br = Br
        self.Cr = Cr
        self.Dr = Dr

        self.Cr0 = Cr0
        self.Cr2 = Cr2

        self.Cm1 = Cm1
        self.Cm2 = Cm2

        car_l = (lf + lr)/2 * scale
        car_w = car_l/2
        self.car_shape = np.asfortranarray(np.array([ [ car_l, car_w, 1.],
                                    [ car_l,-car_w, 1.],
                                    [-car_l,-car_w, 1.],
                                    [-car_l, car_w, 1.],
                                    [ car_l, car_w, 1.]], dtype=np.float64))
        self.sim_dt     = sim_dt
        self.control_dt = control_dt
        self.iterNum = int(self.control_dt/self.sim_dt)

    @property
    def shape(self):
        shape = np.dot(self.car_shape,
                       np.asfortranarray(np.array([
                           [ np.cos(self.x[2]),  np.sin(self.x[2]), 0.],
                           [-np.sin(self.x[2]),  np.cos(self.x[2]), 0.],
                           [ self.x[0]        ,  self.x[1]        , 1.]], dtype=np.float64)))
        return shape[:,:2]

    def ODE_rh_eq(self, x, dq, u):
        return self.ODE_rh_eq_x(x, dq, u), self.ODE_rh_eq_dq(x, dq, u)

    def ODE_rh_eq_dq(self, x, dq, u):
        ddq = np.zeros_like(dq)

        alpha_f =  np.arctan2(dq[1] + self.lf * dq[2], dq[0]) - u[1]
        alpha_r =  np.arctan2(dq[1] - self.lr * dq[2], dq[0])

        Ffy = - self.Df * np.sin(self.Cf*np.arctan(self.Bf*alpha_f))
        Fry = - self.Dr * np.sin(self.Cr*np.arctan(self.Br*alpha_r))
        Frx = self.Cm1*u[0]-self.Cm2*u[0]*dq[0]-self.Cr0-self.Cr2*dq[0]**2

        ddq[0] =  dq[1] * dq[2] + 1.0/self.m * (Frx - Ffy*np.sin(u[1]) )
        ddq[1] = -dq[0] * dq[2] + 1.0/self.m * (Fry + Ffy*np.cos(u[1]) )
        ddq[2] = 1.0/self.Iz * (Ffy * self.lf * np.cos(u[1]) - Fry * self.lr)

        return ddq

    def ODE_rh_eq_x(self, x, dq, u):
        dx = np.zeros_like(x)
        dx[0:2] = np.dot(np.array([[np.cos(x[2]), -np.sin(x[2])],
                                   [np.sin(x[2]),  np.cos(x[2])]]),
                         dq[0:2])
        dx[2] = dq[2]
        return dx

    def RK4(self, u, dt):
        k1x, k1dq = self.ODE_rh_eq(self.x, self.dq, u)
        k2x, k2dq = self.ODE_rh_eq(self.x+k1x*dt/2, self.dq+k1dq*dt/2, u)
        k3x, k3dq = self.ODE_rh_eq(self.x+k2x*dt/2, self.dq+k2dq*dt/2, u)
        k4x, k4dq = self.ODE_rh_eq(self.x+k3x*dt/2, self.dq+k3dq*dt/2, u)

        self.x  = self.x  + dt * (k1x/6 + k2x/3 + k3x/3 + k4x/6)
        self.dq = self.dq + dt * (k1dq/6 + k2dq/3 + k3dq/3 + k4dq/6)

    def LeapFrog(self, u, dt):
        dx = self.ODE_rh_eq_x(self.x, self.dq, u)
        ddq = self.ODE_rh_eq_dq(self.x, self.dq, u)
        self.x = self.x + dx * dt + ddq*dt**2/2

        ddq2= self.ODE_rh_eq_dq(self.x, self.dq, u)
        self.dq = self.dq + (ddq + ddq2)*dt/2

    def sim_step(self, u):
        for i in range(self.iterNum):
            self.RK4(u, self.sim_dt/1000)
            #self.LeapFrog(u, self.sim_dt/1000)

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import matplotlib.animation as animation

    v = VehicleSimModel(scale=0.01)

    v.dq[0] = 0.2
    v.u[1] = np.deg2rad(20)

    x = np.array([v.x[0]])
    y = np.array([v.x[1]])
    car = []
    vhit = []

    for i in range(100):
        v.sim_step(v.u)
        x = np.append(x, v.x[0])
        y = np.append(y, v.x[1])

        carX, carY = v.shape.T
        car.append((carX, carY))
        vhit.append(v.dq[0])

    plt.figure(0)
    plt.plot(x, y, "ob-", label="trajectory")
    for carX, carY in car:
        plt.plot(carX, carY, 'k-')
    plt.axis("equal")
    plt.figure(1)
    plt.plot(range(len(vhit)), vhit)
    plt.grid(True)
    plt.show()

max_stars_repo_name: Chachay/MPCC

max_issues_repo_name: Chachay/MPCC

max_forks_repo_name: Chachay/MPCC

reason: import numpy,from numba


text:
import cv2
import numpy as np
import os.path
import copy

# 椒盐噪声
def SaltAndPepper(src, percetage):
    SP_NoiseImg = src.copy()
    SP_NoiseNum = int(percetage * src.shape[0] * src.shape[1])
    for i in range(SP_NoiseNum):
        randR = np.random.randint(0, src.shape[0] - 1)
        randG = np.random.randint(0, src.shape[1] - 1)
        randB = np.random.randint(0, 3)
        if np.random.randint(0, 1) == 0:
            SP_NoiseImg[randR, randG, randB] = 0
        else:
            SP_NoiseImg[randR, randG, randB] = 255
    return SP_NoiseImg

# 高斯噪声
def addGaussianNoise(image, percetage):
    G_Noiseimg = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    G_NoiseNum = int(percetage * image.shape[0] * image.shape[1])
    for i in range(G_NoiseNum):
        temp_x = np.random.randint(0, h)
        temp_y = np.random.randint(0, w)
        G_Noiseimg[temp_x][temp_y][np.random.randint(3)] = np.random.randn(1)[0]
    return G_Noiseimg

# 昏暗
def darker(image, percetage=0.9):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    # get darker
    for xi in range(0, w):
        for xj in range(0, h):
            image_copy[xj, xi, 0] = int(image[xj, xi, 0] * percetage)
            image_copy[xj, xi, 1] = int(image[xj, xi, 1] * percetage)
            image_copy[xj, xi, 2] = int(image[xj, xi, 2] * percetage)
    return image_copy

# 亮度
def brighter(image, percetage=1.5):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    # get brighter
    for xi in range(0, w):
        for xj in range(0, h):
            image_copy[xj, xi, 0] = np.clip(int(image[xj, xi, 0] * percetage), a_max=255, a_min=0)
            image_copy[xj, xi, 1] = np.clip(int(image[xj, xi, 1] * percetage), a_max=255, a_min=0)
            image_copy[xj, xi, 2] = np.clip(int(image[xj, xi, 2] * percetage), a_max=255, a_min=0)
    return image_copy

# 旋转
def rotate(image, angle, center=None, scale=1.0):
    (h, w) = image.shape[:2]
    # If no rotation center is specified, the center of the image is set as the rotation center
    if center is None:
        center = (w / 2, h / 2)
    m = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(image, m, (w, h))
    return rotated

# 翻转
def flip(image):
    flipped_image = np.fliplr(image)
    return flipped_image

# 更改为自己图片文件夹路径，扩增后数据集会输入到同个文件夹下
file_dir = r'C:/Users/Quinton/Desktop/sentry/'
for img_name in os.listdir(file_dir):
    img_path = file_dir + img_name
    img = cv2.imread(img_path)
    rotated_90 = rotate(img, 90)
    cv2.imwrite(file_dir + img_name[0:-4] + '_r90.jpg', rotated_90)
    rotated_180 = rotate(img, 180)
    cv2.imwrite(file_dir + img_name[0:-4] + '_r180.jpg', rotated_180)

for img_name in os.listdir(file_dir):
    img_path = file_dir + img_name
    img = cv2.imread(img_path)
    # 镜像
    flipped_img = flip(img)
    cv2.imwrite(file_dir + img_name[0:-4] + '_fli.jpg', flipped_img)

    # 增加噪声
    img_gauss = addGaussianNoise(img, 0.3)
    cv2.imwrite(file_dir + img_name[0:-4] + '_noise.jpg', img_gauss)

    # 变亮、变暗
    img_darker = darker(img)
    cv2.imwrite(file_dir + img_name[0:-4] + '_darker.jpg', img_darker)
    img_brighter = brighter(img)
    cv2.imwrite(file_dir + img_name[0:-4] + '_brighter.jpg', img_brighter)

    # cv2.GaussianBlur(图像，卷积核，标准差）
    blur = cv2.GaussianBlur(img, (7, 7), 1.5)
    cv2.imwrite(file_dir + img_name[0:-4] + '_blur.jpg', blur)


text:
import networkx as nx

from ._approx import KN20Approximator

class ExactRangeQuery:
    def __init__(self, data, eps, metric):
        self.data = data
        self.eps = eps
        self.metric = metric

    def query(self, q_idx):
        q = self.data[q_idx]
        return [i for i, p in enumerate(self.data)
                if i != q_idx and self.metric(p, q) <= self.eps]

class ApproximateRangeQuery:
    def __init__(self, data, eps, metric, spanner_eps):
        approximator = KN20Approximator(metric)
        self.approx_metric = approximator.approximate(data, spanner_eps)

        lengths = nx.all_pairs_dijkstra_path_length(
            self.approx_metric.spanner,
            cutoff=eps,
        )

        self.neighborhoods = [None] * len(data)
        for neighbors in lengths:
            src, targets = neighbors
            self.neighborhoods[src] = [k for k, v in targets.items() if v != 0]

    def query(self, q_idx):
        return self.neighborhoods[q_idx]

    @property
    def spanner(self):
        return self.approx_metric.spanner

max_stars_repo_name: compTAG/dbspan

max_issues_repo_name: compTAG/dbspan

max_forks_repo_name: compTAG/dbspan

reason: import networkx


text:
# structure_ex1.py
import numpy as np
from civpy.structures import *

section = CrossSection(
    name='dummy',
    area=32.9,
    inertia_x=236,
    inertia_y=716,
    inertia_j=15.1
)

material = Material(
    name='dummy',
    elasticity=29000,
    rigidity=11500
)

group = ElementGroup(
    name='dummy',
    section=section,
    material=material
)

nodes = [
    Node('1', 0, 0, 0),
    Node('2', -240, 0, 0).fixed(),
    Node('3', 0, -240, 0).fixed(),
    Node('4', 0, 0, -240).fixed(),
]

elements =[
    Element('1', '2', '1', group),
    Element('2', '3', '1', group, roll=np.deg2rad(-90)),
    Element('3', '4', '1', group, roll=np.deg2rad(-30)),
]

nloads = [
    NodeLoad('1', mx=-1800, mz=1800)
]

eloads = [
    ElementLoad('1', fy=-0.25)
]

lc = LoadCase('1', nloads, eloads)

struct = Structure(
    name='dummy',
    nodes=nodes,
    elements=elements
)

struct.plot_3d()
result = struct.linear_analysis(lc)

result['glob']
#   load_case node   force_x    force_y   force_z     moment_x   moment_y  \
# 0         1    1  0.000000   0.000000  0.000000 -1800.000000   0.000000
# 1         1    2  5.375736  44.106293 -0.742724     2.172151  58.987351
# 2         1    3 -4.624913  11.117379 -6.460651  -515.545730  -0.764719
# 3         1    4 -0.750823   4.776328  7.203376  -383.501559 -60.166419

#       moment_z    defl_x    defl_y    defl_z     rot_x     rot_y     rot_z
# 0  1800.000000 -0.001352 -0.002797 -0.001812 -0.003002  0.001057  0.006499
# 1  2330.519663  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
# 2   369.671654  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000
# 3    -4.701994  0.000000  0.000000  0.000000  0.000000  0.000000  0.000000

max_stars_repo_name: mpewsey/civpy

max_issues_repo_name: mpewsey/civpy

max_forks_repo_name: mpewsey/civpy


text:
import os
import numpy as np
import glob
import nibabel as nib

from myconfig import *

def rotate_1(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.rot90(img[j,:,:,i],1)
        for j in range(c2):
            lab_rot[j,:,:,i] = np.rot90(lab[j,:,:,i],1)
    return img_rot, lab_rot

def rotate_2(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.rot90(img[j,:,:,i],2)
        for j in range(c2):
            lab_rot[j,:,:,i] = np.rot90(lab[j,:,:,i],2)
    return img_rot, lab_rot

def rotate_3(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.rot90(img[j,:,:,i],3)
        for j in range(c2):
            lab_rot[j,:,:,i] = np.rot90(lab[j,:,:,i],3)
    return img_rot, lab_rot

def rotate_4(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.fliplr(img[j,:,:,i])
        for j in range(c2):
            lab_rot[j,:,:,i] = np.fliplr(lab[j,:,:,i])
    return img_rot, lab_rot

def rotate_5(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.flipud(img[j,:,:,i])
        for j in range(c2):
            lab_rot[j,:,:,i] = np.flipud(lab[j,:,:,i])
    return img_rot, lab_rot

def rotate_6(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.flipud(np.rot90(img[j,:,:,i],1))
        for j in range(c2):
            lab_rot[j,:,:,i] = np.flipud(np.rot90(lab[j,:,:,i],1))
    return img_rot, lab_rot

def rotate_7(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.flipud(np.rot90(img[j,:,:,i],3))
        for j in range(c2):
            lab_rot[j,:,:,i] = np.flipud(np.rot90(lab[j,:,:,i],3))
    return img_rot, lab_rot

def rotate_8(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.flipud(np.rot90(img[j,:,:,i],2))
        for j in range(c2):
            lab_rot[j,:,:,i] = np.flipud(np.rot90(lab[j,:,:,i],2))
    return img_rot, lab_rot

def rotate_9(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.fliplr(np.rot90(img[j,:,:,i],1))
        for j in range(c2):
            lab_rot[j,:,:,i] = np.fliplr(np.rot90(lab[j,:,:,i],1))
    return img_rot, lab_rot

def rotate_10(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.fliplr(np.rot90(img[j,:,:,i],3))
        for j in range(c2):
            lab_rot[j,:,:,i] = np.fliplr(np.rot90(lab[j,:,:,i],3))
    return img_rot, lab_rot

def rotate_11(img,lab):
    # img,lab must be [channels,height,width,depth]
    c1,h,w,d = img.shape
    c2 = lab.shape[0]
    img_rot = np.zeros((c1,h,w,d),dtype=np.float32)
    lab_rot = np.zeros((c2,h,w,d),dtype=np.uint8)
    for i in range(d):
        for j in range(c1):
            img_rot[j,:,:,i] = np.fliplr(np.rot90(img[j,:,:,i],2))
        for j in range(c2):
            lab_rot[j,:,:,i] = np.fliplr(np.rot90(lab[j,:,:,i],2))
    return img_rot, lab_rot

HEIGHT = 96
WIDTH = 96
DEPTH = 48
data_root = DATA_ROOT
sub = 'train'
names = glob.glob(data_root + '/' + sub + '/A-*')
for name in names:
    # [channels, depth, height, width]

    image = np.zeros((3, HEIGHT, WIDTH, DEPTH),dtype=np.float32)
    label = np.zeros((2, HEIGHT, WIDTH, DEPTH),dtype=np.uint8)
    image_prefixes = ['InputCT_ROI', 'InputPET_ROI', 'InputPET_SUV_ROI']
    label_prefixes = ['GTV_Primary_ROI_CT{}'.format(GT_POSTFIX), 'GTV_Primary_ROI_PET{}'.format(GT_POSTFIX)]
    for i, prefix in enumerate(image_prefixes):
        f = nib.load('{}/{}.nii.gz'.format(name, prefix))
        image[i, ...] = f.get_data().astype(np.float32)
    for i, prefix in enumerate(label_prefixes):
        f = nib.load('{}/{}.nii.gz'.format(name, prefix))
        label[i, ...] = f.get_data().astype(np.uint8)

    c1 = 3
    c2 = 2
    num = 7
    image_raw_rot = np.zeros((num,c1,HEIGHT, WIDTH, DEPTH),dtype=np.float32)
    label_raw_rot = np.zeros((num,c2,HEIGHT, WIDTH, DEPTH),dtype=np.uint8)
    image_raw_rot[0,...], label_raw_rot[0,...] = rotate_1(image, label)
    image_raw_rot[1,...], label_raw_rot[1,...] = rotate_2(image, label)
    image_raw_rot[2,...], label_raw_rot[2,...] = rotate_3(image, label)
    image_raw_rot[3,...], label_raw_rot[3,...] = rotate_4(image, label)
    image_raw_rot[4,...], label_raw_rot[4,...] = rotate_5(image, label)
    image_raw_rot[5,...], label_raw_rot[5,...] = rotate_6(image, label)
    image_raw_rot[6,...], label_raw_rot[6,...] = rotate_7(image, label)
    # image_raw_rot[7,...], label_raw_rot[7,...] = rotate_8(image, label) # equal to rotate_4
    # image_raw_rot[8,...], label_raw_rot[8,...] = rotate_9(image, label) # equal to rotate_7
    # image_raw_rot[9,...], label_raw_rot[9,...] = rotate_10(image, label) # equal to rotate_6
    # image_raw_rot[10,...], label_raw_rot[10,...] = rotate_11(image, label) # equal to rotate_5

    affine = np.eye(4)
    for i in range(num):
        save_dir = name+'_rot'+str(i)
        os.makedirs(save_dir)
        for j, prefix in enumerate(image_prefixes):
            img = nib.Nifti1Image(image_raw_rot[i,j,:,:,:], affine)
            nib.save(img, '{}/{}.nii.gz'.format(save_dir, prefix))
        for j, prefix in enumerate(label_prefixes):
            img = nib.Nifti1Image(label_raw_rot[i,j,:,:,:], affine)
            nib.save(img, '{}/{}.nii.gz'.format(save_dir, prefix))

    print(name)
    # break


text:
#!/usr/bin/env python3

# Retipy - Retinal Image Processing on Python
# Copyright (C) 2017  Alejandro Valdes
#
# 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
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
script to estimate the linear tortuosity of a set of retinal images, it will output the values
to a file in the output folder defined in the configuration. The output will only have the
estimated value and it is sorted by image file name.
"""

import argparse
import glob
# import numpy as np
import os
import h5py
import shutil
import pandas as pd
# import scipy.stats as stats

from retipy import configuration, retina, tortuosity_measures

if os.path.exists('../../Results/M2/artery_vein/artery_binary_skeleton/.ipynb_checkpoints'):
    shutil.rmtree('../../Results/M2/artery_vein/artery_binary_skeleton/.ipynb_checkpoints')
if os.path.exists('../../Results/M2/binary_vessel/binary_skeleton/.ipynb_checkpoints'):
    shutil.rmtree('../../Results/M2/binary_vessel/binary_skeleton/.ipynb_checkpoints')
if os.path.exists('../../Results/M2/artery_vein/vein_binary_skeleton/.ipynb_checkpoints'):
    shutil.rmtree('../../Results/M2/artery_vein/vein_binary_skeleton/.ipynb_checkpoints')
if not os.path.exists('../../Results/M3/Macular_centred/Width/'):
    os.makedirs('../../Results/M3/Macular_centred/Width/')

#if os.path.exists('./DDR/av_seg/raw/.ipynb_checkpoints'):
#    shutil.rmtree('./DDR/av_seg/raw/.ipynb_checkpoints')

parser = argparse.ArgumentParser()

parser.add_argument(
    "-c",
    "--configuration",
    help="the configuration file location",
    default="resources/retipy.config")
args = parser.parse_args()

CONFIG = configuration.Configuration(args.configuration)
binary_FD_binary,binary_VD_binary,binary_Average_width,binary_t2_list,binary_t4_list,binary_t5_list = [],[],[],[],[],[]
artery_FD_binary,artery_VD_binary,artery_Average_width,artery_t2_list,artery_t4_list,artery_t5_list = [],[],[],[],[],[]
vein_FD_binary,vein_VD_binary,vein_Average_width,vein_t2_list,vein_t4_list,vein_t5_list = [],[],[],[],[],[]
name_list = []

Artery_PATH = '../../Results/M2/artery_vein/macular_centred_artery_skeleton'
Vein_PATH = '../../Results/M2/artery_vein/macular_centred_vein_skeleton'
Binary_PATH = '../../Results/M2/binary_vessel/macular_centred_binary_skeleton'

for filename in sorted(glob.glob(os.path.join(Binary_PATH, '*.png'))):
    segmentedImage = retina.Retina(None, filename, store_path='../../Results/M2/binary_vessel/macular_centred_binary_process')
    #segmentedImage.threshold_image()
    #segmentedImage.reshape_square()
    #window_sizes = segmentedImage.get_window_sizes()
    window_sizes = [912]
    window = retina.Window(
        segmentedImage, window_sizes[-1], min_pixels=CONFIG.pixels_per_window)
    FD_binary,VD_binary,Average_width, t2, t4, td = tortuosity_measures.evaluate_window(window, CONFIG.pixels_per_window, CONFIG.sampling_size, CONFIG.r_2_threshold,store_path='../../Results/M2/binary_vessel/macular_centred_binary_process/')
    #print(window.tags)
    binary_t2_list.append(t2)
    binary_t4_list.append(t4)
    binary_t5_list.append(td)
    binary_FD_binary.append(FD_binary)
    binary_VD_binary.append(VD_binary)
    binary_Average_width.append(Average_width)
    name_list.append(filename.split('/')[-1])

for filename in sorted(glob.glob(os.path.join(Artery_PATH, '*.png'))):

    segmentedImage = retina.Retina(None, filename,store_path='../../Results/M2/artery_vein/macular_centred_artery_process')
    window_sizes = [912]
    window = retina.Window(
        segmentedImage, window_sizes[-1], min_pixels=CONFIG.pixels_per_window)
    FD_binary,VD_binary,Average_width, t2, t4, td = tortuosity_measures.evaluate_window(window, CONFIG.pixels_per_window, CONFIG.sampling_size, CONFIG.r_2_threshold,store_path='../../Results/M2/artery_vein/macular_centred_artery_process/')
    #print(window.tags)
    artery_t2_list.append(t2)
    artery_t4_list.append(t4)
    artery_t5_list.append(td)
    artery_FD_binary.append(FD_binary)
    artery_VD_binary.append(VD_binary)
    artery_Average_width.append(Average_width)

for filename in sorted(glob.glob(os.path.join(Vein_PATH, '*.png'))):

    segmentedImage = retina.Retina(None, filename,store_path='../../Results/M2/artery_vein/macular_centred_vein_process')
    window_sizes = [912]
    window = retina.Window(
        segmentedImage, window_sizes[-1], min_pixels=CONFIG.pixels_per_window)
    FD_binary,VD_binary,Average_width, t2, t4, td = tortuosity_measures.evaluate_window(window, CONFIG.pixels_per_window, CONFIG.sampling_size, CONFIG.r_2_threshold,store_path='../../Results/M2/artery_vein/macular_centred_vein_process/')
    #print(window.tags)
    vein_t2_list.append(t2)
    vein_t4_list.append(t4)
    vein_t5_list.append(td)
    vein_FD_binary.append(FD_binary)
    vein_VD_binary.append(VD_binary)
    vein_Average_width.append(Average_width)

Disc_file = pd.read_csv('../../Results/M3/Macular_centred/Disc_cup_results.csv')

Data4stage2 = pd.DataFrame({'Fractal_dimension':binary_FD_binary, 'Vessel_density':binary_VD_binary, 'Average_width':binary_Average_width,'Distance_tortuosity':binary_t2_list, 'Squared_curvature_tortuosity':binary_t4_list, 'Tortuosity_density':binary_t5_list, 'Artery_Fractal_dimension':artery_FD_binary, 'Artery_Vessel_density':artery_VD_binary, 'Artery_Average_width':artery_Average_width,'Artery_Distance_tortuosity':artery_t2_list, 'Artery_Squared_curvature_tortuosity':artery_t4_list, 'Artery_Tortuosity_density':artery_t5_list, 'Vein_Fractal_dimension':vein_FD_binary, 'Vein_Vessel_density':vein_VD_binary, 'Vein_Average_width':vein_Average_width,'Vein_Distance_tortuosity':vein_t2_list, 'Vein_Squared_curvature_tortuosity':vein_t4_list, 'Vein_Tortuosity_density':vein_t5_list})
#Data4stage2 = pd.concat([Disc_file[Disc_file['Name'].isin(name_list)].reset_index().drop(columns=['index', 'Name']), Data4stage2] ,axis=1)
Data4stage2 = pd.concat([Disc_file[Disc_file['Name'].isin(name_list)].reset_index(), Data4stage2] ,axis=1)
Data4stage2.to_csv('../../Results/M3/Macular_centred/Macular_Measurement.csv', index = None, encoding='utf8')

max_stars_repo_name: rmaphoh/AutoMorph

max_issues_repo_name: rmaphoh/AutoMorph

max_forks_repo_name: rmaphoh/AutoMorph

reason: import numpy,import scipy


text:
import numpy as np
import matplotlib.pyplot as plt

from moviepy.editor import *
from moviepy.video.io.bindings import mplfig_to_npimage

if __name__ == '__main__':

    #Boolean to set whether to include audio or not in final product
    include_audio = False

    # Read in data generated from previous video analysis
    scenes = np.loadtxt('MV_scenelist.csv', delimiter=',')

    # Break these columns apart to make it a little easier to index later
    scene_frames = scenes[:, 0]
    scene_times = scenes[:, 1]
    scene_durs = scenes[:, 2]

    # Calculate the time that passes each frame
    msec_per_frame = scene_times[-1] / scene_frames[-1]

    # Initialize a list to keep track of edits/sec
    rolling_average = []

    # Size of the rolling window to average over
    window_sec = 5.0

    # Keep track of number of scenes to have passed
    scene_count = []

    for i in range(int(np.max(scene_frames)) + 1):

        # First frame, starts with 0 edits per second
        if i == 0:
            rolling_average.append(0)
            scene_count.append(1)
            continue

        # Find current time in msec
        current_time = i * msec_per_frame

        # Find all scenes that have happened prior to current frame
        in_window_scenes = scene_times[np.where(current_time >= scene_times)]

        # Keep track of total number of scenes to have passed
        scene_count.append(len(in_window_scenes))

        # Then filter that down to scenes that have happened within the rolling
        # window prior to current frame
        in_window_scenes = len(in_window_scenes[np.where(
            current_time - window_sec * 1000. <= in_window_scenes)])

        # Find the rate from number of scenes
        scenes_per_sec = in_window_scenes / window_sec

        # Add this frame's rate to the list
        rolling_average.append(scenes_per_sec)

    # Calculate the average rate at which edits are made
    avg_rate = scene_count[-1] / (scene_times[-1] / 1000.0)

    # duration = 30
    duration = scene_times[-1] / 1000.0

    # Make the first figure that will keep track of the rate of transitions
    fig1, ax = plt.subplots(1, figsize=(4, 4), facecolor='white')
    ax.set_title("Rate of Scene Transitions \n (%0d sec Rolling Average)" % window_sec)
    ax.set_ylim(0, max(rolling_average))
    ax.set_xlim(0, duration)
    ax.set_xlabel('Time (sec)')
    ax.set_ylabel('Detected Rate of Transitions (changes/sec)')
    line, = ax.plot(0, 0, 'k-')
    line2, = ax.plot([0, duration], [avg_rate, avg_rate], 'b-')
    plt.tight_layout()

    # Initialize lists to keep track of things
    times = []
    rates = []

    def make_frame1(t):
        """Function to make a graph of the rate of scene transitions."""

        times.append(t)

        # Find all scenes that have happened prior to current frame
        in_window_scenes = scene_times[np.where(t * 1000.0 >= scene_times)]

        # Then filter that down to scenes that have happened within the rolling
        # window prior to current frame
        in_window_scenes = len(in_window_scenes[np.where(
            t * 1000.0 - window_sec * 1000.0 <= in_window_scenes)])

        # Find the rate from number of scenes
        scenes_per_sec = in_window_scenes / window_sec

        # Add the rate to the list
        rates.append(scenes_per_sec)

        # Update the graph
        line.set_xdata(times)
        line.set_ydata(rates)

        return mplfig_to_npimage(fig1)

    # Use our function to animate a video and save it to file
    animation1 = VideoClip(make_frame1, duration=duration).resize(height=540)
    # animation1.write_videofile('rate_animation.mp4', fps=23.976)
    plt.close()

    # Animate plot of total scene transitions detected
    fig2, ax = plt.subplots(1, figsize=(4, 4), facecolor='white')
    ax.set_title("Number of Scene Transitions")
    ax.set_ylim(0, max(scene_count))
    ax.set_xlim(0, duration)
    ax.set_xlabel('Time (sec)')
    ax.set_ylabel('Total Number of Detected Scenes')
    line, = ax.plot(0, 1, 'k-')
    line2, = ax.plot([0, duration], [1, max(scene_count)], 'b-')
    plt.tight_layout()

    # Initialize lists
    times = []
    scenes = []

    def make_frame2(t):
        """Function to graph the total number of scene transitions over time."""

        # Keep track of the time
        times.append(t)

        # Find all scenes that have happened prior to current frame
        in_window_scenes = scene_times[np.where(t * 1000.0 >= scene_times)]

        # Keep track of total number of scenes to have passed
        scenes.append(len(in_window_scenes) + 1)

        # Update the graph
        line.set_xdata(times)
        line.set_ydata(scenes)

        return mplfig_to_npimage(fig2)

    # Animate the graph and save it to file
    animation2 = VideoClip(make_frame2, duration=duration).resize(height=540)
    # animation2.write_videofile('total_animation.mp4', fps=23.976)

    # Reload saved videos of graphs as they cannot be composited together
    # until they are each rendered and saved independently
    animation1 = VideoFileClip('rate_animation.mp4')
    animation2 = VideoFileClip('total_animation.mp4')

    # Stack the two graphs on top of each other
    animation_array = clips_array([[animation1], [animation2]])
    # animation_array.write_videofile('stacked_animation.mp4', fps=23.976)

    # Load the main video
    mv_video = (VideoFileClip('BTS_2017_DNA_Annotated_33_10.mp4').
                resize(width=1920 - animation_array.w))

    #Get rid of audio if set
    if not include_audio:
        mv_video = mv_video.set_audio(None)

    # Stick the videos together and save the result
    final_array = clips_array([[mv_video, animation_array]])
    # final_array.write_videofile('final_composition.mp4', fps=23.976)

    # Make a final results screen

    # Calculate final stats
    sec_per_scene = duration / float(max(scene_count))

    result_string1 = "Results:                              \n"
    result_string2 = "%03d scenes in %03.1f seconds           " % (max(scene_count), duration)
    result_string3 = "Average of %0.2f transitions per second" % avg_rate
    result_string4 = "Average of %0.2f seconds per scene     " % sec_per_scene

    result_text = [result_string1, result_string2,
                   result_string3, result_string4]

    final_result_text = "\n".join(result_text)

    result_screen_text = TextClip(final_result_text, fontsize=72,
                                  font="FreeMono-Bold", color='white',
                                  size=(final_array.w, final_array.h)
                                  ).set_duration(7.5).set_pos('center')

    video_result = concatenate_videoclips([final_array, result_screen_text])

    # Tack on the longest scene at the end

    # Find the longest scene
    scene_idx = np.argmax(scene_durs)

    # Figure out the timing of the longest scene
    scene_start = scene_times[scene_idx - 1] / 1000.0
    scene_end = (scene_times[scene_idx] - msec_per_frame) / 1000.0
    scene_duration = scene_end - scene_start

    # Make the text for the top of the screen
    scene_text = (TextClip("Longest Scene", fontsize=144, font="FreeMono-Bold",
                           stroke_color='black', stroke_width=3, color='white').
                           set_duration(scene_duration).set_opacity(0.6))
    scene_text = scene_text.set_pos("center").set_pos('top')

    # Make the text for the bottom of the screen
    dur_text = (TextClip("%0.3f seconds long" % scene_duration, fontsize=144,
                         font="FreeMono-Bold", stroke_color='black',
                         color='white', stroke_width=3).
                         set_duration(scene_duration).set_opacity(0.6))
    dur_text = dur_text.set_pos('center').set_pos('bottom')

    # Load the longest scene from the previously annotated video
    video_file = '/home/walter/Videos/Music Videos/BTS_2017_DNA.mkv'
    longest_scene = (VideoFileClip(video_file).
                     subclip(scene_start, scene_end))

    #Get rid of audio if set
    if not include_audio:
        longest_scene = longest_scene.set_audio(None)

    # Combine the text and the longest scene together
    final_scene = CompositeVideoClip([longest_scene, scene_text, dur_text])

    # Add the longest scene onto the end of the annotated video
    added_scene = concatenate_videoclips([video_result, final_scene])
    added_scene.write_videofile('added_scene.mp4', fps=23.976, preset='medium')


text:
import os
import pickle
from copy import deepcopy
import numpy as np

from .import features as elf_feats
from .import learning as elf_learn
from .import lifted_multicut as elf_lmc
from .import multicut as elf_mc
from .import watershed as elf_ws

FEATURE_NAMES = {'raw-edge-features', 'raw-region-features', 'boundary-edge-features'}
DEFAULT_WS_KWARGS = {'threshold': 0.25, 'sigma_seeds': 2., 'sigma_weights': 2.,
                     'min_size': 100, 'alpha': 0.9, 'pixel_pitch': None,
                     'apply_nonmax_suppression': False}
DEFAULT_RF_KWARGS = {'ignore_label': None, 'n_estimators': 200, 'max_depth': 10}

def _compute_watershed(boundaries, use_2dws, mask, ws_kwargs, n_threads):
    if use_2dws:
        bd_ = boundaries if isinstance(boundaries, np.ndarray) else boundaries[0]
        return elf_ws.stacked_watershed(bd_, mask=mask, n_threads=n_threads, **ws_kwargs)[0]
    else:
        return elf_ws.distance_transform_watershed(boundaries, mask=mask, **ws_kwargs)[0]

def _compute_xyz_boundary_features(rag, boundaries, watershed, use_2dws, n_threads):
    bd_xy, bd_z = boundaries
    z_edges = elf_feats.compute_z_edge_mask(rag, watershed)
    feats = elf_feats.compute_boundary_features_with_filters(rag, bd_xy, use_2dws, n_threads=n_threads)
    feats_z = elf_feats.compute_boundary_features_with_filters(rag, bd_z, use_2dws, n_threads=n_threads)
    feats[z_edges] = feats_z[z_edges]
    return feats

def _compute_features(raw, boundaries, watershed, feature_names, use_2dws, n_threads):
    if (len(FEATURE_NAMES - set(feature_names)) > 0) or (len(feature_names) == 0):
        raise ValueError("Invalid feature set")
    if isinstance(boundaries, np.ndarray):
        if raw.shape != boundaries.shape:
            raise ValueError("Shapes %s and %s do not match" % (str(raw.shape), str(boundaries.shape)))
    if raw.shape != watershed.shape:
        raise ValueError("Shapes %s and %s do not match" % (str(raw.shape), str(watershed.shape)))

    rag = elf_feats.compute_rag(watershed, n_threads=n_threads)

    features = []
    if 'raw-edge-features' in feature_names:
        feats = elf_feats.compute_boundary_features_with_filters(rag, raw, use_2dws,
                                                                 n_threads=n_threads)
        features.append(feats)
    if 'boundary-edge-features' in feature_names:
        if isinstance(boundaries, np.ndarray):
            feats = elf_feats.compute_boundary_features_with_filters(rag, boundaries, use_2dws,
                                                                     n_threads=n_threads)
        else:
            if not use_2dws:
                raise ValueError("Separate boundary maps for xy and z edges are only available if use_2dws is True")
            if len(boundaries) != 2 or\
               (not isinstance(boundaries[0], np.ndarray)) or\
               (not isinstance(boundaries[1], np.ndarray)):
                raise ValueError("Invalid boundary maps")
            feats = _compute_xyz_boundary_features(rag, boundaries, watershed, use_2dws, n_threads)
        features.append(feats)
    if 'raw-region-features' in feature_names:
        feats = elf_feats.compute_region_features(rag.uvIds(), raw, watershed.astype('uint32'),
                                                  n_threads=n_threads)
        features.append(feats)

    # for now, we always append the length as one other feature
    # eventually, it would be nice to add topolgy features, cf.
    # https://github.com/ilastik/nature_methods_multicut_pipeline/blob/master/software/multicut_src/DataSet.py#L954
    # https://github.com/DerThorsten/nifty/blob/master/src/python/lib/graph/rag/accumulate.cxx#L361
    edge_len = elf_feats.compute_boundary_mean_and_length(rag, raw, n_threads=n_threads)[:, 1]
    features.append(edge_len[:, None])

    features = np.concatenate(features, axis=1)
    assert len(features) == rag.numberOfEdges
    return rag, features

def _compute_features_and_labels(raw, boundaries, watershed, labels, mask,
                                 use_2dws, feature_names, ignore_label, n_threads):
    rag, features = _compute_features(raw, boundaries, watershed, feature_names, use_2dws, n_threads)
    edge_mask = np.ones(len(features), dtype='bool')

    if ignore_label is None:
        edge_labels = elf_learn.compute_edge_labels(rag, labels, n_threads=n_threads)
    else:
        edge_labels, edge_mask1 = elf_learn.compute_edge_labels(rag, labels, ignore_label, n_threads)
        edge_mask = np.logical_and(edge_mask, edge_mask1)

    if mask is not None:
        edge_mask2 = (rag.uvIds() != 0).any(axis=1)
        edge_mask = np.logical_and(edge_mask, edge_mask2)

    features, edge_labels = features[edge_mask], edge_labels[edge_mask]
    assert len(features) == len(edge_labels)

    if use_2dws:
        z_edges = elf_feats.compute_z_edge_mask(rag, watershed)
        z_edges = z_edges[edge_mask]
        assert len(z_edges) == len(features)
    else:
        z_edges = None

    return features, edge_labels, z_edges

def _get_solver(solver):
    if isinstance(solver, str):
        solver = elf_mc.get_multicut_solver(solver)
    elif not callable(solver):
        raise ValueError("Invalid multicut solver")
    return solver

def _mask_edges(rag, edge_costs):
    uv_ids = rag.uvIds()
    ignore_edges = (uv_ids == 0).any(axis=1)
    edge_costs[ignore_edges] = - np.abs(edge_costs).max()
    return edge_costs

#
# training workflows
#

def edge_training(raw, boundaries, labels, use_2dws, watershed=None, mask=None,
                  feature_names=FEATURE_NAMES, ws_kwargs=DEFAULT_WS_KWARGS,
                  learning_kwargs=DEFAULT_RF_KWARGS, n_threads=None):
    """ Train random forest classifier for edges.

    Arguments:
        raw [np.ndarray or list[np.ndarray]] - raw data or list of raw data arrays
        boundaries [np.ndarray or list[np.ndarray]] - boundary maps or list of boundary maps
        labels [np.ndarray or list[np.ndarray]] - groundtruth segmentation
        use_2dws [bool] - whether to run watersheds in 2d and stack them
            or if the watersheds passed are 2d and stacked
        watershed [np.ndarray or list[np.ndarray]] - watershed segmentation.
            will be computed if None is passed (default: None)
        mask [np.ndarray or list[np.ndarray]] - mask (default: None)
        feature_names [list] - names of features that will be computed (default: all available)
        ws_kwargs [dict] - keyword arguments for watershed function
        learning_kwargs [dict] - keyword arguments for random forest
        n_threads [int] - number of threads (default: None)
    """

    # we store the ignore label(s) in the random forest kwargs,
    # but they need to be removed before this is passed to the rf implementation
    rf_kwargs = deepcopy(learning_kwargs)
    ignore_label = rf_kwargs.pop('ignore_label', None)

    # if the raw data is a numpy array, we assume that we have a single training set
    if isinstance(raw, np.ndarray):
        if (not isinstance(boundaries, np.ndarray)) or (not isinstance(labels, np.ndarray)):
            raise ValueError("Expect raw data, boundaries and labels to be either all numpy arrays or lists")
        if mask is not None and not isinstance(mask, np.ndarray):
            raise ValueError("Invalid mask")

        # compute the watersheds segmentation if it was not passed
        if watershed is None:
            watershed = _compute_watershed(boundaries, use_2dws, mask, ws_kwargs, n_threads)
        features, edge_labels, z_edges = _compute_features_and_labels(raw, boundaries, watershed,
                                                                      labels, mask,
                                                                      use_2dws, feature_names,
                                                                      ignore_label, n_threads)

    # otherwise, we assume to get listlike data for raw data, boundaries and labels,
    # corresponding to multiple training data-sets
    else:
        if not (len(raw) == len(boundaries) == len(labels)):
            raise ValueError("Expect same number of raw data, boundary and label arrays")
        if watershed is not None and len(watershed) != len(raw):
            raise ValueError("Expect same number of watershed arrays as raw data")
        if mask is not None and not len(mask) == len(raw):
            raise ValueError("Expect same number of mask arrays as raw data")

        features = []
        edge_labels = []
        z_edges = []
        # compute features and labels for all training data-sets
        for train_id, (this_raw, this_boundaries, this_labels) in enumerate(zip(raw, boundaries, labels)):
            # compute the watersheds segmentation if it was not passed
            if watershed is None:
                this_mask = None if mask is None else mask[train_id]
                this_watershed = _compute_watershed(this_boundaries, use_2dws, this_mask, ws_kwargs, n_threads)
            else:
                this_watershed = watershed[train_id]
                this_mask = None

            this_features, this_edge_labels, this_z_edges = _compute_features_and_labels(this_raw, this_boundaries,
                                                                                this_watershed, this_labels,
                                                                                this_mask, use_2dws,
                                                                                feature_names,
                                                                                ignore_label, n_threads)
            features.append(this_features)
            edge_labels.append(this_edge_labels)
            # we only get z-edges if we have 2d watersheds
            if use_2dws:
                assert this_z_edges is not None
                z_edges.append(this_z_edges)

        features = np.concatenate(features, axis=0)
        edge_labels = np.concatenate(edge_labels, axis=0)
        if use_2dws:
            z_edges = np.concatenate(z_edges, axis=0)

    assert len(features) == len(edge_labels), "%i, %i" % (len(features),
                                                          len(edge_labels))

    # train the random forest (2 separate random forests for in-plane and between plane edges for stacked 2d watersheds)
    if use_2dws:
        assert len(features) == len(z_edges)
        rf = elf_learn.learn_random_forests_for_xyz_edges(features, edge_labels, z_edges, n_threads=n_threads,
                                                          **rf_kwargs)
    else:
        rf = elf_learn.learn_edge_random_forest(features, edge_labels, n_threads=n_threads, **rf_kwargs)
    return rf

def _load_rf(rf, use_2dws):

    def _load(rfp):
        if isinstance(rfp, str):
            if not os.path.exists(rfp):
                raise ValueError("Invalid random foerst path %s" % rfp)
            with open(rfp, 'rb') as f:
                rfp = pickle.load(f)
        return rfp

    if use_2dws:
        if len(rf) != 2:
            raise ValueError("Expect two RFs")
        rf_xy, rf_z = rf
        rf_xy, rf_z = _load(rf_xy), _load(rf_z)
        return rf_xy, rf_z
    else:
        rf = _load(rf)
        return rf

# TODO
# - multicut_workflow_from_config(config_file)
#
# multicut workflows
#

def multicut_segmentation(raw, boundaries, rf, use_2dws,
                          multicut_solver='kernighan-lin', watershed=None, mask=None,
                          feature_names=FEATURE_NAMES, weighting_scheme=None,
                          ws_kwargs=DEFAULT_WS_KWARGS, solver_kwargs={},
                          beta=0.5, n_threads=None, return_intermediates=False,
                          cost_callback=None):
    """ Instance segmentation with multicut with edge costs
    derived from random forest predictions.

    Arguments:
        raw [np.ndarray] - raw data
        boundaries [np.ndarray] - boundary maps
        rf [RandomForestClassifier or str] - edge classifier, path to pickled rf is also possible
        use_2dws [bool] - whether to run watersheds in 2d and stack them
            or if the watersheds passed are 2d and stacked
        multicut_solver [str or callable] - name of multicut solver in elf.segmentation.multicut
        or custom solver function (default: 'kernighan-lin')
        watershed [np.ndarray] - watershed segmentation.
            will be computed if None is passed (default: None)
        mask [np.ndarray] - mask to ignore in segmentation (default: None)
        feature_names [list] - names of features that will be computed (default: all available)
        weighting_scheme [str] - strategy to weight multicut edge costs by size (default: None)
        ws_kwargs [dict] - keyword arguments for watershed function
        solver_kwargs [dict] - keyword arguments for multicut solver (default: {})
        beta [float] - boundary bias (default: 0.5)
        n_threads [int] - number of threads (default: None)
        return_intermediates [bool] - whether to also return intermediate results (default: False)
        cost_callback [callable] - callback to modify the costs.
            It will be passed the costs, rag and watershed (default: None)
    """

    # get the multicut solver
    solver = _get_solver(multicut_solver)

    # load random forest(s) if they are given as path
    if use_2dws:
        rf_xy, rf_z = _load_rf(rf, use_2dws)
    else:
        rf = _load_rf(rf, use_2dws)

    # compute watersheds if none were given
    if watershed is None:
        print("Compute watershed ...")
        watershed = _compute_watershed(boundaries, use_2dws, mask, ws_kwargs, n_threads)

    # compute rag and features
    print("Compute rag and features ...")
    rag, features = _compute_features(raw, boundaries, watershed,
                                      feature_names, use_2dws, n_threads)

    # use random forest to compute edge probabilties.
    # if we have stacked 2d watersheds, we expect two random forests,
    # one for the in-plane (xy) and one for the between-plane (z) edges
    if use_2dws:
        print("Predict edges for 2d watershed ...")
        z_edges = elf_feats.compute_z_edge_mask(rag, watershed)
        edge_probs = elf_learn.predict_edge_random_forests_for_xyz_edges(rf_xy, rf_z,
                                                                         features, z_edges,
                                                                         n_threads)
    else:
        print("Predict edges for 3d watershed ...")
        edge_probs = elf_learn.predict_edge_random_forest(rf, features, n_threads)
        z_edges = None

    # derive the edge costs from random forst probabilities
    print("Compute edge sizes ...")
    edge_sizes = elf_feats.compute_boundary_mean_and_length(rag, raw, n_threads)[:, 1]
    print("Compute edge costs ...")
    edge_costs = elf_mc.compute_edge_costs(edge_probs, edge_sizes=edge_sizes, beta=beta,
                                           z_edge_mask=z_edges, weighting_scheme=weighting_scheme)

    # apply the cost call back if given
    # and mask edges connecting to the mask label if mask was given
    if cost_callback is not None:
        edge_costs = cost_callback(edge_costs, rag, watershed)
    if mask is not None:
        edge_costs = _mask_edges(rag, edge_costs)

    # compute multicut and project to pixels
    # we pass the watershed to the solver as well, because it is needed for
    # the blockwise-multicut solver
    print("Run multicut solver ...")
    node_labels = solver(rag, edge_costs, n_threads=n_threads,
                         segmentation=watershed, **solver_kwargs)
    print("Project to volume ...")
    seg = elf_feats.project_node_labels_to_pixels(rag, node_labels, n_threads)

    if return_intermediates:
        return {'watershed': watershed,
                'rag': rag,
                'features': features,
                'edge_costs': edge_costs,
                'node_labels': node_labels,
                'segmentation': seg}
    else:
        return seg

def multicut_workflow(train_raw, train_boundaries, train_labels,
                      raw, boundaries, use_2dws, multicut_solver='kernighan-lin',
                      train_watershed=None, watershed=None, train_mask=None, mask=None,
                      feature_names=FEATURE_NAMES, weighting_scheme=None,
                      ws_kwargs=DEFAULT_WS_KWARGS, learning_kwargs=DEFAULT_RF_KWARGS,
                      solver_kwargs={}, beta=0.5, n_threads=None):
    """ Run workflow for multicut segmentation based on boundary maps with edge weights learned via random forest.

    Based on "Multicut brings automated neurite segmentation closer to human performance":
    https://hci.iwr.uni-heidelberg.de/sites/default/files/publications/files/217205318/beier_17_multicut.pdf

    Arguments:
        train_raw [np.ndarray or list[np.ndarray]] - raw data for training data-sets,
            can be list if there are multiple training sets
        train_boundaries [np.ndarray or list[np.ndarray]] - boundary maps for training data-sets,
            can be list if there are multiple training sets
        train_labels [np.ndarray or list[np.ndarray]] - segmentation ground-truth for training data-sets,
            can be list if there are multiple training sets
        raw [np.ndarray] - raw data of the data-set to be segmented
        boundaries [np.ndarray] - boundary maps of the data-set to be segmented
        use_2dws [bool] - whether to run watersheds in 2d and stack them
            or if the watersheds passed are 2d and stacked
        multicut_solver [str or callable] - name of multicut solver in elf.segmentation.multicut
            or custom solver function (default: 'kernighan-lin')
        train_watershed [np.ndarray] - watershed segmentation for the training data-sets
            can be list if there are multiple training sets,
            will be computed if None is passed (default: None)
        watershed [np.ndarray] - watershed segmentation for the data to be segmented.
            will be computed if None is passed (default: None)
        train_mask [np.ndarray] - mask to ignre in waterheds for training datasets (default: None)
        mask [np.ndarray] - mask to ignore in segmentation (default: None)
        feature_names [list] - names of features that will be computed (default: all available)
        weighting_scheme [str] - strategy to weight multicut edge costs by size (default: None)
        ws_kwargs [dict] - keyword arguments for watershed function
        learning_kwargs [dict] - keyword arguments for the random forest
        solver_kwargs [dict] - keyword arguments for multicut solver (default: {})
        beta [float] - boundary bias (default: 0.5)
        n_threads [int] - number of threads (default: None)
    """
    # train random forest for edge classification based on the training data-sets
    rf = edge_training(train_raw, train_boundaries, train_labels, use_2dws,
                       watershed=train_watershed, mask=train_mask, feature_names=feature_names,
                       ws_kwargs=ws_kwargs, learning_kwargs=learning_kwargs,
                       n_threads=n_threads)

    # run multicut to obtain segmenation for the data to be segmented
    seg = multicut_segmentation(raw, boundaries, rf,
                                use_2dws, multicut_solver,
                                watershed=watershed, mask=mask, feature_names=feature_names,
                                weighting_scheme=weighting_scheme, ws_kwargs=ws_kwargs,
                                solver_kwargs=solver_kwargs, beta=beta, n_threads=n_threads)
    return seg

def _simple_mc_problem(input_, use_2dws, watershed, mask, weighting_scheme, ws_kwargs,
                       beta, offsets, n_threads, cost_callback):
    # check if input is affinities or boundaries and determine the spatial shape
    have_affs = input_.ndim == 4
    spatial_shape = input_.shape[1:] if have_affs else input_.shape
    if have_affs:
        offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]] if offsets is None else offsets
        if input_.shape[0] != len(offsets):
            raise ValueError("Incosistent number of offsets and affinity channels: %i, %i" % (len(offsets),
                                                                                input_.shape[0]))
    if mask is not None:
        if mask.shape != spatial_shape or mask.dtype != np.dtype('bool'):
            raise ValueError("Invalid mask was passed")

    # compute watersheds if they were not passed
    if watershed is None:
        if have_affs:
            bmap = np.max(input_[1:3], axis=0) if use_2dws else np.max(input_[:3], axis=0)
        else:
            bmap = input_
        watershed = _compute_watershed(bmap, use_2dws, mask, ws_kwargs, n_threads)
        assert watershed.shape == spatial_shape
    else:
        if watershed.shape != spatial_shape:
            raise ValueError("Invalid shape of watersheds: got %s, expected %s" % str(watershed.shape,
                                                                                spatial_shape))

    # compute rag and edge costs
    rag = elf_feats.compute_rag(watershed, n_threads=n_threads)
    if have_affs:
        edge_probs = elf_feats.compute_affinity_features(rag, input_, offsets, n_threads=n_threads)
    else:
        edge_probs = elf_feats.compute_boundary_mean_and_length(rag, input_, n_threads=n_threads)

    if use_2dws:
        z_edges = elf_feats.compute_z_edge_mask(rag, watershed)
    else:
        z_edges = None

    edge_probs, edge_sizes = edge_probs[:, 0], edge_probs[:, -1]
    edge_costs = elf_mc.compute_edge_costs(edge_probs, edge_sizes=edge_sizes, beta=beta,
                                           weighting_scheme=weighting_scheme, z_edge_mask=z_edges)

    if cost_callback is not None:
        edge_costs = cost_callback(edge_costs, rag, watershed)

    if mask is not None:
        edge_costs = _mask_edges(rag, edge_costs)
    return watershed, rag, edge_probs, edge_costs

def simple_multicut_workflow(input_, use_2dws, multicut_solver='kernighan-lin', watershed=None, mask=None,
                             weighting_scheme=None, ws_kwargs=DEFAULT_WS_KWARGS,
                             solver_kwargs={}, beta=0.5, offsets=None, n_threads=None,
                             return_intermediates=False, cost_callback=None):
    """ Run simplified multicut segmentation workflow from affinity or boundary maps.

    Adapted from "Multicut brings automated neurite segmentation closer to human performance":
    https://hci.iwr.uni-heidelberg.de/sites/default/files/publications/files/217205318/beier_17_multicut.pdf

    Arguments:
        input_ [np.ndarray] - input data, either boundary or affinity map
        use_2dws [bool] - whether to run watersheds in 2d and stack them
            or if the watersheds passed are 2d and stacked
        multicut_solver [str or callable] - name of multicut solver in elf.segmentation.multicut
            or custom solver function (default: 'kernighan-lin')
        watershed [np.ndarray] - watershed over-segmentation, if None is passed,
            the overr-segmentation witll be computted (default: None)
        mask [np.ndarray] - binary mask to exclude from the segmentation;
            values that are False will be excluded (default: None)
        weighting_scheme [str] - strategy to weight multicut edge costs by size (default: None)
        ws_kwargs [dict] - keyword arguments for the watershed function (default: None)
        solver_kwargs [dict] - keyword arguments for multicut solver (default: {})
        beta [float] - boundary bias (default: 0.5)
        offsets [list] - pixel offsets for affintities, by default assume nearest neighbor offsets (default: None)
        n_threads [int] - number of threads (default: None)
        return_intermediates [bool] - whether to also return intermediate results (default: False)
        cost_callback [callable] - callback to modify the costs.
            It will be passed the costs, rag and watershed (default: None)
    """
    # get the multicut solver
    solver = _get_solver(multicut_solver)

    # create the multicut problem
    watershed, rag, edge_probs, edge_costs = _simple_mc_problem(
        input_, use_2dws, watershed, mask, weighting_scheme, ws_kwargs,
        beta, offsets, n_threads, cost_callback
    )

    # solve the multicut problem and project back to segmentation
    # we pass the watershed to the solver as well, because it is needed for
    # the blockwise-multicut solver
    node_labels = solver(rag, edge_costs, n_threads=n_threads,
                         segmentation=watershed, **solver_kwargs)
    seg = elf_feats.project_node_labels_to_pixels(rag, node_labels, n_threads)

    if return_intermediates:
        return {'watershed': watershed,
                'rag': rag,
                'probs': edge_probs,
                'costs': edge_costs,
                'node_labels': node_labels,
                'segmentation': seg}
    else:
        return seg

# TODO:
# - lifted_multicut_from_probabilities_workflow
# - learning based lifted multicut workflows
# - lifted_multicut_workflow_from_config(config_file)
#
# lifted multicut workflows
#

def _get_lifted_solver(solver):
    if isinstance(solver, str):
        solver = elf_lmc.get_lifted_multicut_solver(solver)
    elif not callable(solver):
        raise ValueError("Invalid lifted multicut solver")
    return solver

def lifted_multicut_from_segmentation_workflow(input_, segmentation, use_2dws,
                                               overlap_threshold, graph_depth,
                                               same_segment_cost, different_segment_cost,
                                               lifted_multicut_solver='kernighan-lin', watershed=None, mask=None,
                                               weighting_scheme=None, ws_kwargs=DEFAULT_WS_KWARGS, mode='all',
                                               solver_kwargs={}, beta=0.5, offsets=None, n_threads=None,
                                               return_intermediates=False, cost_callback=None):
    """ Run simple lifted multicut segmentation workflow from affinity or boundary maps and using lifted edges
    derived from priors by mapping a separate segmentation to the watersheds.

    Adapted from "Leveraging domain knowledge to improve microscopy image segmentation with lifted multicuts":
    https://doi.org/10.3389/fcomp.2019.00006

    Arguments:
        input_ [np.ndarray] - input data, either boundary or affinity map
        segmentation [np.ndarray] - the segmentation used to derive priors
        use_2dws [bool] - whether to run watersheds in 2d and stack them
            or if the watersheds passed are 2d and stacked
        overlap_threshold [float] - minimal overlap to assign a segment id to node
        graph_depth [int] - maximal graph depth up to which
            lifted edges will be included
        same_segment_cost [float] - costs for edges between nodes with same segment id attribution
        different_segment_cost [float] - costs for edges between nodes with different segment id attribution
        lifted_multicut_solver [str or callable] - name of multicut solver in elf.segmentation.multicut
            or custom solver function (default: 'kernighan-lin')
        watershed [np.ndarray] - watershed over-segmentation, if None is passed,
            the overr-segmentation witll be computted (default: None)
        mask [np.ndarray] - binary mask to exclude from the segmentation;
            values that are False will be excluded (default: None)
        weighting_scheme [str] - strategy to weight multicut edge costs by size (default: None)
        ws_kwargs [dict] - keyword arguments for the watershed function (default: None)
        mode [str] - mode for insertion of lifted edges. Can be
            "all" - lifted edges will be inserted in between all nodes with attribution
            "different" - lifted edges will only be inserted in between nodes attributed to different classes
            "same" - lifted edges will only be inserted in between nodes attribted to the same class
            (default: "different")
        solver_kwargs [dict] - keyword arguments for the lifted multicut solver (default: {})
        beta [float] - boundary bias (default: 0.5)
        offsets [list] - pixel offsets for affintities, by default assume nearest neighbor offsets (default: None)
        n_threads [int] - number of threads (default: None)
        return_intermediates [bool] - whether to also return intermediate results (default: False)
        cost_callback [callable] - callback to modify the costs.
            It will be passed the costs, rag and watershed (default: None)
    """
    # get the multicut solver
    solver = _get_lifted_solver(lifted_multicut_solver)

    # create the multicut problem
    watershed, rag, edge_probs, edge_costs = _simple_mc_problem(
        input_, use_2dws, watershed, mask, weighting_scheme, ws_kwargs,
        beta, offsets, n_threads, cost_callback
    )

    # get the lifted edges and costs
    lifted_uvs, lifted_costs = elf_feats.lifted_problem_from_segmentation(rag, watershed, segmentation,
                                                                          overlap_threshold, graph_depth,
                                                                          same_segment_cost, different_segment_cost,
                                                                          mode=mode, n_threads=n_threads)

    # solve the multicut problem and project back to segmentation
    # we pass the watershed to the solver as well, because it is needed for
    # the blockwise-multicut solver
    node_labels = solver(rag, edge_costs, lifted_uvs, lifted_costs,
                         n_threads=n_threads, segmentation=watershed, **solver_kwargs)
    seg = elf_feats.project_node_labels_to_pixels(rag, node_labels, n_threads)

    if return_intermediates:
        return {'watershed': watershed,
                'rag': rag,
                'probs': edge_probs,
                'costs': edge_costs,
                'lifted_edges': lifted_uvs,
                'lifted_costs': lifted_costs,
                'node_labels': node_labels,
                'segmentation': seg}
    else:
        return seg

max_stars_repo_name: qin-yu/elf

max_issues_repo_name: qin-yu/elf

max_forks_repo_name: qin-yu/elf