text:
'''
This package provides a function and a base class:
    unify and AbstractSegmentation.

unify(data_sets, directory, size=None, segmentation_size=None, crop=False,
        min_frequency=None, min_coverage=None, synonyms=None,
        test_limit=None, single_process=False, verbose=False):

unify creates a stadnard format multichannel segementation
out of heterogenous labeled image sources.

Input:

    data_sets is a list of instances of subclasses of
    AbstractSegmentation (or similar objects).  Five methods
    must be provided on each object.

    names(cat, j) - a list of names that describe label j
        as it is used within category cat in this dataset.

    size() - number of images in the dataset.

    filename(i) - the jpg filename for the i'th image in the dataset.

    metadata(i) - a small pickleable record from which the
         segmentation of the ith item can be derived, withou reference
         to a specific corpus instance.

    @classmethod
    resolve_segmentation(m) - resolves to a dict and a shape of a segmentation.
         The dict contains string key names and numpy-array vals describing
         the segmentation of the image drived from metadata m.  The keys
         of this dict are arbitrary short strings, but for example could
         look like:
       {
         'scene': 1293, # 0-d: applies one label to the whole image
         'object': (1024x1280 int16), # 2-d: one label per pixel
         'part': (3x1024x1280 int16), # 3-d: layers of labels, dominant first
         'texture': [14,22] # 1-d: multiple whole-image labels, dominant first
       }

    The reason for the separation between metadata(i) and
    resolve_segmentation(m) so that we can run metadata(i) quickly on the
    main process, and then shard resolve_segmentation(m) out to worker
    processes.

Output is the following directory structure:
'''
README_TEXT= '''
This directory contains the following data and metadata files:

    images/[datasource]/...
         images drawn from a specific datasource are reformatted
         and scaled and saved as jpeg and png files in subdirectories
         of the images directory.

    index.csv
        contains a list of all the images in the dataset, together with
        available labeled data, in the form:

        image,split,ih,iw,sh,sw,[color,object,material,part,scene,texture]

        for examplle:
        dtd/g_0106.jpg,train,346,368,346,368,dtd/g_0106_color.png,,,,,314;194

        The first column is always the original image filename relative to
        the images/ subdirectory; then image height and width and segmentation
        heigh and width dimensions are followed by six columns for label
        data in the six categories.  Label data can be per-pixel or per-image
        (assumed to apply to every pixel in the image), and 0 or more labels
        can be specified per category.  If there are no labels, the field is ''.
        If there are multiple labels, they are separated by semicolons,
        and it is assumed that dominant interpretations are listed first.
        Per-image labels are represented by a decimal number; and per-image
        labels are represented by a filename for an image which encodes
        the per-pixel labels in the (red + 256 * green) channels.

    category.csv
        name,first,last,count,frequency

        for example:
        object,12,1138,529,208688

        In the generic case there may not be six categories; this directory
        may contain any set of categories of segmentations.  This file
        lists all segmentation categories found in the data sources,
        along with statistics on now many class labels apply to
        each category, and in what range; as well as how many
        images mention a label of that category

    label.csv
        number,name,category,frequency,coverage,syns

        for example:
        10,red-c,color(289),289,9.140027,
        21,fabric,material(36);object(3),39,4.225474,cloth

        This lists all label numbers and corresponding names, along
        with some statistics including the number of images for
        which the label appears in each category; the total number
        of images which have the label; and the pixel portions
        of images that have the label.

    c_[category].csv (for example, map_color.csv)
        code,number,name,frequency,coverage

        for example:
        4,31,glass,27,0.910724454131

        Although labels are store under a unified code, this file
        lists a standard dense coding that can be used for a
        specific subcategory of labels.
'''

import codecs
from functools import partial
import itertools
from multiprocessing import Pool, cpu_count
import numpy
import operator
import os
import re
import subprocess
import time
from PIL import Image
from collections import OrderedDict
from scipy.misc import imread, imresize, imsave
from scipy.ndimage.interpolation import zoom
import signal
import shutil
import sys
from unicsv import DictUnicodeWriter

def unify(data_sets, directory, size=None, segmentation_size=None, crop=False,
        splits=None, min_frequency=None, min_coverage=None,
        synonyms=None, test_limit=None, single_process=False, verbose=False):
    # Make sure we have a directory to work in
    directory = os.path.expanduser(directory)
    ensure_dir(directory)

    # Step 0: write a README file with generated information.
    write_readme_file([
        ('data_sets', data_sets), ('size', size),
        ('segmentation_size', segmentation_size), ('crop', crop),
        ('splits', splits),
        ('min_frequency', min_frequency), ('min_coverage', min_coverage),
        ('synonyms', synonyms), ('test_limit', test_limit),
        ('single_process', single_process)],
        directory=directory, verbose=verbose)

    # Clear old images data
    try:
        if verbose:
            print 'Removing old images.'
        shutil.rmtree(os.path.join(directory, 'images'), ignore_errors=True)
    except:
        pass
    ensure_dir(os.path.join(directory, 'images'))
    # Phase 1: Count label statistics
    # frequency = number of images touched by each label
    # coverage  = total portion of images covered by each label
    frequency, coverage = gather_label_statistics(
            data_sets, test_limit, single_process,
            segmentation_size, crop, verbose)
    # Phase 2: Sort, collapse, and filter labels
    labnames, syns = normalize_labels(data_sets, frequency, coverage, synonyms)
    report_aliases(directory, labnames, data_sets, frequency, verbose)
    # Phase 3: Filter by frequncy, and assign numbers
    names, assignments = assign_labels(
            labnames, frequency, coverage, min_frequency, min_coverage, verbose)
    # Phase 4: Collate and output label stats
    cats = write_label_files(
            directory, names, assignments, frequency, coverage, syns, verbose)
    # Phase 5: Create normalized segmentation files
    create_segmentations(
            directory, data_sets, splits, assignments, size, segmentation_size,
            crop, cats, test_limit, single_process, verbose)

def gather_label_statistics(data_sets, test_limit, single_process,
        segmentation_size, crop, verbose):
    '''
    Phase 1 of unification.  Counts label statistics.
    '''
    # Count frequency and coverage for each individual image
    stats = map_in_pool(partial(count_label_statistics,
        segmentation_size=segmentation_size,
        crop=crop,
        verbose=verbose),
            all_dataset_segmentations(data_sets, test_limit),
            single_process=single_process,
            verbose=verbose)
    # Add them up
    frequency, coverage = (sum_histogram(d) for d in zip(*stats))
    # TODO: also, blacklist images that have insufficient labled pixels
    return frequency, coverage

def normalize_labels(data_sets, frequency, coverage, synonyms):
    '''
    Phase 2 of unification.
    Assigns unique label names and resolves duplicates by name.
    '''
    # Sort by frequency, descending, and assign a name to each label
    top_syns = {}
    labname = {}
    freq_items = zip(*sorted((-f, lab) for lab, f in frequency.items()))[1]
    for lab in freq_items:
        dataset, category, label = lab
        names = [n.lower() for n in data_sets[dataset].all_names(
              category, label) if len(n) and n != '-']
        if synonyms:
            names = synonyms(names)
        # Claim the synonyms that have not already been taken
        for name in names:
            if name not in top_syns:
                top_syns[name] = lab
        # best_labname may decide to collapse two labels because they
        # have the same names and seem to mean the same thing
        labname[lab], unique = best_labname(
                lab, names, labname, top_syns, coverage, frequency)
    return labname, top_syns

def report_aliases(directory, labnames, data_sets, frequency, verbose):
    '''
    Phase 2.5
    Report the aliases.  These are printed into 'syns.txt'
    '''
    show_all = True  # Print all details to help debugging
    name_to_lab = invert_dict(labnames)
    with codecs.open(os.path.join(directory, 'syns.txt'), 'w', 'utf-8') as f:
        def report(txt):
            f.write('%s\n' % txt)
            if verbose:
                print txt
        for name in sorted(name_to_lab.keys(),
                key=lambda n: (-len(name_to_lab[n]), n)):
            keys = name_to_lab[name]
            if not show_all and len(keys) <= 1:
                break
            # Don't bother reporting aliases if all use the same index;
            # that is probably due to an underlying shared code.
            if not show_all and len(set(i for d, c, i in keys)) <= 1:
                continue
            report('name: %s' % name)
            for ds, cat, i in keys:
                names = ';'.join(data_sets[ds].all_names(cat, i))
                freq = frequency[(ds, cat ,i)]
                report('%s/%s#%d: %d, (%s)' % (ds, cat, i, freq, names))
            report('')

def assign_labels(
        labnames, frequency, coverage, min_frequency, min_coverage, verbose):
    '''
    Phase 3 of unification.
    Filter names that are too infrequent, then assign numbers.
    '''
    # Collect by-name frequency and coverage
    name_frequency = join_histogram(frequency, labnames)
    name_coverage = join_histogram(coverage, labnames)
    names = name_frequency.keys()
    if min_frequency is not None:
        names = [n for n in names if name_frequency[n] >= min_frequency]
    if min_coverage is not None:
        names = [n for n in names if name_coverage[n] >= min_coverage]
    # Put '-' at zero
    names = [n for n in names if n != '-']
    names = ['-'] + sorted(names,
            key=lambda x: (-name_frequency[x], -name_coverage[x]))
    nums = dict((n, i) for i, n in enumerate(names))
    assignments = dict((k, nums.get(v, 0)) for k, v in labnames.items())
    return names, assignments

def write_label_files(
        directory, names, assignments, frequency, coverage, syns, verbose):
    '''
    Phase 4 of unification.
    Collate some stats and then write then to two metadata files.
    '''
    # Make lists of synonyms claimed by each label
    synmap = invert_dict(dict((w, assignments[lab]) for w, lab in syns.items()))
    # We need an (index, category) count
    ic_freq = join_histogram_fn(frequency, lambda x: (assignments[x], x[1]))
    ic_cov = join_histogram_fn(coverage, lambda x: (assignments[x], x[1]))
    for z in [(j, cat) for j, cat in ic_freq if j == 0]:
        del ic_freq[z]
        del ic_cov[z]
    catstats = [[] for n in names]
    # For each index, get a (category, frequency) list in descending order
    for (ind, cat), f in sorted(ic_freq.items(), key=lambda x: -x[1]):
        catstats[ind].append((cat, f))
    index_coverage = join_histogram(coverage, assignments)
    with open(os.path.join(directory, 'label.csv'), 'w') as csvfile:
        fields = ['number', 'name', 'category', 'frequency', 'coverage', 'syns']
        writer = DictUnicodeWriter(csvfile, fieldnames=fields)
        writer.writeheader()
        for ind, name in enumerate(names):
            if ind == 0:
                continue
            writer.writerow(dict(
                number='%d' % ind,
                name=name,
                category=';'.join('%s(%d)' % s for s in catstats[ind]),
                frequency='%d' % sum(f for c, f in catstats[ind]),
                coverage='%f' % index_coverage[ind],
                syns=';'.join([s for s in synmap[ind] if s != name])
            ))
    # For each category, figure the first, last, and other stats
    cat_ind = [(cat, ind) for ind, cat in ic_freq.keys()]
    first_index = build_histogram(cat_ind, min)
    last_index = build_histogram(cat_ind, max)
    count_labels = build_histogram([(cat, 1) for cat, _ in cat_ind])
    cat_freq = join_histogram_fn(ic_freq, lambda x: x[1])
    cats = sorted(first_index.keys(), key=lambda x: first_index[x])
    with open(os.path.join(directory, 'category.csv'), 'w') as csvfile:
        fields = ['name', 'first', 'last', 'count', 'frequency']
        writer = DictUnicodeWriter(csvfile, fieldnames=fields)
        writer.writeheader()
        for cat in cats:
            writer.writerow(dict(
                name=cat,
                first=first_index[cat],
                last=last_index[cat],
                count=count_labels[cat],
                frequency=cat_freq[cat]))
    # And for each category, create a dense coding file.
    for cat in cats:
        dense_code = [0] + sorted([i for i, c in ic_freq if c == cat],
                key=lambda i: (-ic_freq[(i, cat)], -ic_cov[(i, cat)]))
        fields = ['code', 'number', 'name', 'frequency', 'coverage']
        with open(os.path.join(directory, 'c_%s.csv' % cat), 'w') as csvfile:
            writer = DictUnicodeWriter(csvfile, fieldnames=fields)
            writer.writeheader()
            for code, i in enumerate(dense_code):
                if code == 0:
                    continue
                writer.writerow(dict(
                    code=code,
                    number=i,
                    name=names[i],
                    frequency=ic_freq[(i, cat)],
                    coverage=ic_cov[(i, cat)]))
    return cats

def create_segmentations(directory, data_sets, splits, assignments, size,
        segmentation_size, crop, cats, test_limit, single_process, verbose):
    '''
    Phase 5 of unification.  Create the normalized segmentation files
    '''
    if size is not None and segmentation_size is None:
        segmentation_size = size
    # Get assignments into a nice form, once, here.
    # (dataset, category): [numpy array with new indexes]
    index_max = build_histogram(
            [((ds, cat), i) for ds, cat, i in assignments.keys()], max)
    index_mapping = dict([k, numpy.zeros(i + 1, dtype=numpy.int16)]
            for k, i in index_max.items())
    for (ds, cat, oldindex), newindex in assignments.items():
        index_mapping[(ds, cat)][oldindex] = newindex
    # Count frequency and coverage for each individual image
    segmented = map_in_pool(
            partial(translate_segmentation,
                directory=directory,
                mapping=index_mapping,
                size=size,
                segmentation_size=segmentation_size,
                categories=cats,
                crop=crop,
                verbose=verbose),
            all_dataset_segmentations(data_sets, test_limit),
            single_process=single_process,
            verbose=verbose)
    # Sort nonempty itesm randomly+reproducibly by md5 hash of the filename.
    ordered = sorted([(hashed_float(r['image']), r) for r in segmented if r])
    # Assign splits, pullout out last 20% for validation.
    cutoffs = cumulative_splits(splits)
    for floathash, record in ordered:
        for name, cutoff in cutoffs:
            if floathash <= cutoff:
                record['split'] = name
                break
        else:
            assert False, 'hash %f exceeds last split %f' % (floathash, c)

    # Now write one row per image and one column per category
    with open(os.path.join(directory, 'index.csv'), 'w') as csvfile:
        fields = ['image', 'split', 'ih', 'iw', 'sh', 'sw'] + cats
        writer = DictUnicodeWriter(csvfile, fieldnames=fields)
        writer.writeheader()
        for f, record in ordered:
            writer.writerow(record)

def translate_segmentation(record, directory, mapping, size,
        segmentation_size, categories, crop, verbose):
    '''
    Translates a single segmentation.
    '''
    dataset, index, seg_class, filename, md = record
    basename = os.path.basename(filename)
    if verbose:
        print 'Processing #%d %s %s' % (index, dataset, basename)
    full_seg, shape = seg_class.resolve_segmentation(md)
    # Rows can be omitted by returning no segmentations.
    if not full_seg:
        return None
    jpg = imread(filename)
    if size is not None:
        jpg = scale_image(jpg, size, crop)
        for cat in full_seg:
            full_seg[cat] = scale_segmentation(
                    full_seg[cat], segmentation_size, crop)
    else:
        size = jpg.shape[:2]
        segmentation_size = (1,1)
        for cat in full_seg:
            if len(numpy.shape(full_seg[cat])) >= 2:
                segmentation_size = numpy.shape(full_seg[cat])
                break

    imagedir = os.path.join(directory, 'images')
    ensure_dir(os.path.join(imagedir, dataset))
    fn = save_image(jpg, imagedir, dataset, basename)
    result = {
            'image': os.path.join(dataset, fn),
            'ih': size[0],
            'iw': size[1],
            'sh': segmentation_size[0],
            'sw': segmentation_size[1]
    }
    for cat in full_seg:
        if cat not in categories:
            continue  # skip categories with no data globally
        result[cat] = ';'.join(save_segmentation(full_seg[cat],
                imagedir, dataset, fn, cat, mapping[(dataset, cat)]))
    return result

def best_labname(lab, names, assignments, top_syns, coverage, frequency):
    '''
    Among the given names, chooses the best name to assign, given
    information about previous assignments, synonyms, and stats
    '''
    # Best shot: get my own name, different from other names.
    if 'dog' in names:
        print names
    for name in names:
        if top_syns[name] == lab:
            return name, True
    if len(names) == 0 or len(names) == '1' and names[0] in ['', '-']:
        # If there are no names, we must use '-' and map to 0.
        return '-', False
    elif len(names) == 1:
        # If we have a conflict without ambiguity, then we just merge names.
        other = top_syns[names[0]]
    else:
        # If we need to merge and we have multiple synonyms, let's score them.
        scores = []
        # Get the average size of an object of my type
        size = coverage[lab] / frequency[lab]
        for name in names:
            other = top_syns[name]
            # Compare it to the average size of objects of other types
            other_size = coverage[other] / frequency[other]
            scores.append((abs(size - other_size), -frequency[other], other))
        # Choose the synonyms that has the most similar average size.
        # (breaking ties according to maximum frequency)
        other = min(scores)[2]
    name = assignments[other]
    return name, False

def build_histogram(pairs, reducer=operator.add):
    '''Creates a histogram by combining a list of key-value pairs.'''
    result = {}
    for k, v in pairs:
        if k not in result:
            result[k] = v
        else:
            result[k] = reducer(result[k], v)
    return result

def join_histogram_fn(histogram, makekey):
    '''Rekeys histogram according to makekey fn, summing joined buckets.'''
    result = {}
    for oldkey, val in histogram.iteritems():
        newkey = makekey(oldkey)
        if newkey not in result:
            result[newkey] = val
        else:
            result[newkey] += val
    return result

def join_histogram(histogram, newkeys):
    '''Rekeys histogram according to newkeys map, summing joined buckets.'''
    result = {}
    for oldkey, newkey in newkeys.iteritems():
        if newkey not in result:
            result[newkey] = histogram[oldkey]
        else:
            result[newkey] += histogram[oldkey]
    return result

def sum_histogram(histogram_list):
    '''Adds histogram dictionaries elementwise.'''
    result = {}
    for d in histogram_list:
        for k, v in d.iteritems():
            if k not in result:
                result[k] = v
            else:
                result[k] += v
    return result

def invert_dict(d):
    '''Transforms {k: v} to {v: [k,k..]}'''
    result = {}
    for k, v in d.iteritems():
        if v not in result:
            result[v] = [k]
        else:
            result[v].append(k)
    return result

def count_label_statistics(record, segmentation_size, crop, verbose):
    '''
    Resolves the segmentation record, and then counts all nonzero
    labeled pixels in the resuting segmentation.  Returns two maps:
      freq[(dataset, category, label)] = 1, if the label is present
      coverage[(dataset, category, label)] = p, for portion of pixels covered
    '''
    dataset, index, seg_class, fn, md = record
    if verbose:
        print 'Counting #%d %s %s' % (index, dataset, os.path.basename(fn))
    full_seg, shape = seg_class.resolve_segmentation(md)
    freq = {}
    coverage = {}
    for category, seg in full_seg.items():
        if seg is None:
            continue
        dims = len(numpy.shape(seg))
        if dims <= 1:
            for label in (seg if dims else (seg,)):
                key = (dataset, category, int(label))
                freq[key] = 1
                coverage[key] = 1.0
        elif dims >= 2:
            # We do _not_ scale the segmentation for counting purposes!
            # Different scales should produce the same labels and label order.
            # seg = scale_segmentation(seg, segmentation_size, crop=crop)
            bc = numpy.bincount(seg.ravel())
            pixels = numpy.prod(seg.shape[-2:])
            for label in bc.nonzero()[0]:
                if label > 0:
                    key = (dataset, category, int(label))
                    freq[key] = 1
                    coverage[key] = float(bc[label]) / pixels
    return freq, coverage

def all_dataset_segmentations(data_sets, test_limit=None):
    '''
    Returns an iterator for metadata over all segmentations
    for all images in all the datasets.  The iterator iterates over
    (dataset_name, global_index, dataset_resolver, metadata(i))
    '''
    j = 0
    for name, ds in data_sets.items():
        for i in truncate_range(range(ds.size()), test_limit):
            yield (name, j, ds.__class__, ds.filename(i), ds.metadata(i))
            j += 1

def truncate_range(data, limit):
    '''For testing, if limit is not None, limits data by slicing it.'''
    if limit is None:
        return data
    if isinstance(limit, slice):
        return data[limit]
    return data[:limit]

def map_in_pool(fn, data, single_process=False, verbose=False):
    '''
    Our multiprocessing solution; wrapped to stop on ctrl-C well.
    '''
    if single_process:
        return map(fn, data)
    n_procs = min(cpu_count(), 12)
    original_sigint_handler = setup_sigint()
    pool = Pool(processes=n_procs, initializer=setup_sigint)
    restore_sigint(original_sigint_handler)
    try:
        if verbose:
            print 'Mapping with %d processes' % n_procs
        res = pool.map_async(fn, data)
        return res.get(31536000)
    except KeyboardInterrupt:
        print("Caught KeyboardInterrupt, terminating workers")
        pool.terminate()
        raise
    else:
        pool.close()
        pool.join()

def setup_sigint():
    return signal.signal(signal.SIGINT, signal.SIG_IGN)

def restore_sigint(original):
    signal.signal(signal.SIGINT, original)

def scale_image(im, dims, crop=False):
    '''
    Scales or crops a photographic image using antialiasing.
    '''
    if len(im.shape) == 2:
        # Handle grayscale images by adding an RGB channel
        im = numpy.repeat(im[numpy.newaxis], 3, axis=0)
    if im.shape[0:2] != dims:
        if not crop:
            im = imresize(im, dims)
        else:
            source = im.shape[0:2]
            aspect = float(dims[1]) / dims[0]
            if aspect * source[0] > source[1]:
                width = int(dims[1] / aspect)
                margin = (width - dims[0]) // 2
                im = imresize(im, (width, dims[1]))[
                        margin:margin+dims[0],:,:]
            else:
                height = int(dims[0] * aspect)
                margin = (height - dims[1]) // 2
                im = imresize(im, (dims[0], height))[
                        margin:margin+dims[1],:,:]
    return im

def scale_segmentation(segmentation, dims, crop=False):
    '''
    Zooms a 2d or 3d segmentation to the given dims, using nearest neighbor.
    '''
    shape = numpy.shape(segmentation)
    if len(shape) < 2 or shape[-2:] == dims:
        return segmentation
    peel = (len(shape) == 2)
    if peel:
        segmentation = segmentation[numpy.newaxis]
    levels = segmentation.shape[0]
    result = numpy.zeros((levels, ) + dims,
            dtype=segmentation.dtype)
    ratio = (1,) + tuple(res / float(orig)
            for res, orig in zip(result.shape[1:], segmentation.shape[1:]))
    if not crop:
        safezoom(segmentation, ratio, output=result, order=0)
    else:
        ratio = max(ratio[1:])
        height = int(round(dims[0] / ratio))
        hmargin = (segmentation.shape[0] - height) // 2
        width = int(round(dims[1] / ratio))
        wmargin = (segmentation.shape[1] - height) // 2
        safezoom(segmentation[:, hmargin:hmargin+height,
            wmargin:wmargin+width],
            (1, ratio, ratio), output=result, order=0)
    if peel:
        result = result[0]
    return result

def safezoom(array, ratio, output=None, order=0):
    '''Like numpy.zoom, but does not crash when the first dimension
    of the array is of size 1, as happens often with segmentations'''
    dtype = array.dtype
    if array.dtype == numpy.float16:
        array = array.astype(numpy.float32)
    if array.shape[0] == 1:
        if output is not None:
            output = output[0,...]
        result = zoom(array[0,...], ratio[1:],
                output=output, order=order)
        if output is None:
            output = result[numpy.newaxis]
    else:
        result = zoom(array, ratio, output=output, order=order)
        if output is None:
            output = result
    return output.astype(dtype)

def save_image(im, imagedir, dataset, filename):
    '''
    Try to pick a unique name and save the image with that name.
    This is not race-safe, so the given name should already be unique.
    '''
    trynum = 1
    fn = filename
    while os.path.exists(os.path.join(imagedir, dataset, fn)):
        trynum += 1
        fn = re.sub('(?:\.jpg)?$', '%d.jpg' % trynum, filename)
    imsave(os.path.join(imagedir, dataset, fn), im)
    return fn

def save_segmentation(seg, imagedir, dataset, filename, category, translation):
    '''
    Saves the segmentation in a file or files if it is large, recording
    the filenames.  Or serializes it as decimal numbers to record if it
    is small.  Then returns the array of strings recorded.
    '''
    if seg is None:
        return None
    shape = numpy.shape(seg)
    if len(shape) < 2:
        # Save numbers as decimal strings; and omit zero labels.
        return ['%d' % translation[t]
                for t in (seg if len(shape) else [seg]) if t]

    result = []
    for channel in ([()] if len(shape) == 2 else range(shape[0])):
        # Save bitwise channels as filenames of PNGs; and save the files.
        im = encodeRG(translation[seg[channel]])
        if len(shape) == 2:
            fn = re.sub('(?:\.jpg)?$', '_%s.png' % category, filename)
        else:
            fn = re.sub('(?:\.jpg)?$', '_%s_%d.png' %
                    (category, channel + 1), filename)
        result.append(os.path.join(dataset, fn))
        imsave(os.path.join(imagedir, dataset, fn), im)
    return result

def encodeRG(channel):
    '''Encodes a 16-bit per-pixel code using the red and green channels.'''
    result = numpy.zeros(channel.shape + (3,), dtype=numpy.uint8)
    result[:,:,0] = channel % 256
    result[:,:,1] = (channel // 256)
    return result

def ensure_dir(targetdir):
    if not os.path.isdir(targetdir):
        try:
            os.makedirs(targetdir)
        except:
            pass

def hashed_float(s):
    # Inspired by http://code.activestate.com/recipes/391413/ by Ori Peleg
    '''Hashes a string to a float in the range [0, 1).'''
    import md5, struct
    [number] = struct.unpack(">Q", md5.new(s).digest()[:8])
    return number / (2.0 ** 64)  # python will constant-fold this denominator.

def cumulative_splits(splits):
    '''Converts [0.8, 0.1, 0.1] to [0.8, 0.9, 1.0]'''
    if splits is None:
        return [('train', 1.0)]  # Default to just one split.
    result = []
    c = 0.0
    for name, s in splits.items():
        c += s
        result.append((name, c))
    # Eliminate any fp rounding problem: round last split up to 1.0
    if result[-1][1] < 1.0 - len(result) * sys.float_info.epsilon:
        raise ValueError('splits must add to 1.0, but %s add to %s' % (
            repr(splits), result[-1][1]))
    result[-1] = (result[-1][0], 1.0)
    return result

def write_readme_file(args, directory, verbose):
    '''
    Writes a README.txt that describes the settings used to geenrate the ds.
    '''
    with codecs.open(os.path.join(directory, 'README.txt'), 'w', 'utf-8') as f:
        def report(txt):
            f.write('%s\n' % txt)
            if verbose:
                print txt
        title = '%s joined segmentation data set' % os.path.basename(directory)
        report('%s\n%s' % (title, '=' * len(title)))
        for key, val in args:
            if key == 'data_sets':
                report('Joins the following data sets:')
                for name, kind in val.items():
                    report('    %s: %d images' % (name, kind.size()))
            else:
                report('%s: %r' % (key, val))
        report('\ngenerated at: %s' % time.strftime("%Y-%m-%d %H:%M"))
        try:
            label = subprocess.check_output(['git', 'rev-parse', 'HEAD'])
            report('git label: %s' % label)
        except:
            pass
        f.write('\n')
        # Add boilerplate readme text
        f.write(README_TEXT)

if __name__ == '__main__':
    import argparse
    import adeseg
    import dtdseg
    import osseg
    import pascalseg
    from synonym import synonyms

    parser = argparse.ArgumentParser(
        description='Generate broden dataset.')
    parser.add_argument(
            '--size',
            type=int, default=224,
            help='pixel size for input images, e.g., 224 or 227')
    args = parser.parse_args()

    image_size = (args.size, args.size)
    seg_size = (args.size // 2, args.size // 2)

    print 'CREATING NEW SEGMENTATION OF SIZE %d.\n' % args.size
    print 'Loading source segmentations.'
    ade = adeseg.AdeSegmentation('dataset/ade20k', 'ADE20K_2016_07_26')
    dtd = dtdseg.DtdSegmentation('dataset/dtd/dtd-r1.0.1')
    # OpenSurfaces is not balanced in scene and object types.
    oss = osseg.OpenSurfaceSegmentation('dataset/opensurfaces/',
            supply=set(['material', 'color']))
    # Remove distinction between upper-arm, lower-arm, etc.
    pascal = pascalseg.PascalSegmentation('dataset/pascal/',
            collapse_adjectives=set([
                'left', 'right', 'front', 'back', 'upper', 'lower', 'side']))
    data = OrderedDict(ade20k=ade, dtd=dtd, opensurfaces=oss, pascal=pascal)
    unify(data,
            splits=OrderedDict(train=0.7, val=0.3),
            size=image_size, segmentation_size=seg_size,
            directory=('dataset/broden1_%d' % args.size),
            synonyms=synonyms,
            min_frequency=10, min_coverage=0.5, verbose=True)

max_stars_repo_name: rathoreanirudh/NetDissect

max_issues_repo_name: xujinglin/NetDissect

max_forks_repo_name: xujinglin/NetDissect

reason: import numpy,from scipy


text:
import torch
import numpy as np
import pandas as pd

from sksurv.functions import StepFunction

def get_numpy(df, long = ["Y1","Y2","Y3"], base = ["X1","X2"], obstime = "obstime", max_len=None):
    df.loc[:,"id_new"] = df.groupby(by="id").grouper.group_info[0] # assign id from 0 to num subjects
    df.loc[:,"roundtime"] = (df.loc[:,obstime] * 2).round() / 2 # round obstime to nearest 0.5
    if "visit" not in df:
        df.loc[:,"visit"] = df.groupby(by="id").cumcount()

    I = len(np.unique(df.loc[:,"id"]))
    if max_len is None:
        max_len = int(np.max(df.loc[:,"roundtime"]) * 2 + 1) # based on 0.5 rounding

    x_long = np.empty((I, max_len, len(long)))
    x_long[:,:,:] = np.NaN
    x_base = np.zeros((I, len(base)))
    for index, row in df.iterrows():
        ii = int(row.loc["id_new"])
        jj = int(row.loc["roundtime"]) * 2 # based on 0.5 rounding
        x_long[ii,jj,:] = row.loc[long]
        if jj==0:
            x_base[ii,:] = row.loc[base]

    e = df.loc[df["visit"]==0,"event"].values.squeeze()
    t = df.loc[df["visit"]==0,"time"].values.squeeze()
    obs_time = np.arange(0, max_len/2, 0.5)

    return x_long, x_base, e, t, obs_time

def sortByTime(x, e, t):
    # sort data by descending survival time
    t_index_desc = torch.argsort(t, descending=True)
    x = x[t_index_desc]
    e = e[t_index_desc]
    t = t[t_index_desc]
    return x, e, t

## Loss function
def surv_loss(risk, event):
    hr = torch.exp(risk)
    log_risk = torch.log(torch.cumsum(hr,0))
    uncensored_LL = risk - log_risk
    censored_LL = uncensored_LL * event.unsqueeze(1)
    neg_likelihood = -torch.sum(censored_LL)# / torch.sum(event)
    return neg_likelihood

def _compute_counts(event, time, order=None):
    """Count right censored and uncensored samples at each unique time point.
    Parameters
    ----------
    event : array
        Boolean event indicator.
    time : array
        Survival time or time of censoring.
    order : array or None
        Indices to order time in ascending order.
        If None, order will be computed.
    Returns
    -------
    times : array
        Unique time points.
    n_events : array
        Number of events at each time point.
    n_at_risk : array
        Number of samples that have not been censored or have not had an event at each time point.
    n_censored : array
        Number of censored samples at each time point.
    """
    n_samples = event.shape[0]

    if order is None:
        order = np.argsort(time, kind="mergesort")

    uniq_times = np.empty(n_samples, dtype=time.dtype)
    uniq_events = np.empty(n_samples, dtype=np.int_)
    uniq_counts = np.empty(n_samples, dtype=np.int_)

    i = 0
    prev_val = time[order[0]]
    j = 0
    while True:
        count_event = 0
        count = 0
        while i < n_samples and prev_val == time[order[i]]:
            if event[order[i]]:
                count_event += 1

            count += 1
            i += 1

        uniq_times[j] = prev_val
        uniq_events[j] = count_event
        uniq_counts[j] = count
        j += 1

        if i == n_samples:
            break

        prev_val = time[order[i]]

    times = np.resize(uniq_times, j)
    n_events = np.resize(uniq_events, j)
    total_count = np.resize(uniq_counts, j)
    n_censored = total_count - n_events

    # offset cumulative sum by one
    total_count = np.r_[0, total_count]
    n_at_risk = n_samples - np.cumsum(total_count)

    return times, n_events, n_at_risk[:-1], n_censored

class BreslowEstimator:
    """Breslow's estimator of the cumulative hazard function.
    Attributes
    ----------
    cum_baseline_hazard_ : :class:`sksurv.functions.StepFunction`
        Cumulative baseline hazard function.
    baseline_survival_ : :class:`sksurv.functions.StepFunction`
        Baseline survival function.
    """

    def fit(self, linear_predictor, event, time):
        """Compute baseline cumulative hazard function.
        Parameters
        ----------
        linear_predictor : array-like, shape = (n_samples,)
            Linear predictor of risk: `X @ coef`.
        event : array-like, shape = (n_samples,)
            Contains binary event indicators.
        time : array-like, shape = (n_samples,)
            Contains event/censoring times.
        Returns
        -------
        self
        """
        risk_score = np.exp(linear_predictor)
        order = np.argsort(time, kind="mergesort")
        risk_score = risk_score[order]
        uniq_times, n_events, n_at_risk, _ = _compute_counts(event, time, order)

        divisor = np.empty(n_at_risk.shape, dtype=np.float_)
        value = np.sum(risk_score)
        divisor[0] = value
        k = 0
        for i in range(1, len(n_at_risk)):
            d = n_at_risk[i - 1] - n_at_risk[i]
            value -= risk_score[k:(k + d)].sum()
            k += d
            divisor[i] = value

        assert k == n_at_risk[0] - n_at_risk[-1]

        y = np.cumsum(n_events / divisor)
        self.cum_baseline_hazard_ = StepFunction(uniq_times, y)
        self.baseline_survival_ = StepFunction(self.cum_baseline_hazard_.x,
                                               np.exp(- self.cum_baseline_hazard_.y))
        return self

    def get_cumulative_hazard_function(self, linear_predictor):
        """Predict cumulative hazard function.
        Parameters
        ----------
        linear_predictor : array-like, shape = (n_samples,)
            Linear predictor of risk: `X @ coef`.
        Returns
        -------
        cum_hazard : ndarray, shape = (n_samples,)
            Predicted cumulative hazard functions.
        """
        risk_score = np.exp(linear_predictor)
        n_samples = risk_score.shape[0]
        funcs = np.empty(n_samples, dtype=np.object_)
        for i in range(n_samples):
            funcs[i] = StepFunction(x=self.cum_baseline_hazard_.x,
                                    y=self.cum_baseline_hazard_.y,
                                    a=risk_score[i])
        return funcs

    def get_survival_function(self, linear_predictor):
        """Predict survival function.
        Parameters
        ----------
        linear_predictor : array-like, shape = (n_samples,)
            Linear predictor of risk: `X @ coef`.
        Returns
        -------
        survival : ndarray, shape = (n_samples,)
            Predicted survival functions.
        """
        risk_score = np.exp(linear_predictor)
        n_samples = risk_score.shape[0]
        funcs = np.empty(n_samples, dtype=np.object_)
        for i in range(n_samples):
            funcs[i] = StepFunction(x=self.baseline_survival_.x,
                                    y=np.power(self.baseline_survival_.y, risk_score[i]))
        return funcs

ext: py

max_stars_repo_name: reylined/TransformerJM

max_issues_repo_name: reylined/TransformerJM

max_forks_repo_name: reylined/TransformerJM

reason: import numpy


text:
import math
import random
import numpy as np

import logging

class Subregion():
    """ Subregion class.

    Each instance of Subregion contains
    the name, and the left and right
    boundaries of a subregion.

    Name and boundaries are set during
    initualization and can be retrieved
    afterwards using helper methods.
    """
    def __init__(self, name, left_boundary, right_boundary, iterations):
        self.name = name
        self.left_boundary = left_boundary
        self.right_boundary = right_boundary
        self.flux_list = np.zeros(iterations, dtype=np.int64)

    def leftBoundary(self):
        return self.left_boundary

    def rightBoundary(self):
        return self.right_boundary

    def getName(self):
        return self.name

    def incrementFlux(self, iteration):
        self.flux_list[iteration] += 1

    def getFluxList(self):
        return self.flux_list

def get_event_particles(e_events, subregions, model_particles, level_limit, height_dependant=False):
    """ Find and return list of particles to be entrained

    Keyword arguments:
    e_events -- number of events requested per subregion
    subregions -- array of Subregion objects
    model_particles -- array of all model particles

    Returns:
    event_particles -- List of particles to be entrained

    """
    if e_events == 0:
        e_events = 1 #???

    event_particles = []
    for subregion in subregions:
        # Take only particles in the boundaries of the current subregion
        subregion_particles = model_particles[
                (model_particles[:,0] >= subregion.leftBoundary())
              & (model_particles[:,0] <= subregion.rightBoundary())]
        # Take only particles that are in-stream (not ghost)
        in_stream_particles = subregion_particles[
                                                subregion_particles[:,0] != -1]
        # Take only particles that are 'active'
        active_particles =  in_stream_particles[
                                                in_stream_particles[:,4] != 0]

        # Do not take any particles that have been selected for entrainment (i.e do not double select)
        # This only happens when particles rest on the boundary.
        active_event, active_idx, event_idx = np.intersect1d(active_particles[:,3], event_particles, return_indices=True)
        active_particles = np.delete(active_particles, active_idx, axis=0)

        subregion_event_ids = []
        if height_dependant: # any particle at the level limit must be entrained
            levels = elevation_list(subregion_particles[:,2], desc=False)
            tip_particles = []
            # find the tip particles -- these are the particles being entrained
            if len(levels) == level_limit:
                tip_particles = active_particles[active_particles[:,2] == levels[level_limit-1]]
            for particle in tip_particles:
                subregion_event_ids.append(particle[3])
                active_particles = active_particles[active_particles[:,2] != particle[2]]
        # If there are not enough particles in the subregion to sample from, alter the sample size
        if e_events > len(active_particles):
            random_sample = random.sample(range(len(active_particles)),
                                        len(active_particles))
        else:
            random_sample = random.sample(range(len(active_particles)),
                                        e_events)
        # TODO: change so that we don't rely on loop index to grab particle
        for index in random_sample:
            subregion_event_ids.append(int(active_particles[index][3])  )

        ghost_particles = np.where(model_particles[:,0] == -1)[0]
        for index in ghost_particles:
            model_particles[index][0] = 0
            subregion_event_ids.append(index)

        if e_events != len(subregion_event_ids):
            msg = (
                     f'Requested {e_events} events in {subregion.getName()} '
                     f'but {len(subregion_event_ids)} are occuring'
            )
            logging.warning(msg)
        event_particles = event_particles + subregion_event_ids
    event_particles = np.array(event_particles, dtype=np.intp)

    return event_particles

def define_subregions(bed_length, num_subregions, iterations):
    """ Define subregion list for model stream.

    Keyword arguments:
    bed_length -- length of the model bed.
    num_subregions -- number of subregions to create
    iterations -- number of iterations for this model simulation

    Returns:
    subregions_arr -- array of initialized subregion objects

    """
    subregion_length = bed_length/num_subregions
    left_boundary = 0.0
    subregions_arr = []
    for region in range(num_subregions):
        right_boundary = left_boundary + subregion_length
        subregion = Subregion(f'subregion-{region}', left_boundary, right_boundary, iterations)
        left_boundary = right_boundary

        subregions_arr.append(subregion)

    return subregions_arr

def build_streambed(x_max, set_diam):
    """ Build the bed particle list.

    Handles calls to add_bed_particle, checks for
    completness of bed and updates the x-extent
    of stream when the packing exceeds/under packs
    within 8mm range.

    Note: the updates to x-extent are only required
    when variable particle diameter is being used.

    Return values:
    bed_particles -- array of all bed particles

    """
    max_particles = int(math.ceil( x_max / set_diam ))
    bed_particles = np.zeros([max_particles, 7],dtype=float)

    particle_id = -1
    centre = (set_diam/2)
    state = 0
    age = 0
    loop_age = 0
    elevation = 0
    while not bed_complete(centre, x_max):
        # index with negative indices... bed particles are built from the final element to the first
        bed_particles[particle_id] = [centre, set_diam, elevation, particle_id, state, age, loop_age]
        centre += set_diam
        particle_id += -1 # Bed particles get negative IDs

    return bed_particles

def bed_complete(pack_idx, x_max):
    """Check to see if bed is complete based on model params."""
    # similarly, if np.count_nonzero(bed_space) == x_max
    if pack_idx >= x_max:
        return 1
    else: return 0

def determine_num_particles(pack_frac, num_vertices):
    """Return the number of model particles to be used, based on
    the packing fraction"""

    num_particles = num_vertices * pack_frac
    num_particles = int(math.ceil(num_particles))

    return num_particles

# Trig from: https://math.stackexchange.com/questions/2293201/
def place_particle(particle, model_particles, bed_particles, h):
    """ Calculate new Y of particle based on location in stream.

    Provided a particle p's location X in stream,
    search for 2 supporting particles that p will rest
    on when placed at X.

    Calculate the Y position of p based on the elevation
    and horizontal placement of both supporting particles.
    The calcuated X might be different up to some decimal
    point, so both X and Y are rounded to 2 decimal places.

    Keyword arguments:
    particle -- array representing the model particle that is being placed
    model_particles -- array of all model particles
    bed_particles -- array of all bed particles

    Return values:
    rounded_x -- rounded float of particle's new x location
    rounded_y -- rounded float of particle's new y location
    left_support -- id of the left support for the placed particle
    right_support -- id of the right support for the placed particle

    """
    left_support, right_support = find_supports(particle, model_particles, bed_particles)
    rounded_x = round(particle[0], 2)
    rounded_y = round(np.add(h, left_support[2]), 2)
    return rounded_x, rounded_y, left_support[3], right_support[3]

def update_particle_states(model_particles, model_supports):
    """ Set/update each model particle's state.

    If any model particle p has a particle
    resting on it in the stream then p must
    be set to inactive indicated by a boolean 0.

    If p does not have any particles resting
    on top of it then it is considered active
    indicated by a boolean 1.

    Note that bed particles are always considered
    inactive.

    Keyword arguments:
    model_particles -- array of all model particles
    model_supports -- array of left/right supporting particles
                                        for each model particle
    bed_particles -- array of all bed particles

    Return values:
    model_particles -- array of all model particles with updated states

    """
    # Start by setting all model particles to active then
    # only set to inactive if there is a particle sitting on top
    model_particles[:,4] = 1
    in_stream_particles = model_particles[model_particles[:,0] != -1]
    inactive_left = np.intersect1d(in_stream_particles[:,3], model_supports[:,0])
    inactive_right = np.intersect1d(in_stream_particles[:,3], model_supports[:,1])

    if inactive_left.size != 0:
        model_particles[inactive_left.astype(int), 4] = 0
    if inactive_right.size != 0:
        model_particles[inactive_right.astype(int), 4] = 0

    return model_particles

def find_supports(particle, model_particles, bed_particles):
    """ Find the 2 supporting particles for a given particle.

    Provided a particle p at location x, this function
    will search the stream for particles that p would
    rest on, if being dropped at x.

    More generally, supporting particles are those
    particles which directly hold up a particle. Supporting
    particles will always have a centre location that is
    exactly a radius length away from p's centre, since
    every particle has the same diameter.

    Keyword arguments:
    particle -- array representing a particle
    model_particles -- model particle list
    bed_particles -- bed particle list

    Returns:
    left_support -- the left supporting particle
    right_support -- the right supporting particle
    """
    all_particles = np.concatenate((model_particles, bed_particles), axis=0)
    # Define location where left and right supporting particles must sit
    # in order to be considered a supporting particle.
    # Note: This limits the model to using same-sized grains.
    left_center = particle[0] - (particle[1] / 2)
    right_center = particle[0] + (particle[1] / 2)

    l_candidates = all_particles[all_particles[:,0] == left_center]
    try:
        left_support = l_candidates[l_candidates[:,2]
                                    == np.max(l_candidates[:,2])]
    except ValueError:
        error_msg = (
                     f'No left supporting particle at {left_center}'
                     f'for particle {particle[3]}'
        )
        logging.error(error_msg)
        raise

    r_candidates = all_particles[all_particles[:,0] == right_center]
    try:
        right_support = r_candidates[r_candidates[:,2] == np.max(r_candidates[:,2])]
    except ValueError:
        error_msg = (
                     f'No right supporting particle at {right_center}'
                     f'for particle {particle[3]}'
        )
        logging.error(error_msg)
        raise
    return left_support[0], right_support[0]

def set_model_particles(bed_particles, available_vertices, set_diam, pack_fraction, h):
    """ Create array of model particles set each particle in stream.

    Locations of model particles are randomly
    assigned based on the vertices that are available
    from the bare bed.

    As a reminder the structure of a model particle is:
        [0] = horizontal location (centre),
        [1] = diameter,
        [2] = elevation (centre),
        [3] = uid,
        [4] = active (boolean)
        [5] = age counter
        [6] = loop age counter
    This should be the exact same as a bed particle.

    Keyword arguments:
    bed_particles -- array of all bed particles
    available_vertices -- array of all available vertices in the stream
    set_diam -- diameter of all particles
    pack_fraction -- float packing density value
    h --

    Return values:
    model_partilces -- array of all model particles
    model_supp -- array of left/right supporting particle ids
                                        for each model particle

     """
    num_placement_loc = np.size(available_vertices)
    # determine the number of model particles that should be introduced into the stream bed
    num_particles = determine_num_particles(pack_fraction, num_placement_loc)
    # create an empty n-6 array to store model particle information
    model_particles = np.zeros([num_particles, 7], dtype='float')
    model_supp = np.zeros([num_particles, 2], dtype='float')

    for particle in range(num_particles):
        # the following lines select a vertex to place the current particle at,
        # and ensure that it is not already occupied by another particle
        random_idx = random.randint(0, np.size(available_vertices)-1)
        vertex = available_vertices[random_idx]
        available_vertices = available_vertices[available_vertices != vertex]

        # intialize the particle information
        model_particles[particle][0] = vertex
        model_particles[particle][1] = set_diam

        model_particles[particle][3] = particle # id number for each particle
        model_particles[particle][4] = 1 # each particle begins as active

        # place particle at the chosen vertex
        p_x, p_y, left_supp, right_supp  = place_particle(model_particles[particle],
                                                            model_particles,
                                                            bed_particles,
                                                            h)
        model_particles[particle][0] = p_x
        model_particles[particle][2] = p_y
        model_particles[particle][5] = 0
        model_particles[particle][6] = 0

        model_supp[particle][0] = left_supp
        model_supp[particle][1] = right_supp

    return model_particles, model_supp

def compute_available_vertices(model_particles, bed_particles, set_diam, level_limit,
                               lifted_particles=None):
    """ Compute the avaliable vertices in the model stream.

    An available vertex is an x location that a
    model particle is able to be entrained to.
    This function identifies the distinct elevations
    present in the stream then looks at subsets of
    particles in decesnding order of their elevation,
    in order to compute available vertices.

    For each elevation group, if a particle is sitting on a vertex
    x, then x cannot be available (it is occupied) and it
    is considered nulled. Then vertices v created by two
    particles touching are considered. If v is not already
    nulled by a particle occupying the location at a higher
    level, then it is considered an available vertex.

    This ends once the bed particles (the lowest elevation)
    have been considered.

    Keyword arguments:
        model_particles -- list of model particles
        bed_particles -- list of bed particles
        lifted_particles  -- idx of the 'lifted' particles. Default None

    Returns:
        available_vertices -- the set of available vertices
    """
    nulled_vertices = []
    avail_vertices = []

    # If we are lifting particles, we need to consider the subset of particles
    # that includes every particles _except_ the particles being
    if lifted_particles is not None:
        model_particles_lifted = np.delete(model_particles,
                                           lifted_particles, 0)
        all_particles = np.concatenate((model_particles_lifted,
                                        bed_particles), axis=0)
    else:
        all_particles = np.concatenate((model_particles,
                                        bed_particles), axis=0)
    # Get unique model particle elevations in stream (descending)
    elevations = elevation_list(all_particles[:,2])

    for idx, elevation in enumerate(elevations):
        tmp_particles = all_particles[all_particles[:,2] == elevation]

        for particle in tmp_particles:
            nulled_vertices.append(particle[0])

        right_vertices = tmp_particles[:,0] + (set_diam / 2)
        left_vertices = tmp_particles[:,0] - (set_diam / 2)
        tmp_shared_vertices = np.intersect1d(left_vertices, right_vertices)

        # Enforce level limit by nulling any vertex above limit:
        if len(elevations) == level_limit+1 and idx==0:
            for vertex in tmp_shared_vertices:
                nulled_vertices.append(vertex)

        for vertex in tmp_shared_vertices:
            if vertex not in nulled_vertices:
                avail_vertices.append(vertex)

        del(tmp_shared_vertices)
        del(tmp_particles)

    available_vertices = np.array(avail_vertices)

    return available_vertices

def elevation_list(elevations, desc=True):
    """ Return a sorted list of unique elevation values """
    ue = np.unique(elevations)
    if desc:
           ue = ue[::-1]
    return ue

def compute_hops(event_particle_ids, model_particles, mu, sigma, normal=False):
    """ Given a list of (event) paritcles, this function will
    add a hop distance to current x locations of all event particles.

    Current + hop = desired hop distance.

    Hop distances are randomly selected from a log-normal or normal
    distribution.

    Keyword arguments:
        event_particle_ids -- list of event particle ids
        model_particles -- the model's np arry of model_particles
        normal -- boolean flag for sampling from Normal (default Flase)

    Returns:
        event_particles -- list of event particles with 'hopped' x-locations

    """
    event_particles = model_particles[event_particle_ids]
    if normal:
        s = np.random.normal(mu, sigma, len(event_particle_ids))
    else:
        s = np.random.lognormal(mu, sigma, len(event_particle_ids))
    s_hop = np.round(s, 1)
    s_hop = list(s_hop)
    event_particles[:,0] = event_particles[:,0] + s_hop

    return event_particles

def move_model_particles(event_particles, model_particles, model_supp, bed_particles, available_vertices, h):
    """ Given an array of event particles and their desired hops, move each
    event particle to the closest valid vertex if its desired hop is not a vertex.
    Update the model particle and support arrays accordingly.

    Keyword arguments:
        event_particles -- array of particles (full 1-7 struct) to be entrained
        model_particles -- array of all model particles
        model_supp -- array of left/right supporting particles for each
                                                            model particle
        bed_particles -- array of all bed particles
        available_particles -- array of available vertices in the stream

    Returns:
        model_particles -- array of model particles with event particle updates
        model_supports -- array of left/right supporting particles for each
                                    model particle, with event particles updated

    """
    # Randomly iterate over event particles
    for particle in np.random.permutation(event_particles):
        orig_x = model_particles[model_particles[:,3] == particle[3]][0][0]
        verified_hop = find_closest_vertex(particle[0], available_vertices)

        if verified_hop == -1:
            exceed_msg = (
                f'Particle {int(particle[3])} exceeded stream...'
                f'sending to -1 axis'
            )
            logging.info(exceed_msg)
            particle[6] = particle[6] + 1
            particle[0] = verified_hop

            model_supp[int(particle[3])][0] = np.nan
            model_supp[int(particle[3])][1] = np.nan
        else:
            hop_msg = (
                f'Particle {int(particle[3])} entrained from {orig_x} '
                f'to {verified_hop}. Desired hop was: {particle[0]}'
            )
            logging.info(hop_msg)
            particle[0] = verified_hop
            available_vertices = available_vertices[available_vertices != verified_hop]

            placed_x, placed_y, left_supp, right_supp = place_particle(particle, model_particles, bed_particles, h)
            particle[0] = placed_x
            particle[2] = placed_y

            model_supp[int(particle[3])][0] = left_supp
            model_supp[int(particle[3])][1] = right_supp

        model_particles[model_particles[:,3] == particle[3]] = particle
    return model_particles, model_supp

def update_flux(initial_positions, final_positions, iteration, subregions):
    """ Given arrays of initial and final positions, this function
    will update each subregion s to indicate how many paticles crossed
    s's downstream boundary.

    Keyword arguments:
    initial_positions -- array of initial x locations
    final_positions -- array of final (verified) x locations
    iteration -- the iteration the that the crossing should be recorded under
    subregions -- array of subregion objects

    Return values:
    subregions -- array of subregion objects with updated flux values
    """
    # This can _most definitely_ be made quicker but for now, it works
    if len(initial_positions) != len(final_positions):
        raise ValueError(f'Initial_positions and final_positions do not contain the same # of elements')

    for position in range(0, len(initial_positions)):

        initial_pos = initial_positions[position]
        final_pos = final_positions[position]

        for idx, subregion in enumerate(subregions):
            if (initial_pos >= subregion.leftBoundary()) and (subregion.rightBoundary() > initial_pos):
                start_idx = idx

        for subregion_idx in range(start_idx, len(subregions)):
            if final_pos >= subregions[subregion_idx].rightBoundary():
                subregions[subregion_idx].incrementFlux(iteration)
            elif final_pos == -1 and subregion_idx == len(subregions)-1:
                subregions[subregion_idx].incrementFlux(iteration)

    return subregions

def find_closest_vertex(desired_hop, available_vertices):
    """ Find the closest downstream (greater than or equal) vertex
    in availbale vertices. If nothing exists, then return -1.

    Keyword arguments:
    desired_hop -- float representing the desired hop location
    available_location -- np array of available vertices in model

    Returns:
    vertex -- the closest available vertex that is >= desired_hop
    """
    if available_vertices.size == 0:
        raise ValueError('Available vertices array is empty, cannot find closest vertex')
    if desired_hop < 0:
        raise ValueError('Desired hop is negative (invalid)')

    available_vertices = np.sort(available_vertices)
    forward_vertices = available_vertices[available_vertices >= desired_hop]

    if forward_vertices.size < 1:
        vertex = -1
    else:
        vertex = forward_vertices[0]
    return vertex

def increment_age(model_particles, e_event_ids):
    """"Increment model particles' age, set event particles to age 0"""

    model_particles[:,5] = model_particles[:,5] + 1
    model_particles[e_event_ids, 5] = 0

    return model_particles


text:
import time, pdb, argparse, subprocess, pickle, os, gzip, glob
import torch
import numpy as np
import time, pdb, argparse, subprocess, os, math, glob
import cv2
import python_speech_features
from tqdm import tqdm
import timeit

from scipy import signal
from scipy.io import wavfile
from SyncNetModel import *
from shutil import rmtree

import scenedetect
from scenedetect.video_manager import VideoManager
from scenedetect.scene_manager import SceneManager
from scenedetect.frame_timecode import FrameTimecode
from scenedetect.stats_manager import StatsManager
from scenedetect.detectors import ContentDetector

import multiprocessing as mp
from concurrent.futures import ThreadPoolExecutor, as_completed

from scipy.interpolate import interp1d
from scipy.io import wavfile
from scipy import signal

from detectors import S3FD

# ========== ========== ========== ==========
# # PARSE ARGS
# ========== ========== ========== ==========

parser = argparse.ArgumentParser(description = "SyncNet");
parser.add_argument('--syncnet_model', type=str, required=True);
parser.add_argument('--batch_size', type=int, required=True, help='');
parser.add_argument('--vshift', type=int, default='15', help='');
#parser.add_argument('--data_dir', type=str, default='data/work', help='');
#parser.add_argument('--videofile', type=str, required=True);
parser.add_argument('--data_path', type=str, required=True);
parser.add_argument('--tmp_dir', type=str, required=True);
parser.add_argument('--facedet_scale',  type=float, default=0.25, help='Scale factor for face detection');
parser.add_argument('--min_track', type=int, required=True,  help='Minimum facetrack duration');
parser.add_argument('--num_failed_det', type=int, default=25,   help='Number of missed detections allowed before tracking is stopped');
parser.add_argument('--min_face_size',  type=int, required=True,  help='Minimum face size in pixels');
parser.add_argument('--frame_rate',     type=int, default=25,   help='Frame rate');
parser.add_argument('--crop_scale',     type=float, default=0.40, help='Scale bounding box');
parser.add_argument('--verbose',    default=False, help='');
parser.add_argument('--nthreads', type=int,default=1);
parser.add_argument('--sample_duration',type=float, required=True)
parser.add_argument('--parts_num', type=int, required=True)
parser.add_argument('--part', type=int, required=True)
parser.add_argument('--output_path', type=str, required=True);

opt = parser.parse_args();

setattr(opt,'avi_dir',os.path.join(opt.tmp_dir,'pyavi'))
#setattr(opt,'pytmp_dir',os.path.join(opt.tmp_dir,'pytmp'))
setattr(opt,'work_dir',os.path.join(opt.tmp_dir,'pywork'))
setattr(opt,'crop_dir',os.path.join(opt.tmp_dir,'pycrop'))
setattr(opt,'frames_dir',os.path.join(opt.tmp_dir,'pyframes'))

# ========== ========== ========== ==========
# # SyncnetInstanceAnalyze Class
# ========== ========== ========== ==========

class SyncNetInstanceAnalyze(torch.nn.Module):

    def __init__(self, dropout = 0, num_layers_in_fc_layers = 1024):
        super(SyncNetInstanceAnalyze, self).__init__();

        self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers).cuda();

    def evaluate(self, opt,reference, videofile):

        self.__S__.eval();

        # ==================== Convert Files ====================

        if os.path.exists(os.path.join(opt.tmp_dir,reference)):
          rmtree(os.path.join(opt.tmp_dir,reference))

        os.makedirs(os.path.join(opt.tmp_dir,reference))

        command = ("ffmpeg -loglevel error -y -i %s -threads 1 -f image2 %s" % (videofile,os.path.join(opt.tmp_dir,reference,'%06d.jpg')))
        output = subprocess.call(command, shell=True, stdout=None)

        #command = ("ffmpeg -loglevel error -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'audio.wav')))
        #output = subprocess.call(command, shell=True, stdout=None)

        # ==================== Load Videos ====================

        images = []

        flist = glob.glob(os.path.join(opt.tmp_dir,reference,'*.jpg'))
        flist.sort()

        for fname in flist:
            images.append(cv2.imread(fname))

        im = np.stack(images,axis=3)
        im = np.expand_dims(im,axis=0)
        im = np.transpose(im,(0,3,4,1,2))

        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())

        # ==================== Load Audios ====================

        sample_rate, audio = wavfile.read(os.path.join(opt.avi_dir,reference,'audio.wav'))
        mfcc = zip(*python_speech_features.mfcc(audio,sample_rate))
        mfcc = np.stack([np.array(i) for i in mfcc])

        cc = np.expand_dims(np.expand_dims(mfcc,axis=0),axis=0)
        cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float())

        # ==================== Check Audio and Video input length ====================

        #if (float(len(audio))/16000) != (float(len(images))/25) :
            #print("WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."%(float(len(audio))/16000,float(len(images))/25))

        min_length = min(len(images),math.floor(len(audio)/640))

        # ==================== Generate Video and Audio feats ====================

        lastframe = min_length-5
        im_feat = []
        cc_feat = []

        tS = time.time()
        for i in range(0,lastframe,opt.batch_size):

            im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
            im_in = torch.cat(im_batch,0)
            im_out  = self.__S__.forward_lip(im_in.cuda());
            im_feat.append(im_out.data.cpu())

            cc_batch = [ cct[:,:,:,vframe*4:vframe*4+20] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
            cc_in = torch.cat(cc_batch,0)
            cc_out  = self.__S__.forward_aud(cc_in.cuda())
            cc_feat.append(cc_out.data.cpu())

        im_feat = torch.cat(im_feat,0)
        cc_feat = torch.cat(cc_feat,0)

        # ==================== Compute Offset ====================

        #print('Compute time %.3f sec.' % (time.time()-tS))

        dists = calc_pdist(im_feat,cc_feat,vshift=opt.vshift)
        mdist = torch.mean(torch.stack(dists,1),1)

        minval, minidx = torch.min(mdist,0)

        offset = opt.vshift-minidx
        conf   = torch.median(mdist) - minval

        fdist   = np.stack([dist[minidx].numpy() for dist in dists])
        # fdist   = numpy.pad(fdist, (3,3), 'constant', constant_values=15)
        fconf   = torch.median(mdist).numpy() - fdist
        fconfm  = signal.medfilt(fconf,kernel_size=9)

        #np.set_printoptions(formatter={'float': '{: 0.3f}'.format})
        #print('Framewise conf: ')
        #print(fconfm)
        #print('AV offset: \t%d \nMin dist: \t%.3f\nConfidence: \t%.3f' % (offset,minval,conf))

        dists_npy = np.array([ dist.numpy() for dist in dists ])
        #return offset.numpy(), conf.numpy(), dists_npy
        return offset.numpy(), conf.numpy(), minval.numpy()

    def extract_feature(self, opt, videofile):

        self.__S__.eval();

        # ==================== Load videos ====================
        cap = cv2.VideoCapture(videofile)

        frame_num = 1;
        images = []
        while frame_num:
            frame_num += 1
            ret, image = cap.read()
            if ret == 0:
                break

            images.append(image)

        im = np.stack(images,axis=3)
        im = np.expand_dims(im,axis=0)
        im = np.transpose(im,(0,3,4,1,2))

        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())

        # ==================== Generate Video Feats ====================

        lastframe = len(images)-4
        im_feat = []

        tS = time.time()
        for i in range(0,lastframe,opt.batch_size):

            im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ]
            im_in = torch.cat(im_batch,0)
            im_out  = self.__S__.forward_lipfeat(im_in.cuda());
            im_feat.append(im_out.data.cpu())

        im_feat = torch.cat(im_feat,0)

            # ==================== Return video feat ====================

        #print('Compute time %.3f sec.' % (time.time()-tS))

        return im_feat

    def loadParameters(self, path):
        loaded_state = torch.load(path, map_location=lambda storage, loc: storage);

        self_state = self.__S__.state_dict();

        for name, param in loaded_state.items():

            self_state[name].copy_(param);

# ========== ========== ========== ==========
# # FACE DETECTION
# ========== ========== ========== ==========

def inference_video(opt,reference,DET):

  flist = glob.glob(os.path.join(opt.frames_dir,reference,'*.jpg'))
  flist.sort()

  dets = []

  for fidx, fname in enumerate(flist):

    start_time = time.time()

    image = cv2.imread(fname)

    image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[opt.facedet_scale])

    dets.append([]);
    for bbox in bboxes:
      dets[-1].append({'frame':fidx, 'bbox':(bbox[:-1]).tolist(), 'conf':bbox[-1]})

    elapsed_time = time.time() - start_time

    #dprint('%s-%05d; %d dets; %.2f Hz' % (os.path.join(opt.avi_dir,opt.reference,'video.avi'),fidx,len(dets[-1]),(1/elapsed_time)))

  savepath = os.path.join(opt.work_dir,reference,'faces.pckl')

  with open(savepath, 'wb') as fil:
    pickle.dump(dets, fil)

  return dets

# ========== ========== ========== ==========
# # SCENE DETECTION
# ========== ========== ========== ==========

def scene_detect(opt, reference):
  video_manager = VideoManager([os.path.join(opt.avi_dir,reference,'video.avi')])
  stats_manager = StatsManager()
  scene_manager = SceneManager(stats_manager)
  # Add ContentDetector algorithm (constructor takes detector options like threshold).
  scene_manager.add_detector(ContentDetector())
  base_timecode = video_manager.get_base_timecode()

  video_manager.set_downscale_factor()

  video_manager.start()

  scene_manager.detect_scenes(frame_source=video_manager,show_progress=False)

  scene_list = scene_manager.get_scene_list(base_timecode)

  savepath = os.path.join(opt.work_dir,reference,'scene.pckl')

  if scene_list == []:
    scene_list = [(video_manager.get_base_timecode(),video_manager.get_current_timecode())]

  with open(savepath, 'wb') as fil:
    pickle.dump(scene_list, fil)
  #print('%s - scenes detected %d'%(os.path.join(opt.avi_dir,opt.reference,'video.avi'),len(scene_list)))

  return scene_list

# ========== ========== ========== ==========
# # FACE TRACKING
# ========== ========== ========== ==========

def track_shot(opt,scenefaces):

  iouThres  = 0.5     # Minimum IOU between consecutive face detections
  tracks    = []
  face_sizes = []

  while True:
    track     = []
    for framefaces in scenefaces:
      for face in framefaces:
        if track == []:
          track.append(face)
          framefaces.remove(face)
        elif face['frame'] - track[-1]['frame'] <= opt.num_failed_det:
          iou = bb_intersection_over_union(face['bbox'], track[-1]['bbox'])
          if iou > iouThres:
            track.append(face)
            framefaces.remove(face)
            continue
        else:
          break

    if track == []:
      break
    elif len(track) > opt.min_track:

      framenum    = np.array([ f['frame'] for f in track ])
      bboxes      = np.array([np.array(f['bbox']) for f in track])

      frame_i   = np.arange(framenum[0],framenum[-1]+1)

      bboxes_i    = []
      for ij in range(0,4):
        interpfn  = interp1d(framenum, bboxes[:,ij])
        bboxes_i.append(interpfn(frame_i))
      bboxes_i  = np.stack(bboxes_i, axis=1)

      if max(np.mean(bboxes_i[:,2]-bboxes_i[:,0]), np.mean(bboxes_i[:,3]-bboxes_i[:,1])) > opt.min_face_size:
        tracks.append({'frame':frame_i,'bbox':bboxes_i})
        face_sizes.append(max(np.mean(bboxes_i[:,2]-bboxes_i[:,0]), np.mean(bboxes_i[:,3]-bboxes_i[:,1])))
        #print("face size : {}".format(max(np.mean(bboxes_i[:,2]-bboxes_i[:,0]), np.mean(bboxes_i[:,3]-bboxes_i[:,1]))))

  return tracks, face_sizes

# ========== ========== ========== ==========
# # IOU FUNCTION
# ========== ========== ========== ==========

def bb_intersection_over_union(boxA, boxB):

  xA = max(boxA[0], boxB[0])
  yA = max(boxA[1], boxB[1])
  xB = min(boxA[2], boxB[2])
  yB = min(boxA[3], boxB[3])

  interArea = max(0, xB - xA) * max(0, yB - yA)

  boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
  boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])

  iou = interArea / float(boxAArea + boxBArea - interArea)

  return iou

# ========== ========== ========== ==========
# # VIDEO CROP AND SAVE
# ========== ========== ========== ==========

def crop_video(opt,reference,track,cropfile):

  flist = glob.glob(os.path.join(opt.frames_dir,reference,'*.jpg'))
  flist.sort()

  fourcc = cv2.VideoWriter_fourcc(*'XVID')
  vOut = cv2.VideoWriter(cropfile+'t.avi', fourcc, opt.frame_rate, (224,224))

  dets = {'x':[], 'y':[], 's':[]}

  for det in track['bbox']:

    dets['s'].append(max((det[3]-det[1]),(det[2]-det[0]))/2)
    dets['y'].append((det[1]+det[3])/2) # crop center x
    dets['x'].append((det[0]+det[2])/2) # crop center y

  # Smooth detections
  dets['s'] = signal.medfilt(dets['s'],kernel_size=13)
  dets['x'] = signal.medfilt(dets['x'],kernel_size=13)
  dets['y'] = signal.medfilt(dets['y'],kernel_size=13)

  for fidx, frame in enumerate(track['frame']):

    cs  = opt.crop_scale

    bs  = dets['s'][fidx]   # Detection box size
    bsi = int(bs*(1+2*cs))  # Pad videos by this amount

    image = cv2.imread(flist[frame])

    frame = np.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))
    my  = dets['y'][fidx]+bsi  # BBox center Y
    mx  = dets['x'][fidx]+bsi  # BBox center X

    face = frame[int(my-bs):int(my+bs*(1+2*cs)),int(mx-bs*(1+cs)):int(mx+bs*(1+cs))]

    vOut.write(cv2.resize(face,(224,224)))

  audiotmp    = os.path.join(opt.tmp_dir,reference,'audio.wav')
  audiostart  = (track['frame'][0])/opt.frame_rate
  audioend    = (track['frame'][-1]+1)/opt.frame_rate

  vOut.release()

# ========== ========== ========== ==========
# # Get OFFSET
# ========== ========== ========== ==========
def calc_pdist(feat1, feat2, vshift=10):

    win_size = vshift*2+1

    feat2p = torch.nn.functional.pad(feat2,(0,0,vshift,vshift))

    dists = []

    for i in range(0,len(feat1)):

        dists.append(torch.nn.functional.pairwise_distance(feat1[[i],:].repeat(win_size, 1), feat2p[i:i+win_size,:]))

    return dists

# ========== ========== ========== ==========
# # Get begining and end time for clip sample
# ========== ========== ========== ==========
def calc_beginingandend(videofile_path,requested_sample_duration):

    command=("ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 {}".format(videofile_path))
    video_duration = float((subprocess.check_output(command, shell=True)).decode('utf-8').strip())
    sample_duration=min(video_duration,requested_sample_duration)
    begining=video_duration/2-sample_duration/2
    end=video_duration/2+sample_duration/2
    #print("begining: {} ; end : {} ".format(begining,end))

    return begining,end

# ========== ========== ========== ==========
# # MP HANDLER
# ========== ========== ========== ==========

def mp_handler(job):
    try:
        videofile, opt, gpu_id = job
        if opt.verbose :
            start = timeit.default_timer()
        dir_path=os.path.dirname(videofile)
        videofile = os.path.basename(videofile)
        reference=videofile[:-4]

        # ========== DELETE EXISTING DIRECTORIES ==========

        if os.path.exists(os.path.join(opt.work_dir,reference)):
          rmtree(os.path.join(opt.work_dir,reference))

        if os.path.exists(os.path.join(opt.crop_dir,reference)):
          rmtree(os.path.join(opt.crop_dir,reference))

        if os.path.exists(os.path.join(opt.avi_dir,reference)):
          rmtree(os.path.join(opt.avi_dir,reference))

        if os.path.exists(os.path.join(opt.frames_dir,reference)):
          rmtree(os.path.join(opt.frames_dir,reference))

        if os.path.exists(os.path.join(opt.tmp_dir,reference)):
          rmtree(os.path.join(opt.tmp_dir,reference))

        # ========== MAKE NEW DIRECTORIES ==========

        os.makedirs(os.path.join(opt.work_dir,reference))
        os.makedirs(os.path.join(opt.crop_dir,reference))
        os.makedirs(os.path.join(opt.avi_dir,reference))
        os.makedirs(os.path.join(opt.frames_dir,reference))
        os.makedirs(os.path.join(opt.tmp_dir,reference))
        if opt.verbose :
            print('Time to delete and make directories: ', timeit.default_timer() - start)
            start = timeit.default_timer()

        # ========== CONVERT VIDEO AND EXTRACT FRAMES ==========
        begining,end = calc_beginingandend(os.path.join(dir_path,videofile),opt.sample_duration)

        command = ("ffmpeg -loglevel error -y -ss %f -i %s -t %f -qscale:v 2 -async 1 -r 25 %s" % (begining,os.path.join(dir_path,videofile),end-begining,os.path.join(opt.avi_dir,reference,'video.avi')))

        #command = ("ffmpeg -loglevel error -y -i %s -qscale:v 2 -async 1 -r 25 %s" % (os.path.join(dir_path,videofile),os.path.join(opt.avi_dir,reference,'video.avi')))
        output = subprocess.call(command, shell=True, stdout=None)

        command = ("ffmpeg -loglevel error -y -i %s -qscale:v 2 -threads 1 -f image2 %s" % (os.path.join(opt.avi_dir,reference,'video.avi'),os.path.join(opt.frames_dir,reference,'%06d.jpg')))
        output = subprocess.call(command, shell=True, stdout=None)

        command = ("ffmpeg -loglevel error -y -i %s -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % (os.path.join(opt.avi_dir,reference,'video.avi'),os.path.join(opt.avi_dir,reference,'audio.wav')))
        output = subprocess.call(command, shell=True, stdout=None)

        if opt.verbose :
            print('Convert video and extract frames: ', timeit.default_timer() - start)
            start = timeit.default_timer()

        # ========== FACE DETECTION ==========
        faces = inference_video(opt,reference,DET)
        if opt.verbose :
            print('Face Detection: ', timeit.default_timer() - start)
            start = timeit.default_timer()

        # ========== SCENE DETECTION ==========
        scene = scene_detect(opt,reference)
        if opt.verbose:
            print('Scene Detection: ', timeit.default_timer() - start)
            start = timeit.default_timer()

        # ========== FACE TRACKING ==========
        alltracks = []
        vidtracks = []
        face_sizes = []

        for shot in scene:

          if shot[1].frame_num - shot[0].frame_num >= opt.min_track :
            new_track, face_size=track_shot(opt,faces[shot[0].frame_num:shot[1].frame_num])
            #alltracks.extend(track_shot(opt,faces[shot[0].frame_num:shot[1].frame_num]))
            if len(face_size)>0:
                alltracks.extend(new_track)
                face_sizes.append(face_size)
        if opt.verbose:
            print('Face Tracking: ', timeit.default_timer() - start)
            start = timeit.default_timer()

        # ========== FACE TRACK CROP ==========

        for ii, track in enumerate(alltracks):
            vidtracks.append(crop_video(opt,reference,track,os.path.join(opt.crop_dir,reference,'%05d'%ii)))
        if opt.verbose:
            start = timeit.default_timer()
        # ========== SAVE RESULTS ==========

        savepath = os.path.join(opt.work_dir,reference,'tracks.pckl')

        with open(savepath, 'wb') as fil:
            pickle.dump(vidtracks, fil)
        rmtree(os.path.join(opt.tmp_dir,reference))

        if opt.verbose:
            print('Save Results: ', timeit.default_timer() - start)
            start = timeit.default_timer()
        # ==================== LOAD FILE LIST ====================

        flist = glob.glob(os.path.join(opt.crop_dir,reference,'0*.avi'))
        flist.sort()
        if opt.verbose:
            print('Load File List: ', timeit.default_timer() - start)
            start = timeit.default_timer()
        # ==================== GET OFFSETS ====================

        dists = []
    #for idx, fname in enumerate(flist):
    #    print("idx: {}".format(idx))
    #    offset, conf, dist = s.evaluate(opt,reference, videofile=fname)
    #    dists.append(dist)
    #    with open("output.txt", "a") as file_object:
    #        file_object.write("filename : {} ; offset : {} ; conf : {} ; mdist : {}; face_size : {}".format(videofile,offset,conf,dist,face_sizes[0][idx]))
    #        file_object.write("\n")
        if len(flist)==1 :
            fname = flist[0]
            offset, conf, dist = s.evaluate(opt,reference, videofile=fname)
            dists.append(dist)
            with open(opt.output_path, "a") as file_object:
                file_object.write("part : {} ; filename : {} ; offset : {} ; conf : {} ; mdist : {}; face_size : {}\n".format(opt.part,(dir_path+"/")+videofile,offset,conf,dist,face_sizes[0][0]))
        if opt.verbose:
            print('Get & print offset : ', timeit.default_timer() - start)

   # ========== DELETE USED DIRECTORIES ==========
        if os.path.exists(os.path.join(opt.work_dir,reference)):
          rmtree(os.path.join(opt.work_dir,reference))

        if os.path.exists(os.path.join(opt.crop_dir,reference)):
          rmtree(os.path.join(opt.crop_dir,reference))

        if os.path.exists(os.path.join(opt.avi_dir,reference)):
          rmtree(os.path.join(opt.avi_dir,reference))

        if os.path.exists(os.path.join(opt.frames_dir,reference)):
          rmtree(os.path.join(opt.frames_dir,reference))
    except:
        pass

# ========== ========== ========== ==========
# # MAIN
# ========== ========== ========== ==========

# ========== LOADING MANDATORY STUFF ==========
DET = S3FD(device='cuda')
s = SyncNetInstanceAnalyze();
s.loadParameters(opt.syncnet_model);
#print("Model %s loaded."%opt.syncnet_model);

# ========== FOR LOOP OVER ALL EXISTING MP4 FILES IN DATA_PATH ==========

videofiles = glob.glob(opt.data_path+'*/*.mp4')
videofiles = [v for i, v in enumerate(videofiles) if (i%opt.parts_num)==opt.part]

#print('Started processing for {} GPUs'.format(opt.nthreads))

jobs = [(vfile, opt, i%opt.nthreads) for i, vfile in enumerate(videofiles)]
if opt.nthreads==1:
    for j in tqdm(jobs):
        mp_handler(j)
else:
    p = ThreadPoolExecutor(opt.nthreads)
    futures = [p.submit(mp_handler, j) for j in jobs]
    _ = [r.result() for r in tqdm(as_completed(futures), total=len(futures))]

# ==================== PRINT RESULTS TO FILE ====================

#with open(os.path.join(opt.work_dir,opt.reference,'activesd.pckl'), 'wb') as fil:
#    pickle.dump(dists, fil)


text:
import os

import numpy as np
import scipy.io
from os.path import join

import torch
import torchvision
from matplotlib import pyplot as plt
from torchvision import transforms
from torchvision.datasets import VisionDataset
from torchvision.datasets.folder import default_loader
from torchvision.datasets.utils import download_url, list_dir
import tqdm

class Dogs(VisionDataset):
    download_url_prefix = 'http://vision.stanford.edu/aditya86/ImageNetDogs'

    def __init__(self, root, train=True, transform=None, target_transform=None, download=False, rand_number=0, imb_factor=1, imb_type='exp', new_class_idx_sorted=None):
        super(Dogs, self).__init__(root, transform=transform, target_transform=target_transform)
        np.random.seed(rand_number)

        self.loader = default_loader
        self.train = train

        self.root = os.path.join(self.root, 'dogs')
        if download:
            self.download()

        split = self.load_split()
        self.images_folder = join(self.root, 'Images')
        self.annotations_folder = join(self.root, 'Annotation')
        self._breeds = list_dir(self.images_folder)

        self.samples = [(annotation + '.jpg', idx) for annotation, idx in split]
        self.targets = []
        for anno, idx in self.samples:
            self.targets.append(int(idx))
        self.samples = np.array(self.samples)
        self.targets = np.array(self.targets)

        num_in_class = []
        for class_idx in np.unique(self.targets):
            num_in_class.append(len(np.where(self.targets == class_idx)[0]))
        self.num_in_class = num_in_class

        self.cls_num = len(self._breeds)
        # self.sort_dataset(new_class_idx_sorted)
        if train:
            img_num_list = self.get_img_num_per_cls(self.cls_num, imb_type, imb_factor)
            self.gen_imbalanced_data(img_num_list)

    def get_img_num_per_cls(self, cls_num, imb_type, imb_factor):
        img_max = len(self.samples) / cls_num
        img_num_per_cls = []
        if imb_type == 'exp':
            for cls_idx in range(cls_num):
                num = img_max * (imb_factor ** (cls_idx / (cls_num - 1.0)))
                img_num_per_cls.append(int(num))
        elif imb_type == 'step':
            for cls_idx in range(cls_num // 2):
                img_num_per_cls.append(int(img_max))
            for cls_idx in range(cls_num // 2):
                img_num_per_cls.append(int(img_max * imb_factor))
        else:
            img_num_per_cls.extend([int(img_max)] * cls_num)
        return img_num_per_cls

    def gen_imbalanced_data(self, img_num_per_cls):
        new_data = []
        new_targets = []
        targets_np = np.array(self.targets, dtype=np.int64)
        classes = np.unique(targets_np)
        # np.random.shuffle(classes)
        self.num_per_cls_dict = dict()
        self.samples = np.array(self.samples)
        for the_class, the_img_num in zip(classes, img_num_per_cls):
            self.num_per_cls_dict[the_class] = the_img_num
            idx = np.where(targets_np == the_class)[0]
            np.random.shuffle(idx)
            selec_idx = idx[:the_img_num]
            self.num_per_cls_dict[the_class] = len(selec_idx)
            new_data.append(self.samples[selec_idx])
            new_targets.extend([the_class, ] * the_img_num)
        new_data = np.vstack(new_data)
        self.samples = new_data
        self.targets = new_targets
        self.labels = new_targets

    def sort_dataset(self, new_class_idx_sorted=None):
        idx = np.argsort(self.targets)
        self.targets = self.targets[idx]
        self.samples = self.samples[idx]
        if new_class_idx_sorted is None:
            new_class_idx_sorted = np.argsort(self.num_in_class)[::-1]
        for idx, target in enumerate(self.targets):
            self.targets[idx] = np.where(new_class_idx_sorted == target)[0]
        idx = np.argsort(self.targets)
        self.targets = self.targets[idx]
        self.samples = self.samples[idx]
        for idx, sample in enumerate(self.samples):
            self.samples[idx][1] = self.targets[idx]
        self.new_class_idx_sorted = new_class_idx_sorted

    # def get_new_class_idx_sorted(self):
    #     return self.new_class_idx_sorted

    def get_cls_num_list(self):
        cls_num_list = []
        for i in range(self.cls_num):
            cls_num_list.append(self.num_per_cls_dict[i])
        return cls_num_list

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        image_name, target = self.samples[index]
        image_path = join(self.images_folder, image_name)
        image = self.loader(image_path)
        target = torch.tensor(int(target)).long()
        if self.transform is not None:
            if isinstance(self.transform, list):
                sample1 = self.transform[0](image)
                sample2 = self.transform[1](image)
                image = [sample1, sample2]
            else:
                image = self.transform(image)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return image, target, index

    def download(self):
        import tarfile

        if os.path.exists(join(self.root, 'Images')) and os.path.exists(join(self.root, 'Annotation')):
            if len(os.listdir(join(self.root, 'Images'))) == len(os.listdir(join(self.root, 'Annotation'))) == 120:
                print('Files already downloaded and verified')
                return

        for filename in ['images', 'annotation', 'lists']:
            tar_filename = filename + '.tar'
            url = self.download_url_prefix + '/' + tar_filename
            download_url(url, self.root, tar_filename, None)
            print('Extracting downloaded file: ' + join(self.root, tar_filename))
            with tarfile.open(join(self.root, tar_filename), 'r') as tar_file:
                tar_file.extractall(self.root)
            os.remove(join(self.root, tar_filename))

    def load_split(self):
        if self.train:
            split = scipy.io.loadmat(join(self.root, 'train_list.mat'))['annotation_list']
            labels = scipy.io.loadmat(join(self.root, 'train_list.mat'))['labels']
        else:
            split = scipy.io.loadmat(join(self.root, 'test_list.mat'))['annotation_list']
            labels = scipy.io.loadmat(join(self.root, 'test_list.mat'))['labels']

        split = [item[0][0] for item in split]
        labels = [item[0] - 1 for item in labels]
        return list(zip(split, labels))

    def stats(self):
        counts = {}
        for index in range(len(self.samples)):
            image_name, target_class = self.samples[index]
            if target_class not in counts.keys():
                counts[target_class] = 1
            else:
                counts[target_class] += 1

        print("%d samples spanning %d classes (avg %f per class)" % (len(self.samples), len(counts.keys()),
                                                                     float(len(self.samples)) / float(
                                                                         len(counts.keys()))))

        return counts

if __name__ == '__main__':

    train_transform = transforms.Compose([
        transforms.ToTensor(),
    ])

    test_dataset = Dogs('/data', train=False, download=False, transform=train_transform)
    # new_class_idx = test_dataset.get_new_class_idx_sorted()
    train_dataset = Dogs('/data', train=True, download=False, transform=train_transform, imb_factor=0.1)

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=1, shuffle=False,
        num_workers=0, persistent_workers=False, pin_memory=True)
    #
    # classes_freq = np.zeros(train_dataset.cls_num)
    # for x, y in tqdm.tqdm(train_loader):
    #     classes_freq[np.array(y)] += 1
    # print(classes_freq)

    test_loader = torch.utils.data.DataLoader(
        test_dataset, batch_size=1, shuffle=False,
        num_workers=0, persistent_workers=False, pin_memory=True)

    # for i in range(len(train_dataset.get_cls_num_list())):
    #     images = torch.empty(152, 3, 224, 224)
    #     idx = 0
    #     for image, y in train_loader:
    #         if y == i:
    #             images[idx] = image
    #             idx += 1
    #
    #     plt.figure(dpi=300)
    #     plt.title(f'{i}')
    #     plt.clf()
    #     plt.imshow(torchvision.utils.make_grid(images, normalize=True).permute(1, 2, 0))
    #     plt.savefig(f'Dogs_train{i}.png')
    #     if i > 9:
    #         break
    #
    # for i in range(len(train_dataset.get_cls_num_list())):
    #     images = torch.empty(152, 3, 224, 224)
    #     idx = 0
    #     for image, y in test_loader:
    #         if y == i:
    #             images[idx] = image
    #             idx += 1
    #
    #     plt.figure(dpi=300)
    #     plt.title(f'{i}')
    #     plt.clf()
    #     plt.imshow(torchvision.utils.make_grid(images, normalize=True).permute(1, 2, 0))
    #     plt.savefig(f'Dogs_val{i}.png')
    #     if i > 9:
    #         break

    # classes_freq = np.zeros(train_dataset.cls_num)
    # for x, y in tqdm.tqdm(test_loader):
    #     classes_freq[np.array(y)] += 1
    # print(classes_freq)
    # train_loader = torch.utils.data.DataLoader(
    #     train_dataset, batch_size=1, shuffle=False,
    #     num_workers=0, persistent_workers=False, pin_memory=True)
    #
    # print(train_dataset.get_cls_num_list())
    #
    mean = 0.
    std = 0.
    classes_freq = np.zeros(120)
    for images, y in train_loader:
        batch_samples = images.size(0)  # batch size (the last batch can have smaller size!)
        images = images.view(batch_samples, images.size(1), -1)
        mean += images.mean(2).sum(0)
        std += images.std(2).sum(0)
        classes_freq[np.array(y)] += 1
    mean /= len(train_loader.dataset)
    std /= len(train_loader.dataset)
    print(classes_freq)
    print(mean, std)

reason: import numpy,import scipy