netdissect/dissection.py

'''
To run dissection:

1. Load up the convolutional model you wish to dissect, and wrap it in
   an InstrumentedModel; then call imodel.retain_layers([layernames,..])
   to instrument the layers of interest.
2. Load the segmentation dataset using the BrodenDataset class;
   use the transform_image argument to normalize images to be
   suitable for the model, or the size argument to truncate the dataset.
3. Choose a directory in which to write the output, and call
   dissect(outdir, model, dataset).

Example:

    from dissect import InstrumentedModel, dissect
    from broden import BrodenDataset

    model = InstrumentedModel(load_my_model())
    model.eval()
    model.cuda()
    model.retain_layers(['conv1', 'conv2', 'conv3', 'conv4', 'conv5'])
    bds = BrodenDataset('dataset/broden1_227',
            transform_image=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(IMAGE_MEAN, IMAGE_STDEV)]),
            size=1000)
    dissect('result/dissect', model, bds,
            examples_per_unit=10)
'''

import torch, numpy, os, re, json, shutil, types, tempfile, torchvision
# import warnings
# warnings.simplefilter('error', UserWarning)
from PIL import Image
from xml.etree import ElementTree as et
from collections import OrderedDict, defaultdict
from .progress import verbose_progress, default_progress, print_progress
from .progress import desc_progress
from .runningstats import RunningQuantile, RunningTopK
from .runningstats import RunningCrossCovariance, RunningConditionalQuantile
from .sampler import FixedSubsetSampler
from .actviz import activation_visualization
from .segviz import segment_visualization, high_contrast
from .workerpool import WorkerBase, WorkerPool
from .segmenter import UnifiedParsingSegmenter

def dissect(outdir, model, dataset,
        segrunner=None,
        train_dataset=None,
        model_segmenter=None,
        quantile_threshold=0.005,
        iou_threshold=0.05,
        iqr_threshold=0.01,
        examples_per_unit=100,
        batch_size=100,
        num_workers=24,
        seg_batch_size=5,
        make_images=True,
        make_labels=True,
        make_maxiou=False,
        make_covariance=False,
        make_report=True,
        make_row_images=True,
        make_single_images=False,
        rank_all_labels=False,
        netname=None,
        meta=None,
        merge=None,
        settings=None,
        ):
    '''
    Runs net dissection in-memory, using pytorch, and saves visualizations
    and metadata into outdir.
    '''
    assert not model.training, 'Run model.eval() before dissection'
    if netname is None:
        netname = type(model).__name__
    if segrunner is None:
        segrunner = ClassifierSegRunner(dataset)
    if train_dataset is None:
        train_dataset = dataset
    make_iqr = (quantile_threshold == 'iqr')
    with torch.no_grad():
        device = next(model.parameters()).device
        levels = None
        labelnames, catnames = None, None
        maxioudata, iqrdata = None, None
        labeldata = None
        iqrdata, cov = None, None

        labelnames, catnames = segrunner.get_label_and_category_names()
        label_category = [catnames.index(c) if c in catnames else 0
                for l, c in labelnames]

        # First, always collect qunatiles and topk information.
        segloader = torch.utils.data.DataLoader(dataset,
                batch_size=batch_size, num_workers=num_workers,
                pin_memory=(device.type == 'cuda'))
        quantiles, topk = collect_quantiles_and_topk(outdir, model,
            segloader, segrunner, k=examples_per_unit)

        # Thresholds can be automatically chosen by maximizing iqr
        if make_iqr:
            # Get thresholds based on an IQR optimization
            segloader = torch.utils.data.DataLoader(train_dataset,
                    batch_size=1, num_workers=num_workers,
                    pin_memory=(device.type == 'cuda'))
            iqrdata = collect_iqr(outdir, model, segloader, segrunner)
            max_iqr, full_iqr_levels = iqrdata[:2]
            max_iqr_agreement = iqrdata[4]
            # qualified_iqr[max_iqr_quantile[layer] > 0.5] = 0
            levels = {layer: full_iqr_levels[layer][
                    max_iqr[layer].max(0)[1],
                    torch.arange(max_iqr[layer].shape[1])].to(device)
                    for layer in full_iqr_levels}
        else:
            levels = {k: qc.quantiles([1.0 - quantile_threshold])[:,0]
                      for k, qc in quantiles.items()}

        quantiledata = (topk, quantiles, levels, quantile_threshold)

        if make_images:
            segloader = torch.utils.data.DataLoader(dataset,
                    batch_size=batch_size, num_workers=num_workers,
                    pin_memory=(device.type == 'cuda'))
            generate_images(outdir, model, dataset, topk, levels, segrunner,
                    row_length=examples_per_unit, batch_size=seg_batch_size,
                    row_images=make_row_images,
                    single_images=make_single_images,
                    num_workers=num_workers)

        if make_maxiou:
            assert train_dataset, "Need training dataset for maxiou."
            segloader = torch.utils.data.DataLoader(train_dataset,
                    batch_size=1, num_workers=num_workers,
                    pin_memory=(device.type == 'cuda'))
            maxioudata = collect_maxiou(outdir, model, segloader,
                    segrunner)

        if make_labels:
            segloader = torch.utils.data.DataLoader(dataset,
                    batch_size=1, num_workers=num_workers,
                    pin_memory=(device.type == 'cuda'))
            iou_scores, iqr_scores, tcs, lcs, ccs, ics = (
                    collect_bincounts(outdir, model, segloader,
                    levels, segrunner))
            labeldata = (iou_scores, iqr_scores, lcs, ccs, ics, iou_threshold,
                    iqr_threshold)

        if make_covariance:
            segloader = torch.utils.data.DataLoader(dataset,
                    batch_size=seg_batch_size,
                    num_workers=num_workers,
                    pin_memory=(device.type == 'cuda'))
            cov = collect_covariance(outdir, model, segloader, segrunner)

        if make_report:
            generate_report(outdir,
                    quantiledata=quantiledata,
                    labelnames=labelnames,
                    catnames=catnames,
                    labeldata=labeldata,
                    maxioudata=maxioudata,
                    iqrdata=iqrdata,
                    covariancedata=cov,
                    rank_all_labels=rank_all_labels,
                    netname=netname,
                    meta=meta,
                    mergedata=merge,
                    settings=settings)

        return quantiledata, labeldata

def generate_report(outdir, quantiledata, labelnames=None, catnames=None,
        labeldata=None, maxioudata=None, iqrdata=None, covariancedata=None,
        rank_all_labels=False, netname='Model', meta=None, settings=None,
        mergedata=None):
    '''
    Creates dissection.json reports and summary bargraph.svg files in the
    specified output directory, and copies a dissection.html interface
    to go along with it.
    '''
    all_layers = []
    # Current source code directory, for html to copy.
    srcdir = os.path.realpath(
       os.path.join(os.getcwd(), os.path.dirname(__file__)))
    # Unpack arguments
    topk, quantiles, levels, quantile_threshold = quantiledata
    top_record = dict(
            netname=netname,
            meta=meta,
            default_ranking='unit',
            quantile_threshold=quantile_threshold)
    if settings is not None:
        top_record['settings'] = settings
    if labeldata is not None:
        iou_scores, iqr_scores, lcs, ccs, ics, iou_threshold, iqr_threshold = (
                labeldata)
        catorder = {'object': -7, 'scene': -6, 'part': -5,
                    'piece': -4,
                    'material': -3, 'texture': -2, 'color': -1}
        for i, cat in enumerate(c for c in catnames if c not in catorder):
            catorder[cat] = i
        catnumber = {n: i for i, n in enumerate(catnames)}
        catnumber['-'] = 0
        top_record['default_ranking'] = 'label'
        top_record['iou_threshold'] = iou_threshold
        top_record['iqr_threshold'] = iqr_threshold
        labelnumber = dict((name[0], num)
                for num, name in enumerate(labelnames))
    # Make a segmentation color dictionary
    segcolors = {}
    for i, name in enumerate(labelnames):
        key = ','.join(str(s) for s in high_contrast[i % len(high_contrast)])
        if key in segcolors:
            segcolors[key] += '/' + name[0]
        else:
            segcolors[key] = name[0]
    top_record['segcolors'] = segcolors
    for layer in topk.keys():
        units, rankings = [], []
        record = dict(layer=layer, units=units, rankings=rankings)
        # For every unit, we always have basic visualization information.
        topa, topi = topk[layer].result()
        lev = levels[layer]
        for u in range(len(topa)):
            units.append(dict(
                unit=u,
                interp=True,
                level=lev[u].item(),
                top=[dict(imgnum=i.item(), maxact=a.item())
                    for i, a in zip(topi[u], topa[u])],
                ))
        rankings.append(dict(name="unit", score=list([
            u for u in range(len(topa))])))
        # TODO: consider including stats and ranking based on quantiles,
        # variance, connectedness here.

        # if we have labeldata, then every unit also gets a bunch of other info
        if labeldata is not None:
            lscore, qscore, cc, ic = [dat[layer]
                    for dat in [iou_scores, iqr_scores, ccs, ics]]
            if iqrdata is not None:
                # If we have IQR thresholds, assign labels based on that
                max_iqr, max_iqr_level = iqrdata[:2]
                best_label = max_iqr[layer].max(0)[1]
                best_score = lscore[best_label, torch.arange(lscore.shape[1])]
                best_qscore = qscore[best_label, torch.arange(lscore.shape[1])]
            else:
                # Otherwise, assign labels based on max iou
                best_score, best_label = lscore.max(0)
                best_qscore = qscore[best_label, torch.arange(qscore.shape[1])]
            record['iou_threshold'] = iou_threshold,
            for u, urec in enumerate(units):
                score, qscore, label = (
                        best_score[u], best_qscore[u], best_label[u])
                urec.update(dict(
                    iou=score.item(),
                    iou_iqr=qscore.item(),
                    lc=lcs[label].item(),
                    cc=cc[catnumber[labelnames[label][1]], u].item(),
                    ic=ic[label, u].item(),
                    interp=(qscore.item() > iqr_threshold and
                        score.item() > iou_threshold),
                    iou_labelnum=label.item(),
                    iou_label=labelnames[label.item()][0],
                    iou_cat=labelnames[label.item()][1],
                    ))
        if maxioudata is not None:
            max_iou, max_iou_level, max_iou_quantile = maxioudata
            qualified_iou = max_iou[layer].clone()
            # qualified_iou[max_iou_quantile[layer] > 0.75] = 0
            best_score, best_label = qualified_iou.max(0)
            for u, urec in enumerate(units):
                urec.update(dict(
                    maxiou=best_score[u].item(),
                    maxiou_label=labelnames[best_label[u].item()][0],
                    maxiou_cat=labelnames[best_label[u].item()][1],
                    maxiou_level=max_iou_level[layer][best_label[u], u].item(),
                    maxiou_quantile=max_iou_quantile[layer][
                        best_label[u], u].item()))
        if iqrdata is not None:
            [max_iqr, max_iqr_level, max_iqr_quantile,
                    max_iqr_iou, max_iqr_agreement] = iqrdata
            qualified_iqr = max_iqr[layer].clone()
            qualified_iqr[max_iqr_quantile[layer] > 0.5] = 0
            best_score, best_label = qualified_iqr.max(0)
            for u, urec in enumerate(units):
                urec.update(dict(
                    iqr=best_score[u].item(),
                    iqr_label=labelnames[best_label[u].item()][0],
                    iqr_cat=labelnames[best_label[u].item()][1],
                    iqr_level=max_iqr_level[layer][best_label[u], u].item(),
                    iqr_quantile=max_iqr_quantile[layer][
                        best_label[u], u].item(),
                    iqr_iou=max_iqr_iou[layer][best_label[u], u].item()
                    ))
        if covariancedata is not None:
            score = covariancedata[layer].correlation()
            best_score, best_label = score.max(1)
            for u, urec in enumerate(units):
                urec.update(dict(
                    cor=best_score[u].item(),
                    cor_label=labelnames[best_label[u].item()][0],
                    cor_cat=labelnames[best_label[u].item()][1]
                    ))
        if mergedata is not None:
            # Final step: if the user passed any data to merge into the
            # units, merge them now.  This can be used, for example, to
            # indiate that a unit is not interpretable based on some
            # outside analysis of unit statistics.
            for lrec in mergedata.get('layers', []):
                if lrec['layer'] == layer:
                    break
            else:
                lrec = None
            for u, urec in enumerate(lrec.get('units', []) if lrec else []):
                units[u].update(urec)
        # After populating per-unit info, populate per-layer ranking info
        if labeldata is not None:
            # Collect all labeled units
            labelunits = defaultdict(list)
            all_labelunits = defaultdict(list)
            for u, urec in enumerate(units):
                if urec['interp']:
                    labelunits[urec['iou_labelnum']].append(u)
                all_labelunits[urec['iou_labelnum']].append(u)
            # Sort all units in order with most popular label first.
            label_ordering = sorted(units,
                # Sort by:
                key=lambda r: (-1 if r['interp'] else 0,  # interpretable
                    -len(labelunits[r['iou_labelnum']]),  # label freq, score
                    -max([units[u]['iou']
                        for u in labelunits[r['iou_labelnum']]], default=0),
                    r['iou_labelnum'],                    # label
                    -r['iou']))                           # unit score
            # Add label and iou ranking.
            rankings.append(dict(name="label", score=(numpy.argsort(list(
                ur['unit'] for ur in label_ordering))).tolist()))
            rankings.append(dict(name="max iou", metric="iou", score=list(
                -ur['iou'] for ur in units)))
            # Add ranking for top labels
            # for labelnum in [n for n in sorted(
            #     all_labelunits.keys(), key=lambda x:
            #         -len(all_labelunits[x])) if len(all_labelunits[n])]:
            #     label = labelnames[labelnum][0]
            #     rankings.append(dict(name="%s-iou" % label,
            #         concept=label, metric='iou',
            #         score=(-lscore[labelnum, :]).tolist()))
            # Collate labels by category then frequency.
            record['labels'] = [dict(
                        label=labelnames[label][0],
                        labelnum=label,
                        units=labelunits[label],
                        cat=labelnames[label][1])
                    for label in (sorted(labelunits.keys(),
                        # Sort by:
                        key=lambda l: (catorder.get(          # category
                            labelnames[l][1], 0),
                            -len(labelunits[l]),              # label freq
                            -max([units[u]['iou'] for u in labelunits[l]],
                                default=0) # score
                            ))) if len(labelunits[label])]
            # Total number of interpretable units.
            record['interpretable'] = sum(len(group['units'])
                    for group in record['labels'])
            # Make a bargraph of labels
            os.makedirs(os.path.join(outdir, safe_dir_name(layer)),
                    exist_ok=True)
            catgroups = OrderedDict()
            for _, cat in sorted([(v, k) for k, v in catorder.items()]):
                catgroups[cat] = []
            for rec in record['labels']:
                if rec['cat'] not in catgroups:
                    catgroups[rec['cat']] = []
                catgroups[rec['cat']].append(rec['label'])
            make_svg_bargraph(
                    [rec['label'] for rec in record['labels']],
                    [len(rec['units']) for rec in record['labels']],
                    [(cat, len(group)) for cat, group in catgroups.items()],
                    filename=os.path.join(outdir, safe_dir_name(layer),
                        'bargraph.svg'))
            # Only show the bargraph if it is non-empty.
            if len(record['labels']):
                record['bargraph'] = 'bargraph.svg'
        if maxioudata is not None:
            rankings.append(dict(name="max maxiou", metric="maxiou", score=list(
                    -ur['maxiou'] for ur in units)))
        if iqrdata is not None:
            rankings.append(dict(name="max iqr", metric="iqr", score=list(
                    -ur['iqr'] for ur in units)))
        if covariancedata is not None:
            rankings.append(dict(name="max cor", metric="cor", score=list(
                    -ur['cor'] for ur in units)))

        all_layers.append(record)
    # Now add the same rankings to every layer...
    all_labels = None
    if rank_all_labels:
        all_labels = [name for name, cat in labelnames]
    if labeldata is not None:
        # Count layers+quadrants with a given label, and sort by freq
        counted_labels = defaultdict(int)
        for label in [
                re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '', unitrec['iou_label'])
                for record in all_layers for unitrec in record['units']]:
            counted_labels[label] += 1
        if all_labels is None:
            all_labels = [label for count, label in sorted((-v, k)
                for k, v in counted_labels.items())]
        for record in all_layers:
            layer = record['layer']
            for label in all_labels:
                labelnum = labelnumber[label]
                record['rankings'].append(dict(name="%s-iou" % label,
                    concept=label, metric='iou',
                    score=(-iou_scores[layer][labelnum, :]).tolist()))

    if maxioudata is not None:
        if all_labels is None:
            counted_labels = defaultdict(int)
            for label in [
                    re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '',
                        unitrec['maxiou_label'])
                    for record in all_layers for unitrec in record['units']]:
                counted_labels[label] += 1
            all_labels = [label for count, label in sorted((-v, k)
                for k, v in counted_labels.items())]
        qualified_iou = max_iou[layer].clone()
        qualified_iou[max_iou_quantile[layer] > 0.5] = 0
        for record in all_layers:
            layer = record['layer']
            for label in all_labels:
                labelnum = labelnumber[label]
                record['rankings'].append(dict(name="%s-maxiou" % label,
                    concept=label, metric='maxiou',
                    score=(-qualified_iou[labelnum, :]).tolist()))

    if iqrdata is not None:
        if all_labels is None:
            counted_labels = defaultdict(int)
            for label in [
                    re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '',
                        unitrec['iqr_label'])
                    for record in all_layers for unitrec in record['units']]:
                counted_labels[label] += 1
            all_labels = [label for count, label in sorted((-v, k)
                for k, v in counted_labels.items())]
        # qualified_iqr[max_iqr_quantile[layer] > 0.5] = 0
        for record in all_layers:
            layer = record['layer']
            qualified_iqr = max_iqr[layer].clone()
            for label in all_labels:
                labelnum = labelnumber[label]
                record['rankings'].append(dict(name="%s-iqr" % label,
                    concept=label, metric='iqr',
                    score=(-qualified_iqr[labelnum, :]).tolist()))

    if covariancedata is not None:
        if all_labels is None:
            counted_labels = defaultdict(int)
            for label in [
                    re.sub(r'-(?:t|b|l|r|tl|tr|bl|br)$', '',
                        unitrec['cor_label'])
                    for record in all_layers for unitrec in record['units']]:
                counted_labels[label] += 1
            all_labels = [label for count, label in sorted((-v, k)
                for k, v in counted_labels.items())]
        for record in all_layers:
            layer = record['layer']
            score = covariancedata[layer].correlation()
            for label in all_labels:
                labelnum = labelnumber[label]
                record['rankings'].append(dict(name="%s-cor" % label,
                    concept=label, metric='cor',
                    score=(-score[:, labelnum]).tolist()))

    for record in all_layers:
        layer = record['layer']
        # Dump per-layer json inside per-layer directory
        record['dirname'] = '.'
        with open(os.path.join(outdir, safe_dir_name(layer), 'dissect.json'),
                'w') as jsonfile:
            top_record['layers'] = [record]
            json.dump(top_record, jsonfile, indent=1)
        # Copy the per-layer html
        shutil.copy(os.path.join(srcdir, 'dissect.html'),
                os.path.join(outdir, safe_dir_name(layer), 'dissect.html'))
        record['dirname'] = safe_dir_name(layer)

    # Dump all-layer json in parent directory
    with open(os.path.join(outdir, 'dissect.json'), 'w') as jsonfile:
        top_record['layers'] = all_layers
        json.dump(top_record, jsonfile, indent=1)
    # Copy the all-layer html
    shutil.copy(os.path.join(srcdir, 'dissect.html'),
            os.path.join(outdir, 'dissect.html'))
    shutil.copy(os.path.join(srcdir, 'edit.html'),
            os.path.join(outdir, 'edit.html'))


def generate_images(outdir, model, dataset, topk, levels,
        segrunner, row_length=None, gap_pixels=5,
        row_images=True, single_images=False, prefix='',
        batch_size=100, num_workers=24):
    '''
    Creates an image strip file for every unit of every retained layer
    of the model, in the format [outdir]/[layername]/[unitnum]-top.jpg.
    Assumes that the indexes of topk refer to the indexes of dataset.
    Limits each strip to the top row_length images.
    '''
    progress = default_progress()
    needed_images = {}
    if row_images is False:
        row_length = 1
    # Pass 1: needed_images lists all images that are topk for some unit.
    for layer in topk:
        topresult = topk[layer].result()[1].cpu()
        for unit, row in enumerate(topresult):
            for rank, imgnum in enumerate(row[:row_length]):
                imgnum = imgnum.item()
                if imgnum not in needed_images:
                    needed_images[imgnum] = []
                needed_images[imgnum].append((layer, unit, rank))
    levels = {k: v.cpu().numpy() for k, v in levels.items()}
    row_length = len(row[:row_length])
    needed_sample = FixedSubsetSampler(sorted(needed_images.keys()))
    device = next(model.parameters()).device
    segloader = torch.utils.data.DataLoader(dataset,
            batch_size=batch_size, num_workers=num_workers,
            pin_memory=(device.type == 'cuda'),
            sampler=needed_sample)
    vizgrid, maskgrid, origrid, seggrid = [{} for _ in range(4)]
    # Pass 2: populate vizgrid with visualizations of top units.
    pool = None
    for i, batch in enumerate(
            progress(segloader, desc='Making images')):
        # Reverse transformation to get the image in byte form.
        seg, _, byte_im, _ = segrunner.run_and_segment_batch(batch, model,
                want_rgb=True)
        torch_features = model.retained_features()
        scale_offset = getattr(model, 'scale_offset', None)
        if pool is None:
            # Distribute the work across processes: create shared mmaps.
            for layer, tf in torch_features.items():
                [vizgrid[layer], maskgrid[layer], origrid[layer],
                        seggrid[layer]] = [
                    create_temp_mmap_grid((tf.shape[1],
                        byte_im.shape[1], row_length,
                        byte_im.shape[2] + gap_pixels, depth),
                        dtype='uint8',
                        fill=255)
                    for depth in [3, 4, 3, 3]]
            # Pass those mmaps to worker processes.
            pool = WorkerPool(worker=VisualizeImageWorker,
                    memmap_grid_info=[
                        {layer: (g.filename, g.shape, g.dtype)
                            for layer, g in grid.items()}
                        for grid in [vizgrid, maskgrid, origrid, seggrid]])
        byte_im = byte_im.cpu().numpy()
        numpy_seg = seg.cpu().numpy()
        features = {}
        for index in range(len(byte_im)):
            imgnum = needed_sample.samples[index + i*segloader.batch_size]
            for layer, unit, rank in needed_images[imgnum]:
                if layer not in features:
                    features[layer] = torch_features[layer].cpu().numpy()
                pool.add(layer, unit, rank,
                        byte_im[index],
                        features[layer][index, unit],
                        levels[layer][unit],
                        scale_offset[layer] if scale_offset else None,
                        numpy_seg[index])
    pool.join()
    # Pass 3: save image strips as [outdir]/[layer]/[unitnum]-[top/orig].jpg
    pool = WorkerPool(worker=SaveImageWorker)
    for layer, vg in progress(vizgrid.items(), desc='Saving images'):
        os.makedirs(os.path.join(outdir, safe_dir_name(layer),
            prefix + 'image'), exist_ok=True)
        if single_images:
           os.makedirs(os.path.join(outdir, safe_dir_name(layer),
               prefix + 's-image'), exist_ok=True)
        og, sg, mg = origrid[layer], seggrid[layer], maskgrid[layer]
        for unit in progress(range(len(vg)), desc='Units'):
            for suffix, grid in [('top.jpg', vg), ('orig.jpg', og),
                    ('seg.png', sg), ('mask.png', mg)]:
                strip = grid[unit].reshape(
                        (grid.shape[1], grid.shape[2] * grid.shape[3],
                            grid.shape[4]))
                if row_images:
                    filename = os.path.join(outdir, safe_dir_name(layer),
                            prefix + 'image', '%d-%s' % (unit, suffix))
                    pool.add(strip[:,:-gap_pixels,:].copy(), filename)
                    # Image.fromarray(strip[:,:-gap_pixels,:]).save(filename,
                    #        optimize=True, quality=80)
                if single_images:
                    single_filename = os.path.join(outdir, safe_dir_name(layer),
                        prefix + 's-image', '%d-%s' % (unit, suffix))
                    pool.add(strip[:,:strip.shape[1] // row_length
                        - gap_pixels,:].copy(), single_filename)
                    # Image.fromarray(strip[:,:strip.shape[1] // row_length
                    #     - gap_pixels,:]).save(single_filename,
                    #             optimize=True, quality=80)
    pool.join()
    # Delete the shared memory map files
    clear_global_shared_files([g.filename
        for grid in [vizgrid, maskgrid, origrid, seggrid]
        for g in grid.values()])

global_shared_files = {}
def create_temp_mmap_grid(shape, dtype, fill):
    dtype = numpy.dtype(dtype)
    filename = os.path.join(tempfile.mkdtemp(), 'temp-%s-%s.mmap' %
            ('x'.join('%d' % s for s in shape), dtype.name))
    fid = open(filename, mode='w+b')
    original = numpy.memmap(fid, dtype=dtype, mode='w+', shape=shape)
    original.fid = fid
    original[...] = fill
    global_shared_files[filename] = original
    return original

def shared_temp_mmap_grid(filename, shape, dtype):
    if filename not in global_shared_files:
        global_shared_files[filename] = numpy.memmap(
                filename, dtype=dtype, mode='r+', shape=shape)
    return global_shared_files[filename]

def clear_global_shared_files(filenames):
    for fn in filenames:
        if fn in global_shared_files:
            del global_shared_files[fn]
        try:
            os.unlink(fn)
        except OSError:
            pass

class VisualizeImageWorker(WorkerBase):
    def setup(self, memmap_grid_info):
        self.vizgrid, self.maskgrid, self.origrid, self.seggrid = [
                {layer: shared_temp_mmap_grid(*info)
                    for layer, info in grid.items()}
                for grid in memmap_grid_info]
    def work(self, layer, unit, rank,
            byte_im, acts, level, scale_offset, seg):
        self.origrid[layer][unit,:,rank,:byte_im.shape[0],:] = byte_im
        [self.vizgrid[layer][unit,:,rank,:byte_im.shape[0],:],
         self.maskgrid[layer][unit,:,rank,:byte_im.shape[0],:]] = (
                    activation_visualization(
                        byte_im,
                        acts,
                        level,
                        scale_offset=scale_offset,
                        return_mask=True))
        self.seggrid[layer][unit,:,rank,:byte_im.shape[0],:] = (
                    segment_visualization(seg, byte_im.shape[0:2]))

class SaveImageWorker(WorkerBase):
    def work(self, data, filename):
        Image.fromarray(data).save(filename, optimize=True, quality=80)

def score_tally_stats(label_category, tc, truth, cc, ic):
    pred = cc[label_category]
    total = tc[label_category][:, None]
    truth = truth[:, None]
    epsilon = 1e-20 # avoid division-by-zero
    union = pred + truth - ic
    iou = ic.double() / (union.double() + epsilon)
    arr = torch.empty(size=(2, 2) + ic.shape, dtype=ic.dtype, device=ic.device)
    arr[0, 0] = ic
    arr[0, 1] = pred - ic
    arr[1, 0] = truth - ic
    arr[1, 1] = total - union
    arr = arr.double() / total.double()
    mi = mutual_information(arr)
    je = joint_entropy(arr)
    iqr = mi / je
    iqr[torch.isnan(iqr)] = 0 # Zero out any 0/0
    return iou, iqr

def collect_quantiles_and_topk(outdir, model, segloader,
        segrunner, k=100, resolution=1024):
    '''
    Collects (estimated) quantile information and (exact) sorted top-K lists
    for every channel in the retained layers of the model.  Returns
    a map of quantiles (one RunningQuantile for each layer) along with
    a map of topk (one RunningTopK for each layer).
    '''
    device = next(model.parameters()).device
    features = model.retained_features()
    cached_quantiles = {
            layer: load_quantile_if_present(os.path.join(outdir,
                safe_dir_name(layer)), 'quantiles.npz',
                device=torch.device('cpu'))
            for layer in features }
    cached_topks = {
            layer: load_topk_if_present(os.path.join(outdir,
                safe_dir_name(layer)), 'topk.npz',
                device=torch.device('cpu'))
            for layer in features }
    if (all(value is not None for value in cached_quantiles.values()) and
        all(value is not None for value in cached_topks.values())):
        return cached_quantiles, cached_topks

    layer_batch_size = 8
    all_layers = list(features.keys())
    layer_batches = [all_layers[i:i+layer_batch_size]
            for i in range(0, len(all_layers), layer_batch_size)]

    quantiles, topks = {}, {}
    progress = default_progress()
    for layer_batch in layer_batches:
        for i, batch in enumerate(progress(segloader, desc='Quantiles')):
            # We don't actually care about the model output.
            model(batch[0].to(device))
            features = model.retained_features()
            # We care about the retained values
            for key in layer_batch:
                value = features[key]
                if topks.get(key, None) is None:
                    topks[key] = RunningTopK(k)
                if quantiles.get(key, None) is None:
                    quantiles[key] = RunningQuantile(resolution=resolution)
                topvalue = value
                if len(value.shape) > 2:
                    topvalue, _ = value.view(*(value.shape[:2] + (-1,))).max(2)
                    # Put the channel index last.
                    value = value.permute(
                            (0,) + tuple(range(2, len(value.shape))) + (1,)
                            ).contiguous().view(-1, value.shape[1])
                quantiles[key].add(value)
                topks[key].add(topvalue)
        # Save GPU memory
        for key in layer_batch:
            quantiles[key].to_(torch.device('cpu'))
            topks[key].to_(torch.device('cpu'))
    for layer in quantiles:
        save_state_dict(quantiles[layer],
                os.path.join(outdir, safe_dir_name(layer), 'quantiles.npz'))
        save_state_dict(topks[layer],
                os.path.join(outdir, safe_dir_name(layer), 'topk.npz'))
    return quantiles, topks

def collect_bincounts(outdir, model, segloader, levels, segrunner):
    '''
    Returns label_counts, category_activation_counts, and intersection_counts,
    across the data set, counting the pixels of intersection between upsampled,
    thresholded model featuremaps, with segmentation classes in the segloader.

    label_counts (independent of model): pixels across the data set that
        are labeled with the given label.
    category_activation_counts (one per layer): for each feature channel,
        pixels across the dataset where the channel exceeds the level
        threshold.  There is one count per category: activations only
        contribute to the categories for which any category labels are
        present on the images.
    intersection_counts (one per layer): for each feature channel and
        label, pixels across the dataset where the channel exceeds
        the level, and the labeled segmentation class is also present.

    This is a performance-sensitive function.  Best performance is
    achieved with a counting scheme which assumes a segloader with
    batch_size 1.
    '''
    # Load cached data if present
    (iou_scores, iqr_scores,
            total_counts, label_counts, category_activation_counts,
            intersection_counts) = {}, {}, None, None, {}, {}
    found_all = True
    for layer in model.retained_features():
        filename = os.path.join(outdir, safe_dir_name(layer), 'bincounts.npz')
        if os.path.isfile(filename):
            data = numpy.load(filename)
            iou_scores[layer] = torch.from_numpy(data['iou_scores'])
            iqr_scores[layer] = torch.from_numpy(data['iqr_scores'])
            total_counts = torch.from_numpy(data['total_counts'])
            label_counts = torch.from_numpy(data['label_counts'])
            category_activation_counts[layer] = torch.from_numpy(
                    data['category_activation_counts'])
            intersection_counts[layer] = torch.from_numpy(
                    data['intersection_counts'])
        else:
            found_all = False
    if found_all:
        return (iou_scores, iqr_scores,
            total_counts, label_counts, category_activation_counts,
            intersection_counts)

    device = next(model.parameters()).device
    labelcat, categories = segrunner.get_label_and_category_names()
    label_category = [categories.index(c) if c in categories else 0
                for l, c in labelcat]
    num_labels, num_categories = (len(n) for n in [labelcat, categories])

    # One-hot vector of category for each label
    labelcat = torch.zeros(num_labels, num_categories,
            dtype=torch.long, device=device)
    labelcat.scatter_(1, torch.from_numpy(numpy.array(label_category,
        dtype='int64')).to(device)[:,None], 1)
    # Running bincounts
    # activation_counts = {}
    assert segloader.batch_size == 1 # category_activation_counts needs this.
    category_activation_counts = {}
    intersection_counts = {}
    label_counts = torch.zeros(num_labels, dtype=torch.long, device=device)
    total_counts = torch.zeros(num_categories, dtype=torch.long, device=device)
    progress = default_progress()
    scale_offset_map = getattr(model, 'scale_offset', None)
    upsample_grids = {}
    # total_batch_categories = torch.zeros(
    #         labelcat.shape[1], dtype=torch.long, device=device)
    for i, batch in enumerate(progress(segloader, desc='Bincounts')):
        seg, batch_label_counts, _, imshape = segrunner.run_and_segment_batch(
                batch, model, want_bincount=True, want_rgb=True)
        bc = batch_label_counts.cpu()
        batch_label_counts = batch_label_counts.to(device)
        seg = seg.to(device)
        features = model.retained_features()
        # Accumulate bincounts and identify nonzeros
        label_counts += batch_label_counts[0]
        batch_labels = bc[0].nonzero()[:,0]
        batch_categories = labelcat[batch_labels].max(0)[0]
        total_counts += batch_categories * (
                seg.shape[0] * seg.shape[2] * seg.shape[3])
        for key, value in features.items():
            if key not in upsample_grids:
                upsample_grids[key] = upsample_grid(value.shape[2:],
                        seg.shape[2:], imshape,
                        scale_offset=scale_offset_map.get(key, None)
                            if scale_offset_map is not None else None,
                        dtype=value.dtype, device=value.device)
            upsampled = torch.nn.functional.grid_sample(value,
                    upsample_grids[key], padding_mode='border')
            amask = (upsampled > levels[key][None,:,None,None].to(
                upsampled.device))
            ac = amask.int().view(amask.shape[1], -1).sum(1)
            # if key not in activation_counts:
            #     activation_counts[key] = ac
            # else:
            #     activation_counts[key] += ac
            # The fastest approach: sum over each label separately!
            for label in batch_labels.tolist():
                if label == 0:
                    continue  # ignore the background label
                imask = amask * ((seg == label).max(dim=1, keepdim=True)[0])
                ic = imask.int().view(imask.shape[1], -1).sum(1)
                if key not in intersection_counts:
                    intersection_counts[key] = torch.zeros(num_labels,
                            amask.shape[1], dtype=torch.long, device=device)
                intersection_counts[key][label] += ic
            # Count activations within images that have category labels.
            # Note: This only makes sense with batch-size one
            # total_batch_categories += batch_categories
            cc = batch_categories[:,None] * ac[None,:]
            if key not in category_activation_counts:
                category_activation_counts[key] = cc
            else:
                category_activation_counts[key] += cc
    iou_scores = {}
    iqr_scores = {}
    for k in intersection_counts:
        iou_scores[k], iqr_scores[k] = score_tally_stats(
            label_category, total_counts, label_counts,
            category_activation_counts[k], intersection_counts[k])
    for k in intersection_counts:
        numpy.savez(os.path.join(outdir, safe_dir_name(k), 'bincounts.npz'),
                iou_scores=iou_scores[k].cpu().numpy(),
                iqr_scores=iqr_scores[k].cpu().numpy(),
                total_counts=total_counts.cpu().numpy(),
                label_counts=label_counts.cpu().numpy(),
                category_activation_counts=category_activation_counts[k]
                    .cpu().numpy(),
                intersection_counts=intersection_counts[k].cpu().numpy(),
                levels=levels[k].cpu().numpy())
    return (iou_scores, iqr_scores,
            total_counts, label_counts, category_activation_counts,
            intersection_counts)

def collect_cond_quantiles(outdir, model, segloader, segrunner):
    '''
    Returns maxiou and maxiou_level across the data set, one per layer.

    This is a performance-sensitive function.  Best performance is
    achieved with a counting scheme which assumes a segloader with
    batch_size 1.
    '''
    device = next(model.parameters()).device
    cached_cond_quantiles = {
            layer: load_conditional_quantile_if_present(os.path.join(outdir,
                safe_dir_name(layer)), 'cond_quantiles.npz') # on cpu
            for layer in model.retained_features() }
    label_fracs = load_npy_if_present(outdir, 'label_fracs.npy', 'cpu')
    if label_fracs is not None and all(
            value is not None for value in cached_cond_quantiles.values()):
        return cached_cond_quantiles, label_fracs

    labelcat, categories = segrunner.get_label_and_category_names()
    label_category = [categories.index(c) if c in categories else 0
                for l, c in labelcat]
    num_labels, num_categories = (len(n) for n in [labelcat, categories])

    # One-hot vector of category for each label
    labelcat = torch.zeros(num_labels, num_categories,
            dtype=torch.long, device=device)
    labelcat.scatter_(1, torch.from_numpy(numpy.array(label_category,
        dtype='int64')).to(device)[:,None], 1)
    # Running maxiou
    assert segloader.batch_size == 1 # category_activation_counts needs this.
    conditional_quantiles = {}
    label_counts = torch.zeros(num_labels, dtype=torch.long, device=device)
    pixel_count = 0
    progress = default_progress()
    scale_offset_map = getattr(model, 'scale_offset', None)
    upsample_grids = {}
    common_conditions = set()
    if label_fracs is None or label_fracs is 0:
        for i, batch in enumerate(progress(segloader, desc='label fracs')):
            seg, batch_label_counts, im, _ = segrunner.run_and_segment_batch(
                    batch, model, want_bincount=True, want_rgb=True)
            batch_label_counts = batch_label_counts.to(device)
            features = model.retained_features()
            # Accumulate bincounts and identify nonzeros
            label_counts += batch_label_counts[0]
            pixel_count += seg.shape[2] * seg.shape[3]
        label_fracs = (label_counts.cpu().float() / pixel_count)[:, None, None]
        numpy.save(os.path.join(outdir, 'label_fracs.npy'), label_fracs)

    skip_threshold = 1e-4
    skip_labels = set(i.item()
        for i in (label_fracs.view(-1) < skip_threshold).nonzero().view(-1))

    for layer in progress(model.retained_features().keys(), desc='CQ layers'):
        if cached_cond_quantiles.get(layer, None) is not None:
            conditional_quantiles[layer] = cached_cond_quantiles[layer]
            continue

        for i, batch in enumerate(progress(segloader, desc='Condquant')):
            seg, batch_label_counts, _, imshape = (
                    segrunner.run_and_segment_batch(
                         batch, model, want_bincount=True, want_rgb=True))
            bc = batch_label_counts.cpu()
            batch_label_counts = batch_label_counts.to(device)
            features = model.retained_features()
            # Accumulate bincounts and identify nonzeros
            label_counts += batch_label_counts[0]
            pixel_count += seg.shape[2] * seg.shape[3]
            batch_labels = bc[0].nonzero()[:,0]
            batch_categories = labelcat[batch_labels].max(0)[0]
            cpu_seg = None
            value = features[layer]
            if layer not in upsample_grids:
                upsample_grids[layer] = upsample_grid(value.shape[2:],
                        seg.shape[2:], imshape,
                        scale_offset=scale_offset_map.get(layer, None)
                            if scale_offset_map is not None else None,
                        dtype=value.dtype, device=value.device)
            if layer not in conditional_quantiles:
                conditional_quantiles[layer] = RunningConditionalQuantile(
                        resolution=2048)
            upsampled = torch.nn.functional.grid_sample(value,
                    upsample_grids[layer], padding_mode='border').view(
                            value.shape[1], -1)
            conditional_quantiles[layer].add(('all',), upsampled.t())
            cpu_upsampled = None
            for label in batch_labels.tolist():
                if label in skip_labels:
                    continue
                label_key = ('label', label)
                if label_key in common_conditions:
                    imask = (seg == label).max(dim=1)[0].view(-1)
                    intersected = upsampled[:, imask]
                    conditional_quantiles[layer].add(('label', label),
                            intersected.t())
                else:
                    if cpu_seg is None:
                        cpu_seg = seg.cpu()
                    if cpu_upsampled is None:
                        cpu_upsampled = upsampled.cpu()
                    imask = (cpu_seg == label).max(dim=1)[0].view(-1)
                    intersected = cpu_upsampled[:, imask]
                    conditional_quantiles[layer].add(('label', label),
                            intersected.t())
            if num_categories > 1:
                for cat in batch_categories.nonzero()[:,0]:
                    conditional_quantiles[layer].add(('cat', cat.item()),
                            upsampled.t())
            # Move the most common conditions to the GPU.
            if i and not i & (i - 1):  # if i is a power of 2:
                cq = conditional_quantiles[layer]
                common_conditions = set(cq.most_common_conditions(64))
                cq.to_('cpu', [k for k in cq.running_quantiles.keys()
                        if k not in common_conditions])
        # When a layer is done, get it off the GPU
        conditional_quantiles[layer].to_('cpu')

    label_fracs = (label_counts.cpu().float() / pixel_count)[:, None, None]

    for cq in conditional_quantiles.values():
        cq.to_('cpu')

    for layer in conditional_quantiles:
        save_state_dict(conditional_quantiles[layer],
            os.path.join(outdir, safe_dir_name(layer), 'cond_quantiles.npz'))
    numpy.save(os.path.join(outdir, 'label_fracs.npy'), label_fracs)

    return conditional_quantiles, label_fracs


def collect_maxiou(outdir, model, segloader, segrunner):
    '''
    Returns maxiou and maxiou_level across the data set, one per layer.

    This is a performance-sensitive function.  Best performance is
    achieved with a counting scheme which assumes a segloader with
    batch_size 1.
    '''
    device = next(model.parameters()).device
    conditional_quantiles, label_fracs = collect_cond_quantiles(
            outdir, model, segloader, segrunner)

    labelcat, categories = segrunner.get_label_and_category_names()
    label_category = [categories.index(c) if c in categories else 0
                for l, c in labelcat]
    num_labels, num_categories = (len(n) for n in [labelcat, categories])

    label_list = [('label', i) for i in range(num_labels)]
    category_list = [('all',)] if num_categories <= 1 else (
            [('cat', i) for i in range(num_categories)])
    max_iou, max_iou_level, max_iou_quantile = {}, {}, {}
    fracs = torch.logspace(-3, 0, 100)
    progress = default_progress()
    for layer, cq in progress(conditional_quantiles.items(), desc='Maxiou'):
        levels = cq.conditional(('all',)).quantiles(1 - fracs)
        denoms = 1 - cq.collected_normalize(category_list, levels)
        isects = (1 - cq.collected_normalize(label_list, levels)) * label_fracs
        unions = label_fracs + denoms[label_category, :, :] - isects
        iou = isects / unions
        # TODO: erase any for which threshold is bad
        max_iou[layer], level_bucket = iou.max(2)
        max_iou_level[layer] = levels[
                torch.arange(levels.shape[0])[None,:], level_bucket]
        max_iou_quantile[layer] = fracs[level_bucket]
    for layer in model.retained_features():
        numpy.savez(os.path.join(outdir, safe_dir_name(layer), 'max_iou.npz'),
            max_iou=max_iou[layer].cpu().numpy(),
            max_iou_level=max_iou_level[layer].cpu().numpy(),
            max_iou_quantile=max_iou_quantile[layer].cpu().numpy())
    return (max_iou, max_iou_level, max_iou_quantile)

def collect_iqr(outdir, model, segloader, segrunner):
    '''
    Returns iqr and iqr_level.

    This is a performance-sensitive function.  Best performance is
    achieved with a counting scheme which assumes a segloader with
    batch_size 1.
    '''
    max_iqr, max_iqr_level, max_iqr_quantile, max_iqr_iou  = {}, {}, {}, {}
    max_iqr_agreement = {}
    found_all = True
    for layer in model.retained_features():
        filename = os.path.join(outdir, safe_dir_name(layer), 'iqr.npz')
        if os.path.isfile(filename):
            data = numpy.load(filename)
            max_iqr[layer] = torch.from_numpy(data['max_iqr'])
            max_iqr_level[layer] = torch.from_numpy(data['max_iqr_level'])
            max_iqr_quantile[layer] = torch.from_numpy(data['max_iqr_quantile'])
            max_iqr_iou[layer] = torch.from_numpy(data['max_iqr_iou'])
            max_iqr_agreement[layer] = torch.from_numpy(
                    data['max_iqr_agreement'])
        else:
            found_all = False
    if found_all:
        return (max_iqr, max_iqr_level, max_iqr_quantile, max_iqr_iou,
            max_iqr_agreement)


    device = next(model.parameters()).device
    conditional_quantiles, label_fracs = collect_cond_quantiles(
            outdir, model, segloader, segrunner)

    labelcat, categories = segrunner.get_label_and_category_names()
    label_category = [categories.index(c) if c in categories else 0
                for l, c in labelcat]
    num_labels, num_categories = (len(n) for n in [labelcat, categories])

    label_list = [('label', i) for i in range(num_labels)]
    category_list = [('all',)] if num_categories <= 1 else (
            [('cat', i) for i in range(num_categories)])
    full_mi, full_je, full_iqr = {}, {}, {}
    fracs = torch.logspace(-3, 0, 100)
    progress = default_progress()
    for layer, cq in progress(conditional_quantiles.items(), desc='IQR'):
        levels = cq.conditional(('all',)).quantiles(1 - fracs)
        truth = label_fracs.to(device)
        preds = (1 - cq.collected_normalize(category_list, levels)
                )[label_category, :, :].to(device)
        cond_isects = 1 - cq.collected_normalize(label_list, levels).to(device)
        isects = cond_isects * truth
        unions = truth + preds - isects
        arr = torch.empty(size=(2, 2) + isects.shape, dtype=isects.dtype,
                device=device)
        arr[0, 0] = isects
        arr[0, 1] = preds - isects
        arr[1, 0] = truth - isects
        arr[1, 1] = 1 - unions
        arr.clamp_(0, 1)
        mi = mutual_information(arr)
        mi[:,:,-1] = 0  # at the 1.0 quantile should be no MI.
        # Don't trust mi when less than label_frac is less than 1e-3,
        # because our samples are too small.
        mi[label_fracs.view(-1) < 1e-3, :, :] = 0
        je = joint_entropy(arr)
        iqr = mi / je
        iqr[torch.isnan(iqr)] = 0 # Zero out any 0/0
        full_mi[layer] = mi.cpu()
        full_je[layer] = je.cpu()
        full_iqr[layer] = iqr.cpu()
        del mi, je
        agreement = isects + arr[1, 1]
        # When optimizing, maximize only over those pairs where the
        # unit is positively correlated with the label, and where the
        # threshold level is positive
        positive_iqr = iqr
        positive_iqr[agreement <= 0.8] = 0
        positive_iqr[(levels <= 0.0)[None, :, :].expand(positive_iqr.shape)] = 0
        # TODO: erase any for which threshold is bad
        maxiqr, level_bucket = positive_iqr.max(2)
        max_iqr[layer] = maxiqr.cpu()
        max_iqr_level[layer] = levels.to(device)[
                torch.arange(levels.shape[0])[None,:], level_bucket].cpu()
        max_iqr_quantile[layer] = fracs.to(device)[level_bucket].cpu()
        max_iqr_agreement[layer] = agreement[
                torch.arange(agreement.shape[0])[:, None],
                torch.arange(agreement.shape[1])[None, :],
                level_bucket].cpu()

        # Compute the iou that goes with each maximized iqr
        matching_iou = (isects[
                torch.arange(isects.shape[0])[:, None],
                torch.arange(isects.shape[1])[None, :],
                level_bucket] /
            unions[
                torch.arange(unions.shape[0])[:, None],
                torch.arange(unions.shape[1])[None, :],
                level_bucket])
        matching_iou[torch.isnan(matching_iou)] = 0
        max_iqr_iou[layer] = matching_iou.cpu()
    for layer in model.retained_features():
        numpy.savez(os.path.join(outdir, safe_dir_name(layer), 'iqr.npz'),
            max_iqr=max_iqr[layer].cpu().numpy(),
            max_iqr_level=max_iqr_level[layer].cpu().numpy(),
            max_iqr_quantile=max_iqr_quantile[layer].cpu().numpy(),
            max_iqr_iou=max_iqr_iou[layer].cpu().numpy(),
            max_iqr_agreement=max_iqr_agreement[layer].cpu().numpy(),
            full_mi=full_mi[layer].cpu().numpy(),
            full_je=full_je[layer].cpu().numpy(),
            full_iqr=full_iqr[layer].cpu().numpy())
    return (max_iqr, max_iqr_level, max_iqr_quantile, max_iqr_iou,
            max_iqr_agreement)

def mutual_information(arr):
    total = 0
    for j in range(arr.shape[0]):
        for k in range(arr.shape[1]):
            joint = arr[j,k]
            ind = arr[j,:].sum(dim=0) * arr[:,k].sum(dim=0)
            term = joint * (joint / ind).log()
            term[torch.isnan(term)] = 0
            total += term
    return total.clamp_(0)

def joint_entropy(arr):
    total = 0
    for j in range(arr.shape[0]):
        for k in range(arr.shape[1]):
            joint = arr[j,k]
            term = joint * joint.log()
            term[torch.isnan(term)] = 0
            total += term
    return (-total).clamp_(0)

def information_quality_ratio(arr):
    iqr = mutual_information(arr) / joint_entropy(arr)
    iqr[torch.isnan(iqr)] = 0
    return iqr

def collect_covariance(outdir, model, segloader, segrunner):
    '''
    Returns label_mean, label_variance, unit_mean, unit_variance,
    and cross_covariance across the data set.

    label_mean, label_variance (independent of model):
        treating the label as a one-hot, each label's mean and variance.
    unit_mean, unit_variance (one per layer): for each feature channel,
        the mean and variance of the activations in that channel.
    cross_covariance (one per layer): the cross covariance between the
        labels and the units in the layer.
    '''
    device = next(model.parameters()).device
    cached_covariance = {
            layer: load_covariance_if_present(os.path.join(outdir,
                safe_dir_name(layer)), 'covariance.npz', device=device)
            for layer in model.retained_features() }
    if all(value is not None for value in cached_covariance.values()):
        return cached_covariance
    labelcat, categories = segrunner.get_label_and_category_names()
    label_category = [categories.index(c) if c in categories else 0
                for l, c in labelcat]
    num_labels, num_categories = (len(n) for n in [labelcat, categories])

    # Running covariance
    cov = {}
    progress = default_progress()
    scale_offset_map = getattr(model, 'scale_offset', None)
    upsample_grids = {}
    for i, batch in enumerate(progress(segloader, desc='Covariance')):
        seg, _, _, imshape = segrunner.run_and_segment_batch(batch, model,
                want_rgb=True)
        features = model.retained_features()
        ohfeats = multilabel_onehot(seg, num_labels, ignore_index=0)
        # Accumulate bincounts and identify nonzeros
        for key, value in features.items():
            if key not in upsample_grids:
                upsample_grids[key] = upsample_grid(value.shape[2:],
                        seg.shape[2:], imshape,
                        scale_offset=scale_offset_map.get(key, None)
                            if scale_offset_map is not None else None,
                        dtype=value.dtype, device=value.device)
            upsampled = torch.nn.functional.grid_sample(value,
                    upsample_grids[key].expand(
                        (value.shape[0],) + upsample_grids[key].shape[1:]),
                    padding_mode='border')
            if key not in cov:
                cov[key] = RunningCrossCovariance()
            cov[key].add(upsampled, ohfeats)
    for layer in cov:
        save_state_dict(cov[layer],
                os.path.join(outdir, safe_dir_name(layer), 'covariance.npz'))
    return cov

def multilabel_onehot(labels, num_labels, dtype=None, ignore_index=None):
    '''
    Converts a multilabel tensor into a onehot tensor.

    The input labels is a tensor of shape (samples, multilabels, y, x).
    The output is a tensor of shape (samples, num_labels, y, x).
    If ignore_index is specified, labels with that index are ignored.
    Each x in labels should be 0 <= x < num_labels, or x == ignore_index.
    '''
    assert ignore_index is None or ignore_index <= 0
    if dtype is None:
        dtype = torch.float
    device = labels.device
    chans = num_labels + (-ignore_index if ignore_index else 0)
    outshape = (labels.shape[0], chans) + labels.shape[2:]
    result = torch.zeros(outshape, device=device, dtype=dtype)
    if ignore_index and ignore_index < 0:
        labels = labels + (-ignore_index)
    result.scatter_(1, labels, 1)
    if ignore_index and ignore_index < 0:
        result = result[:, -ignore_index:]
    elif ignore_index is not None:
        result[:, ignore_index] = 0
    return result

def load_npy_if_present(outdir, filename, device):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        return torch.from_numpy(data).to(device)
    return 0

def load_npz_if_present(outdir, filename, varnames, device):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        numpy_result = [data[n] for n in varnames]
        return tuple(torch.from_numpy(data).to(device) for data in numpy_result)
    return None

def load_quantile_if_present(outdir, filename, device):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        result = RunningQuantile(state=data)
        result.to_(device)
        return result
    return None

def load_conditional_quantile_if_present(outdir, filename):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        result = RunningConditionalQuantile(state=data)
        return result
    return None

def load_topk_if_present(outdir, filename, device):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        result = RunningTopK(state=data)
        result.to_(device)
        return result
    return None

def load_covariance_if_present(outdir, filename, device):
    filepath = os.path.join(outdir, filename)
    if os.path.isfile(filepath):
        data = numpy.load(filepath)
        result = RunningCrossCovariance(state=data)
        result.to_(device)
        return result
    return None

def save_state_dict(obj, filepath):
    dirname = os.path.dirname(filepath)
    os.makedirs(dirname, exist_ok=True)
    dic = obj.state_dict()
    numpy.savez(filepath, **dic)

def upsample_grid(data_shape, target_shape, input_shape=None,
        scale_offset=None, dtype=torch.float, device=None):
    '''Prepares a grid to use with grid_sample to upsample a batch of
    features in data_shape to the target_shape. Can use scale_offset
    and input_shape to center the grid in a nondefault way: scale_offset
    maps feature pixels to input_shape pixels, and it is assumed that
    the target_shape is a uniform downsampling of input_shape.'''
    # Default is that nothing is resized.
    if target_shape is None:
        target_shape = data_shape
    # Make a default scale_offset to fill the image if there isn't one
    if scale_offset is None:
        scale = tuple(float(ts) / ds
                for ts, ds in zip(target_shape, data_shape))
        offset = tuple(0.5 * s - 0.5 for s in scale)
    else:
        scale, offset = (v for v in zip(*scale_offset))
        # Handle downsampling for different input vs target shape.
        if input_shape is not None:
            scale = tuple(s * (ts - 1) / (ns - 1)
                    for s, ns, ts in zip(scale, input_shape, target_shape))
            offset = tuple(o * (ts - 1) / (ns - 1)
                    for o, ns, ts in zip(offset, input_shape, target_shape))
    # Pytorch needs target coordinates in terms of source coordinates [-1..1]
    ty, tx = (((torch.arange(ts, dtype=dtype, device=device) - o)
                  * (2 / (s * (ss - 1))) - 1)
        for ts, ss, s, o, in zip(target_shape, data_shape, scale, offset))
    # Whoa, note that grid_sample reverses the order y, x -> x, y.
    grid = torch.stack(
        (tx[None,:].expand(target_shape), ty[:,None].expand(target_shape)),2
       )[None,:,:,:].expand((1, target_shape[0], target_shape[1], 2))
    return grid

def safe_dir_name(filename):
    keepcharacters = (' ','.','_','-')
    return ''.join(c
            for c in filename if c.isalnum() or c in keepcharacters).rstrip()

bargraph_palette = [
    ('#4B4CBF', '#B6B6F2'),
    ('#55B05B', '#B6F2BA'),
    ('#50BDAC', '#A5E5DB'),
    ('#81C679', '#C0FF9B'),
    ('#F0883B', '#F2CFB6'),
    ('#D4CF24', '#F2F1B6'),
    ('#D92E2B', '#F2B6B6'),
    ('#AB6BC6', '#CFAAFF'),
]

def make_svg_bargraph(labels, heights, categories,
        barheight=100, barwidth=12, show_labels=True, filename=None):
    # if len(labels) == 0:
    #     return # Nothing to do
    unitheight = float(barheight) / max(max(heights, default=1), 1)
    textheight = barheight if show_labels else 0
    labelsize = float(barwidth)
    gap = float(barwidth) / 4
    textsize = barwidth + gap
    rollup = max(heights, default=1)
    textmargin = float(labelsize) * 2 / 3
    leftmargin = 32
    rightmargin = 8
    svgwidth = len(heights) * (barwidth + gap) + 2 * leftmargin + rightmargin
    svgheight = barheight + textheight

    # create an SVG XML element
    svg = et.Element('svg', width=str(svgwidth), height=str(svgheight),
            version='1.1', xmlns='http://www.w3.org/2000/svg')

    # Draw the bar graph
    basey = svgheight - textheight
    x = leftmargin
    # Add units scale on left
    if len(heights):
        for h in [1, (max(heights) + 1) // 2, max(heights)]:
            et.SubElement(svg, 'text', x='0', y='0',
                style=('font-family:sans-serif;font-size:%dpx;' +
                'text-anchor:end;alignment-baseline:hanging;' +
                'transform:translate(%dpx, %dpx);') %
                (textsize, x - gap, basey - h * unitheight)).text = str(h)
        et.SubElement(svg, 'text', x='0', y='0',
                style=('font-family:sans-serif;font-size:%dpx;' +
                'text-anchor:middle;' +
                'transform:translate(%dpx, %dpx) rotate(-90deg)') %
                (textsize, x - gap - textsize, basey - h * unitheight / 2)
                ).text = 'units'
    # Draw big category background rectangles
    for catindex, (cat, catcount) in enumerate(categories):
        if not catcount:
            continue
        et.SubElement(svg, 'rect', x=str(x), y=str(basey - rollup * unitheight),
                width=(str((barwidth + gap) * catcount - gap)),
                height = str(rollup*unitheight),
                fill=bargraph_palette[catindex % len(bargraph_palette)][1])
        x += (barwidth + gap) * catcount
    # Draw small bars as well as 45degree text labels
    x = leftmargin
    catindex = -1
    catcount = 0
    for label, height in zip(labels, heights):
        while not catcount and catindex <= len(categories):
            catindex += 1
            catcount = categories[catindex][1]
            color = bargraph_palette[catindex % len(bargraph_palette)][0]
        et.SubElement(svg, 'rect', x=str(x), y=str(basey-(height * unitheight)),
                width=str(barwidth), height=str(height * unitheight),
                fill=color)
        x += barwidth
        if show_labels:
            et.SubElement(svg, 'text', x='0', y='0',
                style=('font-family:sans-serif;font-size:%dpx;text-anchor:end;'+
                'transform:translate(%dpx, %dpx) rotate(-45deg);') %
                (labelsize, x, basey + textmargin)).text = readable(label)
        x += gap
        catcount -= 1
    # Text labels for each category
    x = leftmargin
    for cat, catcount in categories:
        if not catcount:
            continue
        et.SubElement(svg, 'text', x='0', y='0',
            style=('font-family:sans-serif;font-size:%dpx;text-anchor:end;'+
            'transform:translate(%dpx, %dpx) rotate(-90deg);') %
            (textsize, x + (barwidth + gap) * catcount - gap,
                basey - rollup * unitheight + gap)).text = '%d %s' % (
                    catcount, readable(cat + ('s' if catcount != 1 else '')))
        x += (barwidth + gap) * catcount
    # Output - this is the bare svg.
    result = et.tostring(svg)
    if filename:
        f = open(filename, 'wb')
        # When writing to a file a special header is needed.
        f.write(''.join([
            '<?xml version=\"1.0\" standalone=\"no\"?>\n',
            '<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n',
            '\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n']
            ).encode('utf-8'))
        f.write(result)
        f.close()
    return result

readable_replacements = [(re.compile(r[0]), r[1]) for r in [
    (r'-[sc]$', ''),
    (r'_', ' '),
    ]]

def readable(label):
    for pattern, subst in readable_replacements:
        label= re.sub(pattern, subst, label)
    return label

def reverse_normalize_from_transform(transform):
    '''
    Crawl around the transforms attached to a dataset looking for a
    Normalize transform, and return it a corresponding ReverseNormalize,
    or None if no normalization is found.
    '''
    if isinstance(transform, torchvision.transforms.Normalize):
        return ReverseNormalize(transform.mean, transform.std)
    t = getattr(transform, 'transform', None)
    if t is not None:
        return reverse_normalize_from_transform(t)
    transforms = getattr(transform, 'transforms', None)
    if transforms is not None:
        for t in reversed(transforms):
            result = reverse_normalize_from_transform(t)
            if result is not None:
                return result
    return None

class ReverseNormalize:
    '''
    Applies the reverse of torchvision.transforms.Normalize.
    '''
    def __init__(self, mean, stdev):
        mean = numpy.array(mean)
        stdev = numpy.array(stdev)
        self.mean = torch.from_numpy(mean)[None,:,None,None].float()
        self.stdev = torch.from_numpy(stdev)[None,:,None,None].float()
    def __call__(self, data):
        device = data.device
        return data.mul(self.stdev.to(device)).add_(self.mean.to(device))

class ImageOnlySegRunner:
    def __init__(self, dataset, recover_image=None):
        if recover_image is None:
            recover_image = reverse_normalize_from_transform(dataset)
        self.recover_image = recover_image
        self.dataset = dataset
    def get_label_and_category_names(self):
        return [('-', '-')], ['-']
    def run_and_segment_batch(self, batch, model,
            want_bincount=False, want_rgb=False):
        [im] = batch
        device = next(model.parameters()).device
        if want_rgb:
            rgb = self.recover_image(im.clone()
                ).permute(0, 2, 3, 1).mul_(255).clamp(0, 255).byte()
        else:
            rgb = None
        # Stubs for seg and bc
        seg = torch.zeros(im.shape[0], 1, 1, 1, dtype=torch.long)
        bc = torch.ones(im.shape[0], 1, dtype=torch.long)
        # Run the model.
        model(im.to(device))
        return seg, bc, rgb, im.shape[2:]

class ClassifierSegRunner:
    def __init__(self, dataset, recover_image=None):
        # The dataset contains explicit segmentations
        if recover_image is None:
            recover_image = reverse_normalize_from_transform(dataset)
        self.recover_image = recover_image
        self.dataset = dataset
    def get_label_and_category_names(self):
        catnames = self.dataset.categories
        label_and_cat_names = [(readable(label),
            catnames[self.dataset.label_category[i]])
                for i, label in enumerate(self.dataset.labels)]
        return label_and_cat_names, catnames
    def run_and_segment_batch(self, batch, model,
            want_bincount=False, want_rgb=False):
        '''
        Runs the dissected model on one batch of the dataset, and
        returns a multilabel semantic segmentation for the data.
        Given a batch of size (n, c, y, x) the segmentation should
        be a (long integer) tensor of size (n, d, y//r, x//r) where
        d is the maximum number of simultaneous labels given to a pixel,
        and where r is some (optional) resolution reduction factor.
        In the segmentation returned, the label `0` is reserved for
        the background "no-label".

        In addition to the segmentation, bc, rgb, and shape are returned
        where bc is a per-image bincount counting returned label pixels,
        rgb is a viewable (n, y, x, rgb) byte image tensor for the data
        for visualizations (reversing normalizations, for example), and
        shape is the (y, x) size of the data.  If want_bincount or
        want_rgb are False, those return values may be None.
        '''
        im, seg, bc = batch
        device = next(model.parameters()).device
        if want_rgb:
            rgb = self.recover_image(im.clone()
                ).permute(0, 2, 3, 1).mul_(255).clamp(0, 255).byte()
        else:
            rgb = None
        # Run the model.
        model(im.to(device))
        return seg, bc, rgb, im.shape[2:]

class GeneratorSegRunner:
    def __init__(self, segmenter):
        # The segmentations are given by an algorithm
        if segmenter is None:
            segmenter = UnifiedParsingSegmenter(segsizes=[256], segdiv='quad')
        self.segmenter = segmenter
        self.num_classes = len(segmenter.get_label_and_category_names()[0])
    def get_label_and_category_names(self):
        return self.segmenter.get_label_and_category_names()
    def run_and_segment_batch(self, batch, model,
            want_bincount=False, want_rgb=False):
        '''
        Runs the dissected model on one batch of the dataset, and
        returns a multilabel semantic segmentation for the data.
        Given a batch of size (n, c, y, x) the segmentation should
        be a (long integer) tensor of size (n, d, y//r, x//r) where
        d is the maximum number of simultaneous labels given to a pixel,
        and where r is some (optional) resolution reduction factor.
        In the segmentation returned, the label `0` is reserved for
        the background "no-label".

        In addition to the segmentation, bc, rgb, and shape are returned
        where bc is a per-image bincount counting returned label pixels,
        rgb is a viewable (n, y, x, rgb) byte image tensor for the data
        for visualizations (reversing normalizations, for example), and
        shape is the (y, x) size of the data.  If want_bincount or
        want_rgb are False, those return values may be None.
        '''
        device = next(model.parameters()).device
        z_batch = batch[0]
        tensor_images = model(z_batch.to(device))
        seg = self.segmenter.segment_batch(tensor_images, downsample=2)
        if want_bincount:
            index = torch.arange(z_batch.shape[0],
                    dtype=torch.long, device=device)
            bc = (seg + index[:, None, None, None] * self.num_classes).view(-1
                ).bincount(minlength=z_batch.shape[0] * self.num_classes)
            bc = bc.view(z_batch.shape[0], self.num_classes)
        else:
            bc = None
        if want_rgb:
            images = ((tensor_images + 1) / 2 * 255)
            rgb = images.permute(0, 2, 3, 1).clamp(0, 255).byte()
        else:
            rgb = None
        return seg, bc, rgb, tensor_images.shape[2:]