text:
"""Histograms of ranked posterior draws, plotted for each chain."""
import numpy as np
import scipy.stats

from ..data import convert_to_dataset
from .plot_utils import (
    _scale_fig_size,
    xarray_var_iter,
    default_grid,
    _create_axes_grid,
    make_label,
)
from ..utils import _var_names

def _sturges_formula(dataset, mult=1):
    """Use Sturges' formula to determine number of bins.

    See https://en.wikipedia.org/wiki/Histogram#Sturges'_formula
    or https://doi.org/10.1080%2F01621459.1926.10502161

    Parameters
    ----------
    dataset: xarray.DataSet
        Must have the `draw` dimension

    mult: float
        Used to scale the number of bins up or down. Default is 1 for Sturges' formula.

    Returns
    -------
    int
        Number of bins to use
    """
    return int(np.ceil(mult * np.log2(dataset.draw.size)) + 1)

def plot_rank(data, var_names=None, coords=None, bins=None, ref_line=True, figsize=None, axes=None):
    """Plot rank order statistics of chains.

    From the paper: Rank plots are histograms of the ranked posterior
    draws (ranked over all chains) plotted separately for each chain.
    If all of the chains are targeting the same posterior, we expect
    the ranks in each chain to be uniform, whereas if one chain has a
    different location or scale parameter, this will be reflected in
    the deviation from uniformity. If rank plots of all chains look
    similar, this indicates good mixing of the chains.

    This plot was introduced by Aki Vehtari, Andrew Gelman, Daniel
    Simpson, Bob Carpenter, Paul-Christian Burkner (2019):
    Rank-normalization, folding, and localization: An improved R-hat
    for assessing convergence of MCMC.
    arXiv preprint https://arxiv.org/abs/1903.08008

    Parameters
    ----------
    data : obj
        Any object that can be converted to an az.InferenceData object
        Refer to documentation of az.convert_to_dataset for details
    var_names : string or list of variable names
        Variables to be plotted
    coords : mapping, optional
        Coordinates of var_names to be plotted. Passed to `Dataset.sel`
    bins : None or passed to np.histogram
        Binning strategy used for histogram. By default uses twice the
        result of Sturges' formula. See `np.histogram` documenation for
        other available arguments.
    ref_line : boolean
        Whether to include a dashed line showing where a uniform
        distribution would lie
    figsize : tuple
        Figure size. If None it will be defined automatically.
    ax : axes
        Matplotlib axes. Defaults to None.

    Returns
    -------
    ax : matplotlib axes

    Examples
    --------
    Show a default rank plot

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> data = az.load_arviz_data('centered_eight')
        >>> az.plot_rank(data)

    Recreate Figure 13 from the arxiv preprint

    .. plot::
        :context: close-figs

        >>> import arviz as az
        >>> data = az.load_arviz_data('centered_eight')
        >>> az.plot_rank(data, var_names='tau')
    """
    posterior_data = convert_to_dataset(data, group="posterior")
    if coords is not None:
        posterior_data = posterior_data.sel(**coords)
    var_names = _var_names(var_names, posterior_data)
    plotters = list(xarray_var_iter(posterior_data, var_names=var_names, combined=True))

    if bins is None:
        # Use double Sturges' formula
        bins = _sturges_formula(posterior_data, mult=2)

    if axes is None:
        rows, cols = default_grid(len(plotters))

        figsize, ax_labelsize, titlesize, _, _, _ = _scale_fig_size(
            figsize, None, rows=rows, cols=cols
        )
        _, axes = _create_axes_grid(len(plotters), rows, cols, figsize=figsize, squeeze=False)

    for ax, (var_name, selection, var_data) in zip(axes.ravel(), plotters):
        ranks = scipy.stats.rankdata(var_data).reshape(var_data.shape)
        all_counts = []
        for row in ranks:
            counts, bin_ary = np.histogram(row, bins=bins, range=(0, ranks.size))
            all_counts.append(counts)
        all_counts = np.array(all_counts)
        gap = all_counts.max() * 1.05
        width = bin_ary[1] - bin_ary[0]

        # Center the bins
        bin_ary = (bin_ary[1:] + bin_ary[:-1]) / 2

        y_ticks = []
        for idx, counts in enumerate(all_counts):
            y_ticks.append(idx * gap)
            if ref_line:
                # Line where data is uniform
                ax.axhline(y=y_ticks[-1] + counts.mean(), linestyle="--", color="C1")
            # fake an x-axis
            ax.axhline(y=y_ticks[-1], color="k", lw=1)
            ax.bar(
                bin_ary,
                counts,
                bottom=y_ticks[-1],
                width=width,
                align="center",
                color="C0",
                edgecolor=ax.get_facecolor(),
            )
        ax.set_xlabel("Rank (all chains)", fontsize=ax_labelsize)
        ax.set_ylabel("Chain", fontsize=ax_labelsize)
        ax.set_yticks(y_ticks)
        ax.set_yticklabels(np.arange(len(y_ticks)))
        ax.set_title(make_label(var_name, selection), fontsize=titlesize)

    return axes

max_stars_repo_name: fehiepsi/arviz

max_issues_repo_name: fehiepsi/arviz

max_forks_repo_name: fehiepsi/arviz

reason: import numpy,import scipy


text:
import gengraph
import random
import torch_geometric
from utils import featgen
import numpy as np
import utils.io_utils as io_utils
from configs import arg_parse
import torch
import torch.nn as nn
from torch.autograd import Variable
from models_pyg import GCNNet
import os
from torch_geometric.utils import from_networkx
from tensorboardX import SummaryWriter

def test(loader, model, args, labels, test_mask):
    model.eval()

    train_ratio = args.train_ratio
    correct = 0
    for data in loader:
        with torch.no_grad():
            pred = model(data)
            # print ('pred:', pred)
            pred = pred.argmax(dim=1)
            # print ('pred:', pred)

        # node classification: only evaluate on nodes in test set
        pred = pred[test_mask]
        # print ('pred:', pred)
        label = labels[test_mask]
        # print ('label:', label)

        correct += pred.eq(label).sum().item()

    total = len(test_mask)
    # print ('correct:', correct)
    return correct / total

def syn_task1(args, writer=None):
    # data
    print ('Generating graph.')
    G, labels, name = gengraph.gen_syn1(
            feature_generator=featgen.ConstFeatureGen(np.ones(args.input_dim, dtype=float)))
    # print ('G.node[0]:', G.node[0]['feat'].dtype)
    # print ('Original labels:', labels)
    pyg_G = from_networkx(G)
    num_classes = max(labels)+1
    labels = torch.LongTensor(labels)
    print ('Done generating graph.')

    # if args.method == 'att':
    # print('Method: att')
    # model = models.GcnEncoderNode(args.input_dim, args.hidden_dim, args.output_dim, num_classes,
    # args.num_gc_layers, bn=args.bn, args=args)

    # else:
    # print('Method:', args.method)
    # model = models.GcnEncoderNode(args.input_dim, args.hidden_dim, args.output_dim, num_classes,
    # args.num_gc_layers, bn=args.bn, args=args)

    model = GCNNet(args.input_dim, args.hidden_dim, num_classes, args.num_gc_layers, args=args)

    if args.gpu:
        model = model.cuda()

    train_ratio = args.train_ratio
    num_train = int(train_ratio * G.number_of_nodes())
    num_test = G.number_of_nodes() - num_train
    shuffle_indices = list(range(G.number_of_nodes()))
    shuffle_indices = np.random.permutation(shuffle_indices)

    train_mask = num_train * [True] + num_test * [False]
    train_mask = torch.BoolTensor([train_mask[i] for i in shuffle_indices])
    test_mask = num_train * [False] + num_test * [True]
    test_mask = torch.BoolTensor([test_mask[i] for i in shuffle_indices])

    loader = torch_geometric.data.DataLoader([pyg_G], batch_size=1)
    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
    for epoch in range(args.num_epochs):
        total_loss = 0
        model.train()
        for batch in loader:
            # print ('batch:', batch.feat)
            opt.zero_grad()
            pred = model(batch)

            pred = pred[train_mask]
            # print ('pred:', pred)
            label = labels[train_mask]
            # print ('label:', label)
            loss = model.loss(pred, label)
            print ('loss:', loss)
            loss.backward()
            opt.step()
            total_loss += loss.item() * 1
        total_loss /= num_train
        writer.add_scalar("loss", total_loss, epoch)

        if epoch % 10 == 0:
            test_acc = test(loader, model, args, labels, test_mask)
            print("Epoch {}. Loss: {:.4f}. Test accuracy: {:.4f}".format(
                epoch, total_loss, test_acc))
            writer.add_scalar("test accuracy", test_acc, epoch)

prog_args = arg_parse()
path = os.path.join(prog_args.logdir, io_utils.gen_prefix(prog_args))
syn_task1(prog_args, writer=SummaryWriter(path))

ext: py

reason: import numpy


text:
import os

import cv2
import imgaug as ia
import numpy as np
import scipy
import torch
import torch.multiprocessing as mp
import torch.nn as nn
from PIL import Image
from torchvision import transforms
from imgaug import augmenters as iaa

from . import utils
from .modeling import deeplab, deeplab_masks

class InferenceOutlines():

    def __init__(self, outlinesWeightsFile, outlinesModel='deeplab_resnet', imgHeight=288, imgWidth=512):
        '''Class to run Inference of the Outlines Prediction Model

        Args:
            outlinesWeightsFile (str): Path to the weights file for the model
            outlinesModel (str): Which model to use. Can be one of ["unet", "deeplab_resnet", "drn"]
            imgHeight (int): The height to which images should be resized to before passing to the model.
            imgWidth (int): The width to which images should be resized to before passing to the model.
        '''

        # Create Model
        if outlinesModel == 'deeplab_resnet':
            print('Creating Deeplabv3-Resnet model for outlines and loading checkpoint')
            self.model = deeplab_masks.DeepLab(num_classes=3, backbone='resnet', sync_bn=True,
                            freeze_bn=False)
        elif outlinesModel == 'drn':
            print('Creating DRN model for outlines and loading checkpoint')
            self.model = deeplab_masks.DeepLab(num_classes=3, backbone='drn', sync_bn=True,
                                               freeze_bn=True)  # output stride is 8 for drn
        else:
            raise NotImplementedError(
                'Model may only be "drn", "deeplab_resnet" or unet". Given model is: {}'.format(outlinesModel))

        # Load Checkpoint and its data (weights file)
        if not os.path.isfile(outlinesWeightsFile):
            raise ValueError('Invalid path to the given weights file in config. The file "{}" does not exist'.format(
                outlinesWeightsFile))

        checkpoint = torch.load(outlinesWeightsFile, map_location='cpu')
        if 'model_state_dict' in checkpoint:
            # Newer checkpoints have multiple dicts within
            self.model.load_state_dict(checkpoint['model_state_dict'])
        else:
            self.model.load_state_dict(checkpoint)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.model.to(self.device)
        self.model.eval()

        # Create Transforms
        self.transform = iaa.Sequential([
            iaa.Resize({
                "height": imgHeight,
                "width": imgWidth
            }, interpolation='nearest'),
        ])

    def runOnNumpyImage(self, img, scaling_factor_weights=3):
        '''Runs inference of Outlines prediction model on a Numpy Image

        Args:
            img (numpy.ndarray): shape=(H, W, 3)
                                 if dtype=np.float32, Expected data in range (0, 1)
                                 if dtype=np.uint8, Expected data in range (0, 255)
                                 Obtained normally by using PIL to open image (in mode RGB) and converting to numpy.
            scaling_factor_weights (int): The power to which outputs of outlines are raised to before scaling to
                                          range 0-1000. The default value, as used by Yinda, is 3.

        Returns:
            numpy.ndarray: Occlusion Weights derived from outlines prediction. Used by depth2depth module.
                           shape=(H, W), dtype=np.uint16
            numpy.ndarray: RGB Visualization of occlusion Weights.
                           shape=(H, W, 3), dtype=np.uint8
            numpy.ndarray: Predicted Outlines RGB image. Value of each color corresponds to the class it belongs to:
                                Red: Background, Green: Depth Boundary, Blue: Surface Normal Gradient
                           shape=(H, W, 3), dtype=np.uint8.
        '''
        with torch.no_grad():
            # Resize Image and convert to tensor
            det_tf = self.transform.to_deterministic()
            img = det_tf.augment_image(img)
            inputs = transforms.ToTensor()(img)
            inputs = torch.unsqueeze(inputs, 0)
            inputs = inputs.to(self.device)

            outputs = self.model(inputs)
            predictions = torch.max(outputs, 1)[1]

            # Generate and Save Occlusion Weights File used by depth2depth
            # calculating occlusion weights
            SCALING_FACTOR_PWR = scaling_factor_weights
            SCALING_FACTOR_MUL = 1000
            output_softmax = nn.Softmax(dim=1)(outputs).squeeze(0).cpu().numpy()
            # TEST - REMOVING UNCERTAIN VALUES FROM OUTLINE WEIGHTS
            # output_softmax[output_softmax < 0.3] = 1e-4
            weight = (1 - output_softmax[1, :, :])  # Occlusion (depth) boundaries is channel 1
            x = np.power(weight, SCALING_FACTOR_PWR)
            x = np.multiply(x, SCALING_FACTOR_MUL)
            occlusion_boundary_weight = x.astype(np.uint16)
            # Increase the min and max values by small amount epsilon so that absolute min/max values
            # don't cause problems in the depth2depth optimization code.
            epsilon = 1
            occlusion_boundary_weight[occlusion_boundary_weight == 0] += epsilon
            occlusion_boundary_weight[occlusion_boundary_weight == SCALING_FACTOR_MUL] -= epsilon

            # Convert prediction to RGB for display (R: background, G: occlusion boundary, B: normals gradient)
            predictions = predictions.squeeze(0).cpu().numpy()
            output_rgb = np.zeros((predictions.shape[0], predictions.shape[1], 3), dtype=np.uint8)
            output_rgb[:, :, 0][predictions == 0] = 255
            output_rgb[:, :, 1][predictions == 1] = 255
            output_rgb[:, :, 2][predictions == 2] = 255

            # Create RGB Visualization of Occlusion Weights
            NORM_FACTOR_UINT8 = SCALING_FACTOR_MUL / 255.0
            occlusion_boundary_weight_rgb = (occlusion_boundary_weight / NORM_FACTOR_UINT8).astype(np.uint8)
            occlusion_boundary_weight_rgb = cv2.applyColorMap(occlusion_boundary_weight_rgb, cv2.COLORMAP_OCEAN)
            occlusion_boundary_weight_rgb = cv2.cvtColor(occlusion_boundary_weight_rgb, cv2.COLOR_BGR2RGB)

        return occlusion_boundary_weight, occlusion_boundary_weight_rgb, output_rgb

class InferenceNormals():

    def __init__(self, normalsWeightsFile, normalsModel='deeplab_resnet', imgHeight=288, imgWidth=512):
        '''Class to run Inference of the normals Prediction Model

        Args:
            normalsWeightsFile (str): Path to the weights file for the model
            normalsModel (str): Which model to use. Can be one of ["unet", "deeplab_resnet", "drn"]
            imgHeight (int): The height to which images should be resized to before passing to the model.
            imgWidth (int): The width to which images should be resized to before passing to the model.
        '''

        # Create Model
        if normalsModel == 'deeplab_resnet':
            print('Creating Deeplabv3-Resnet model for surface normals and loading checkpoint')
            self.model = deeplab.DeepLab(num_classes=3, backbone='resnet', sync_bn=True,
                                         freeze_bn=False)
        elif normalsModel == 'drn':
            print('Creating DRN model for normals and loading checkpoint')
            self.model = deeplab.DeepLab(num_classes=3, backbone='drn', sync_bn=True,
                                         freeze_bn=False)  # output stride is 8 for drn
        else:
            raise NotImplementedError(
                'Model may only be "unet" or "deeplab_resnet". Given model is: {}'.format(normalsModel))

        # Load Checkpoint and its data (weights file)
        if not os.path.isfile(normalsWeightsFile):
            raise ValueError('Invalid path to the given weights file in config. The file "{}" does not exist'.format(
                normalsWeightsFile))

        checkpoint = torch.load(normalsWeightsFile, map_location='cpu')
        if 'model_state_dict' in checkpoint:
            # Newer checkpoints have multiple dicts within
            # checkpoint['model_state_dict'].pop('decoder.last_conv.8.weight')
            # checkpoint['model_state_dict'].pop('decoder.last_conv.8.bias')
            self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)
        elif 'state_dict' in checkpoint:
            checkpoint['state_dict'].pop('decoder.last_conv.8.weight')
            checkpoint['state_dict'].pop('decoder.last_conv.8.bias')
            self.model.load_state_dict(checkpoint['state_dict'], strict=False)
        else:
            self.model.load_state_dict(checkpoint)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.model.to(self.device)
        self.model.eval()

        # Create Transforms
        self.transform = iaa.Sequential([
            iaa.Resize({
                "height": imgHeight,
                "width": imgWidth
            }, interpolation='nearest'),
        ])

    def _normal_to_rgb(self, normals_to_convert):
        '''Converts a surface normals array into an RGB image.
        Surface normals are represented in a range of (-1,1),
        This is converted to a range of (0,255) to be written
        into an image.
        The surface normals are normally in camera co-ords,
        with positive z axis coming out of the page. And the axes are
        mapped as (x,y,z) -> (R,G,B).
        '''
        camera_normal_rgb = normals_to_convert + 1
        camera_normal_rgb *= 127.5
        camera_normal_rgb = camera_normal_rgb.astype(np.uint8)
        return camera_normal_rgb

    def _convert_normals_depth2depth_format(self, normals):
        """This function converts surface normals from our representation to the representation used by depth2depth

        When displaying normals, RGB represents the axes XYZ. Using that demonination, here are the representations
        used by our method and theirs:
          - Ours:        R: Right,  G: Up,               B: Towards Camera  (Y-up notation)
          - depth2depth: R: Right,  G: Away from Camera, B: Up              (Z-up notation)
        Args:
            normals (numpy.ndarray): Surface normals. Shape=(3, H, W), dtype=float32
        """
        if (len(normals.shape) != 3) or (normals.shape[0] != 3):
            raise ValueError('Shape of normals should be (3, H, W). Got shape: {}'.format(normals.shape))

        # Convert normals to shape (N, 3)
        normals_list = np.reshape(normals, (3, -1)).transpose((1, 0))

        # Apply Rotation
        r = scipy.spatial.transform.Rotation.from_euler('x', 90, degrees=True)
        normals_list_rotated = r.apply(normals_list)

        # Convert normals back to shape (3, H, W)
        normals_depth2depth_format = np.reshape(normals_list_rotated.transpose(1, 0), normals.shape)

        return normals_depth2depth_format

    def runOnNumpyImage(self, img):
        '''Runs inference of Outlines prediction model on a Numpy Image

        Args:
            img (numpy.ndarray): shape=(H, W, 3)
                                 if dtype=np.float32, Expected data in range (0, 1)
                                 if dtype=np.uint8, Expected data in range (0, 255)
                                 Obtained normally by using PIL to open image (in mode RGB) and converting to numpy.

        Returns:
            numpy.ndarray: Predicted Surface Normals
                           shape=(3, H, W), dtype=np.float32 with range (-1, 1)
            numpy.ndarray: RGB representation of predicted Surface Normals. (R, G, B) represent (X, Y, Z) axes.
        '''
        with torch.no_grad():
            # Resize Image and convert to tensor
            det_tf = self.transform.to_deterministic()
            img = det_tf.augment_image(img)
            inputs = transforms.ToTensor()(img)
            inputs = torch.unsqueeze(inputs, 0)
            inputs = inputs.to(self.device)

            outputs = self.model(inputs)

            outputs_norm = nn.functional.normalize(outputs, p=2, dim=1)
            surface_normals = outputs_norm.squeeze(0).cpu().numpy()

            # Create RGB Viz of Normals
            surface_normals_rgb = self._normal_to_rgb(surface_normals.transpose((1, 2, 0)))

            # Rotate Surface Normals into depth2depth notation
            surface_normals_rotated = self._convert_normals_depth2depth_format(surface_normals)

            # # Create RGB Viz of Rotated Normals
            # surface_normals_rgb = self._normal_to_rgb(surface_normals_rotated.transpose((1, 2, 0)))

        return surface_normals_rotated, surface_normals_rgb

class InferenceMasks():

    def __init__(self, masksWeightsFile, masksModel='drn', imgHeight=288, imgWidth=512):
        '''Class to run Inference of the Outlines Prediction Model

        Args:
            masksWeightsFile (str): Path to the weights file for the model
            masksModel (str): Which model to use. Can be one of ["deeplab_xception", "deeplab_resnet", "drn"]
            imgHeight (int): The height to which images should be resized to before passing to the model.
            imgWidth (int): The width to which images should be resized to before passing to the model.
        '''

        # Create Model
        if masksModel == 'deeplab_resnet':
            print('Creating Deeplabv3-Resnet model for masks and loading checkpoint')
            self.model = deeplab_masks.DeepLab(num_classes=2, backbone='resnet', sync_bn=True,
                                               freeze_bn=False)
        elif masksModel == 'deeplab_xception':
            self.model = deeplab_masks.DeepLab(num_classes=2, backbone='xception', sync_bn=True,
                                               freeze_bn=False)
        elif masksModel == 'drn':
            print('Creating DRN model for masks and loading checkpoint')
            self.model = deeplab_masks.DeepLab(num_classes=2, backbone='drn', sync_bn=True,
                                               freeze_bn=False)  # output stride is 8 for drn
        else:
            raise NotImplementedError(
                'Model may only be "unet" or "deeplab_resnet". Given model is: {}'.format(masksModel))

        # Load Checkpoint and its data (weights file)
        if not os.path.isfile(masksWeightsFile):
            raise ValueError('Invalid path to the given weights file in config. The file "{}" does not exist'.format(
                masksWeightsFile))

        checkpoint = torch.load(masksWeightsFile, map_location='cpu')
        if 'model_state_dict' in checkpoint:
            # Newer checkpoints have multiple dicts within
            self.model.load_state_dict(checkpoint['model_state_dict'], strict=True)
        elif 'state_dict' in checkpoint:
            checkpoint['state_dict'].pop('decoder.last_conv.8.weight')
            checkpoint['state_dict'].pop('decoder.last_conv.8.bias')
            self.model.load_state_dict(checkpoint['state_dict'], strict=False)
        else:
            raise RuntimeError('Invalid Checkpoint. It does not contain "model_state_dict" in it:\n{}'.format(
                checkpoint.keys()))

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.model.to(self.device)
        self.model.eval()

        # Create Transforms
        self.transform = iaa.Sequential([
            iaa.Resize({
                "height": imgHeight,
                "width": imgWidth
            }, interpolation='nearest'),
        ])

    def runOnNumpyImage(self, img, dilate_mask=False):
        '''Runs inference of Outlines prediction model on a Numpy Image

        Args:
            img (numpy.ndarray): shape=(H, W, 3)
                                 if dtype=np.float32, Expected data in range (0, 1)
                                 if dtype=np.uint8, Expected data in range (0, 255)
                                 Obtained normally by using PIL to open image (in mode RGB) and converting to numpy.

        Returns:
            numpy.ndarray: Binary mask of detected transparent objects. Used to mask input depth.
                           shape=(H, W), dtype=np.uint8, True pixels have value of 255
        '''
        with torch.no_grad():
            # Resize Image and convert to tensor
            det_tf = self.transform.to_deterministic()
            img = det_tf.augment_image(img)
            inputs = transforms.ToTensor()(img)
            inputs = torch.unsqueeze(inputs, 0)
            inputs = inputs.to(self.device)

            outputs = self.model(inputs)
            predictions = torch.max(outputs, 1)[1]
            predictions = predictions.squeeze(0).cpu().numpy()

            mask = np.zeros(predictions.shape, dtype=np.uint8)
            mask[predictions == 1] = 255

            if dilate_mask is True:
                kernel = np.ones((5, 5), np.uint8)
                mask = cv2.dilate(mask, kernel, iterations=1)

        return mask

max_stars_repo_name: AriaPs/cleargrasp

max_issues_repo_name: AriaPs/cleargrasp

max_forks_repo_name: AriaPs/cleargrasp


text:
from itertools import product

import numpy as np
import torch

from alphazero.game import Game

class TicTacToe(Game):
    """Squares are numbered 0-8 from top-left corner to bottom right corner."""

    NUM_ACTIONS = 9

    WINNING_POSITIONS = [
        [0, 1, 2],
        [3, 4, 5],
        [6, 7, 8],
        [0, 3, 6],
        [1, 4, 7],
        [2, 5, 8],
        [0, 4, 8],
        [2, 4, 6]
    ]

    @property
    def terminal(self):
        if len(self.history) < 5:
            return False
        elif len(self.history) == 9:
            return True
        elif self.winner is not None:
            return True
        return False

    @property
    def winner(self):
        if len(self.history) < 5:
            return None
        for player, pos in product((0, 1), TicTacToe.WINNING_POSITIONS):
            if len(set(self.history[player::2]) & set(pos)) == 3:
                return player
        if len(self.history) == 9:
            return -1
        return None

    @property
    def valid_actions(self):
        return sorted(set(range(9)) - set(self.history))

    def render(self, turn=None):
        if turn is None:
            history = self.history
        else:
            assert turn <= len(self.history)
            assert turn >= 0
            history = self.history[:turn]
        img = np.zeros((2, 9), dtype=np.float32)
        next_player_actions = history[self.next_player::2]
        prev_player_actions = history[self.prev_player::2]
        img[0, next_player_actions] = 1
        img[1, prev_player_actions] = 1
        return img.reshape(2, 3, 3)

    def __str__(self):
        h = self.history
        tui = [' ' if i not in h else 'X' if h.index(i) % 2 == 0 else 'O' for i in range(9)]
        s  = ' ' + ' | '.join(tui[0:3]) + '\n'
        s += '-' * 11 + '\n'
        s += ' ' + ' | '.join(tui[3:6]) + '\n'
        s += '-' * 11 + '\n'
        s += ' ' + ' | '.join(tui[6:9])
        return s

    def solve(self):
        BOOK = {
            ():   (0, [ 0,  0,  0,  0,  0,  0,  0,  0,  0]),
            (0,): (0, [    -2, -2, -2,  0, -2, -2, -2, -2]),
            (1,): (0, [ 0,      0, -2,  0, -2, -2,  0, -2]),
            (2,): (0, [-2, -2,     -2,  0, -2, -2, -2, -2]),
            (3,): (0, [ 0, -2, -2,      0,  0,  0, -2, -2]),
            (4,): (0, [ 0, -2,  0, -2,     -2,  0, -2,  0]),
            (5,): (0, [-2, -2,  0,  0,  0,     -2, -2,  0]),
            (6,): (0, [-2, -2, -2, -2,  0, -2,     -2, -2]),
            (7,): (0, [-2,  0, -2, -2,  0, -2,  0,      0]),
            (8,): (0, [-2, -2, -2, -2,  0, -2, -2, -2    ]),
        }
        def minimax(game, root_depth=None, book={}):
            if root_depth is None:
                root_depth = len(game)

            if tuple(game.history) in book:
                return book[tuple(game.history)]

            root_player = root_depth % 2
            cur_player = len(game) % 2

            if game.terminal:
                nmoves = 5 - len(game) // 2
                outcome = game.terminal_value(root_player)
                return nmoves * outcome, []

            scores = []
            for a in game.valid_actions:
                sub_game = game.clone()
                sub_game.take_action(a)
                s, _ = minimax(sub_game, root_depth, book)
                scores.append(s)

            if cur_player == root_player:
                s = max(scores)
            else:
                s = min(scores)
            return s, scores

        s, scores = minimax(self, book=BOOK)
        v = 1 if s > 0 else -1 if s < 0 else 0
        return scores, v


text:
import numpy as np
import polymath as pm

from pathlib import Path
CWD = Path(f"{__file__}").parent
import locale
FLIST_FILE = "filelist_4"
REC_LENGTH = 49
REC_WINDOW = 10
LATITUDE_POS = 28
OPEN = 10000

def np_nn(input_info, file_lengths):
    out_info = {}

    out_info['z'] = np.zeros((REC_WINDOW,))
    target_lat, target_long = input_info['target_lat'], input_info['target_long']
    out_info['entries'] = input_info['entries'].copy()
    out_info['ndists'] = input_info['ndists'].copy()

    start = 0
    input_info['lats'] = np.empty(file_lengths[-1][-1], dtype=np.float64)
    input_info['longs'] = np.empty(file_lengths[-1][-1], dtype=np.float64)
    out_info['first'] = input_info['ndists'].copy()
    while (start + REC_WINDOW) < file_lengths[-1][-1]:

        sandbox = input_info['sandbox'][start: start + REC_WINDOW]

        orig_data = input_info['orig_data'][start: start + REC_WINDOW]
        start += REC_WINDOW

        for i in range(REC_WINDOW):
            if (start - REC_WINDOW) == 0:
                input_info['lats'][i] = sandbox[i][0]
                input_info['longs'][i] = sandbox[i][1]
            tmp_lat = sandbox[i][0]
            tmp_long = sandbox[i][1]
            out_info['z'][i] = np.sqrt(( (tmp_lat-target_lat) * (tmp_lat-target_lat) )+( (tmp_long-target_long) * (tmp_long-target_long)))

        for i in range(REC_WINDOW):
            max_dist = -1
            max_idx = 0
            for j in range(input_info['num']):
                if out_info['ndists'][j] > max_dist:
                    max_dist = out_info['ndists'][j]
                    max_idx = j
    #
            if out_info['z'][i] < out_info['ndists'][max_idx]:
    #
                out_info['entries'][max_idx] = orig_data[i]
                out_info['ndists'][max_idx] = out_info['z'][i]
                if (start - REC_WINDOW) == 0:
                    out_info['first'][max_idx] = out_info['z'][i]
                    print(f"Here: {out_info['ndists']}\n")

    #
    input_info['sandbox'] = input_info['sandbox'][0:REC_WINDOW]
    input_info['orig_data'] = input_info['orig_data'][0:REC_WINDOW]
    input_info['longs'] = input_info['longs'][0:REC_WINDOW]
    input_info['lats'] = input_info['lats'][0:REC_WINDOW]
    input_info.pop("sandbox")
    # for j in range(input_info['num']):
    #     if out_info['ndists'][j] != OPEN:
    #         print(f"{out_info['entries'][j]} --> {out_info['ndists'][j]}")
    return out_info

def read_filelist(listfile):
    with open(f"{listfile}", "r") as f:
        lines = f.readlines()
    lines = [f"{CWD}/data/nn/{Path(l.strip()).name}" for l in lines]
    return lines

def nn_datagen(flist_file, neighbors, tgt_lat, tgt_long, rec_len, rec_window, lowered=False, inf=False):
    full_paths = read_filelist(flist_file)
    input_info = {}
    data = []
    file_lengths = []
    start = 0
    for fname in full_paths:
        with open(fname, "r") as f:
            fdata = f.readlines()
            data += fdata
            file_lengths.append((start, start+len(fdata)))
            start = start+len(fdata) + 1
    assert all([len(d) == rec_len for d in data])
    input_info['orig_data'] = data
    np_data = np.empty((len(data), 2), dtype=np.float64)
    for i, d in enumerate(data):
        lat = float(d[LATITUDE_POS-1:LATITUDE_POS+4])
        long = float(d[LATITUDE_POS+5:LATITUDE_POS+10])
        np_data[i][0] = lat
        np_data[i][1] = long

    input_info['sandbox'] = np_data
    input_info['target_lat'] = tgt_lat
    input_info['target_long'] = tgt_long
    input_info['num'] = neighbors
    input_info['entries'] = np.full((neighbors), '', dtype=np.dtype('<U48'))

    input_info['ndists'] = np.full(neighbors, OPEN, dtype=np.float64)

    out_info = np_nn(input_info, file_lengths)

    input_info.pop("entries")
    input_info.pop("orig_data")
    input_info.pop("num")

    if lowered:
        all_keys = []
        for i in range(rec_window):
            ndist_key = f"ndists/ndists({i},)"
            all_keys.append(ndist_key)
            for j in range(rec_len):
                entry_key = f"entries/entries({i}, {j})"
                all_keys.append(entry_key)
                input_info[entry_key] = input_info["entries"][i,j]
            input_info[ndist_key] = input_info["ndists"][i]
        input_info.pop("entries")
        input_info.pop("ndists")
    else:
        all_keys = ["ndists", "entries"]

    return input_info, all_keys, out_info

def nn_impl(neighbors, latitude,longitude,  coarse=False):

    with pm.Node(name="nn") as graph:
        # num = neighbors
        num = pm.parameter("num")
        target_lat = pm.parameter("target_lat")
        target_long = pm.parameter("target_long")

        lats = pm.input("lats", shape=(REC_WINDOW))
        longs = pm.input("longs", shape=(REC_WINDOW))

        # Initialize to OPEN
        ndists = pm.state("ndists", shape=(num))
        # entries = pm.output("entries", shape=(num))

        i = pm.index(0, REC_WINDOW - 1, name="i")
        j = pm.index(0, num - 1, name="j")
        # tmp_lats = sandbox[i, 0]
        # tmp_longs = sandbox[i, 1]

        z = pm.sqrt(((lats[i]-target_lat) * (lats[i]-target_lat)) + ((longs[i]-target_long) * (longs[i]-target_long)), name="sqrtz")

        # ndists[j] = OPEN

        for idx in range(REC_WINDOW):
            max_dist = pm.max([j], ndists[j], name=f"max_dist{idx}")[0]
            max_idx = (pm.argmax([j], ndists[j])).set_name(f"max_idx{idx}")
            ndists[max_idx] = (((z[idx] < max_dist) * z[idx]) + ((z[idx] >= max_dist) * ndists[max_idx]))

    if coarse:
        in_info, keys, out_info = nn_datagen(f"{CWD}/openmp/nn/filelist_4", neighbors, latitude, longitude,
                                                    REC_LENGTH, REC_WINDOW)

        return graph, in_info, out_info, keys
    else:
        in_info, keys, out_info = nn_datagen(f"{CWD}/openmp/nn/filelist_4", neighbors, latitude, longitude,
                                                    REC_LENGTH, REC_WINDOW, lowered=True)
        shape_val_pass = pm.NormalizeGraph({"num": in_info['num']})
        new_graph = shape_val_pass(graph)
        return new_graph, in_info, out_info, keys

max_stars_repo_name: lite-david/polymath

max_issues_repo_name: lite-david/polymath

max_forks_repo_name: lite-david/polymath


text:
## Cross-correlations of contact times
"""
2D	PLOT_XCORR_MEANED_OVER_WHISKER
    STATS__PLOT_XCORR_MEANED_OVER_WHISKER
    Correlograms of contact times within and across whiskers
"""
import json
import os
import numpy as np
import pandas
import matplotlib.pyplot as plt
import whiskvid
import my
import kkpandas
import my.plot
import matplotlib

## Fonts
my.plot.manuscript_defaults()
my.plot.font_embed()
DK = chr(916) + chr(954)

## Parameters
with open('../parameters') as fi:
    params = json.load(fi)

## Load metadata about sessions
session_df, task2mouse, mouse2task = my.dataload.load_session_metadata(params)

## Load data
# patterns data
big_tm = my.dataload.load_data_from_patterns(
    params, 'big_tm')
C2_whisk_cycles = my.dataload.load_data_from_patterns(
    params, 'big_C2_tip_whisk_cycles')

# Load the contact results
# These have to be manually sliced to remove opto trials
big_ccs_df = pandas.read_pickle(
    os.path.join(params['patterns_dir'], 'big_ccs_df'))

## Use big_tm to slice big_ccs_df
# Get included_trials
included_trials = big_tm.index

# Use those cycles to slice big_ccs_df
big_ccs_df = big_ccs_df.set_index(
    ['trial', 'whisker', 'cycle'], append=True).reorder_levels(
    ['session', 'trial', 'cycle', 'whisker', 'cluster']).sort_index()
big_ccs_df = my.misc.slice_df_by_some_levels(
    big_ccs_df, C2_whisk_cycles.index)

# Return big_ccs_df to its original index
big_ccs_df = big_ccs_df.reset_index().set_index(['session', 'cluster'])

## Stack onsets
session_l = sorted(np.unique(big_ccs_df.index.get_level_values('session')))
whisker_l = ['C1', 'C2', 'C3']
count_l = []
count_keys_l = []
for session in session_l:
    # Get contacts for this session
    session_ccs = big_ccs_df.loc[session]

    # Iterate over whiskers
    for nw0 in range(len(whisker_l)):
        for nw1 in range(len(whisker_l)):
            # Get whiskers
            w0 = whisker_l[nw0]
            w1 = whisker_l[nw1]
            auto = w0 == w1

            # Get contact
            w0_frames = session_ccs.loc[
                session_ccs['whisker'] == w0, 'frame_start'].copy()
            w1_frames = session_ccs.loc[
                session_ccs['whisker'] == w1, 'frame_start'].copy()

            # Skip if nothing
            if len(w0_frames) == 0:
                continue
            if len(w1_frames) == 0:
                continue

            # Jitter to break ties
            w0_frames += 0.1 * np.random.standard_normal(size=w0_frames.size)
            w1_frames += 0.1 * np.random.standard_normal(size=w1_frames.size)

            # Correlate
            count, bins = kkpandas.utility.correlogram(
                w0_frames / 200.,
                w1_frames / 200.,
                bin_width=.005, limit=0.15, auto=auto,
                )

            # Normalize: This was calculating by folding on the spikes in
            # the shorter signal, so divide by the length of the shorter signal
            # to make it into an expectation of the number of spikes per bin.
            rate = count / float(np.min([len(w0_frames), len(w1_frames)]))

            # Store
            count_l.append(rate)
            count_keys_l.append(
                (session, 'auto' if auto else 'cross', w0, w1))

## Bincenters
bincenters = (bins[:-1] + bins[1:]) / 2.0

## Concat the correlograms
corrdf = pandas.DataFrame(
    count_l,
    index=pandas.MultiIndex.from_tuples(
    count_keys_l, names=['session', 'typ', 'w0', 'w1']),
    columns=bincenters,
    )

# Add task level
corrdf = my.misc.insert_mouse_and_task_levels(corrdf, mouse2task)

# Drop detection because so few contacts
# Results for detection are actually basically the same, but there's some
# missing N because of the small number of contacts, so it's annoying
corrdf = corrdf.loc['discrimination']

## Aggregate
# Mean across sessions
mouse_mcorrdf = corrdf.mean(
    level=[lev for lev in corrdf.index.names if lev != 'session'])

# Slice out auto
mouse_auto_mcorrdf = mouse_mcorrdf.xs('auto', level='typ').droplevel('w1')

# Slice out cross, and drop non-adjacent whiskers
mouse_cross_mcorrdf = mouse_mcorrdf.xs('cross', level='typ')
idxdf = mouse_cross_mcorrdf.index.to_frame()
drop_mask = (
    (mouse_cross_mcorrdf.index.get_level_values('w0') == 'C1') &
    (mouse_cross_mcorrdf.index.get_level_values('w1') == 'C3')
    )
mouse_cross_mcorrdf = mouse_cross_mcorrdf.loc[~drop_mask]
drop_mask = (
    (mouse_cross_mcorrdf.index.get_level_values('w0') == 'C3') &
    (mouse_cross_mcorrdf.index.get_level_values('w1') == 'C1')
    )
mouse_cross_mcorrdf = mouse_cross_mcorrdf.loc[~drop_mask]

# Mean across mice, leaving whisker
whisker_auto_mcorrdf = mouse_auto_mcorrdf.mean(level='w0')
whisker_cross_mcorrdf = mouse_cross_mcorrdf.mean(level=['w0', 'w1'])

# Grand mean first across whisker, leaving mice; then mean and sem over mice
# Doing all pairs of whiskers here to make it symmetric
to_agg_auto_mcorrdf = mouse_auto_mcorrdf.mean(level='mouse')
to_agg_cross_mcorrdf = mouse_cross_mcorrdf.mean(level='mouse')
agg_auto_mcorrdf = to_agg_auto_mcorrdf.mean()
agg_cross_mcorrdf = to_agg_cross_mcorrdf.mean()
aggerr_auto_mcorrdf = to_agg_auto_mcorrdf.sem()
aggerr_cross_mcorrdf = to_agg_cross_mcorrdf.sem()

## Plots
PLOT_XCORR_MEANED_OVER_WHISKER = True

if PLOT_XCORR_MEANED_OVER_WHISKER:
    ## Plot auto and cross meaned over whiskers
    f, ax = plt.subplots(figsize=(2.5, 2.25))
    f.subplots_adjust(bottom=.25, left=.27, right=.97, top=.9)

    # Plot auto
    ax.plot(agg_auto_mcorrdf, color='k', lw=1, clip_on=False)
    ax.fill_between(
        x=agg_auto_mcorrdf.index.values,
        y1=(agg_auto_mcorrdf - aggerr_auto_mcorrdf),
        y2=(agg_auto_mcorrdf + aggerr_auto_mcorrdf),
        color='k', alpha=.2, lw=0,
        )

    # Plot cross
    ax.plot(agg_cross_mcorrdf, color='k', lw=1, ls='--',clip_on=False)
    ax.fill_between(
        x=agg_cross_mcorrdf.index.values,
        y1=(agg_cross_mcorrdf - aggerr_cross_mcorrdf),
        y2=(agg_cross_mcorrdf + aggerr_cross_mcorrdf),
        color='k', alpha=.2, lw=0,
        )

    # Pretty
    my.plot.despine(ax)
    ax.set_xlabel('{}t (s)'.format(chr(916)))
    ax.set_ylabel('Pr(contact)')
    ax.set_xlim((-.15, .15))
    ax.set_xticks((-.1, 0, .1))
    ax.set_ylim((0, .2))
    ax.set_yticks((0, .1, .2))

    # Legend
    ax.plot([-.022, .0], [.22, .22], 'k--', lw=1, clip_on=False)
    ax.plot([-.022, .0], [.195, .195], 'k-', lw=1, clip_on=False)
    ax.text(.02, .22, 'across whisker', size=12, ha='left', va='center')
    ax.text(.02, .195, 'within whisker', size=12, ha='left', va='center')

    # Save
    f.savefig('PLOT_XCORR_MEANED_OVER_WHISKER.svg')
    f.savefig('PLOT_XCORR_MEANED_OVER_WHISKER.png', dpi=300)

    # Stats
    with open('STATS__PLOT_XCORR_MEANED_OVER_WHISKER', 'w') as fi:
        # Check same for all
        assert to_agg_auto_mcorrdf.index.equals(to_agg_cross_mcorrdf.index)

        fi.write('STATS__PLOT_XCORR_MEANED_OVER_WHISKER\n')
        fi.write('n = {} mice (all disc.)\n'.format(len(to_agg_auto_mcorrdf)))
        fi.write('error bars: sem\n')

    with open('STATS__PLOT_XCORR_MEANED_OVER_WHISKER', 'r') as fi:
        lines = fi.readlines()
    print(''.join(lines))

plt.show()


text:
"""
    Admin load feat files (.npz one for each model) and compute model similarity
"""
import argparse
import numpy as np
import os
import sys
from numpy.linalg import norm

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('feat_files', nargs='+',
                        help='path to one or multiple features (.npz file) '
                        'by one or multiple models')
    parser.add_argument('--save', required=True, type=str,
                        help='name of or path to a .npy of similarity matrix, '
                        'and a .txt of a list of feature files')
    return parser.parse_args()

def KA(feat1, feat2, remove_mean=True):
    """
        feat1, feat2: n x d
    """
    if remove_mean:
        feat1 -= np.mean(feat1, axis=0, keepdims=1)
        feat2 -= np.mean(feat2, axis=0, keepdims=1)
    norm12 = norm(feat1.T.dot(feat2))**2
    norm11 = norm(feat1.T.dot(feat1))
    norm22 = norm(feat2.T.dot(feat2))
    return norm12 / (norm11 * norm22)

def compute_sim(feat_files):
    N = len(feat_files)
    sim = np.eye(N)
    for i in range(N):
        feat_i = np.load(feat_files[i])
        for j in range(i+1, N):
            feat_j = np.load(feat_files[j])
            sim[i, j] = KA(feat_i, feat_j, remove_mean=True)
            sim[j, i] = sim[i, j]
            print(i, j, sim[i, j], flush=True)
    return sim

if __name__ == '__main__':
    args = parse_args()
    sim = compute_sim(args.feat_files)
    np.save(args.save+'.npy', sim)
    with open(args.save+'.txt', 'w') as f:
        f.write('\n'.join(args.feat_files))

reason: import numpy,from numpy


text:
# -*- coding: utf-8 -*-

"""Copyright 2015 Roger R Labbe Jr.

FilterPy library.
http://github.com/rlabbe/filterpy

Documentation at:
https://filterpy.readthedocs.org

Supporting book at:
https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python

This is licensed under an MIT license. See the readme.MD file
for more information.
"""

#pylint: skip-file

from numpy.random import randn

class RadarSim(object):
    """
    Simulates the radar signal returns from an object flying
    at a constant altityude and velocity in 1D. Assumes random
    process noise - altitude and velocity varies a bit for each call.
    """

    def __init__(self, dt, pos=0., vel=100., alt=1000.):
        self.dt = dt
        self.pos = pos
        self.vel = vel
        self.alt = alt

    def get_range(self, process_err_pct=0.05):
        """
        Returns slant range to the object. Call once for each
        new measurement at dt time from last call.
        """

        vel = self.vel + 5 * randn()
        alt = self.alt + 10 * randn()

        self.pos += vel*self.dt

        err = (self.pos * process_err_pct) * randn()
        slant_range = (self.pos**2 + alt**2)**.5 + err

        return slant_range

reason: from numpy


text:
"""
BOE Semantic segmentation Dataset Loader
"""
import os
import os.path as osp
import json
import numpy as np
from config import cfg
from runx.logx import logx
from datasets.base_loader import BaseLoader
from datasets.utils import make_dataset_folder
from datasets import uniform
ID_TO_NAME = {
0:"construction--flat--road",
1:"construction--flat--sidewalk",
2:"object--street-light",
3:"construction--structure--bridge",
4:"construction--structure--building",
5:"human",
6:"object--support--pole",
7:"marking--continuous--dashed",
8:"marking--continuous--solid",
9:"marking--discrete--crosswalk-zebra",
10:"nature--sand",
11:"nature--sky",
12:"nature--snow",
13:"nature--terrain",
14:"nature--vegetation",
15:"nature--water",
16:"object--vehicle--bicycle",
17:"object--vehicle--boat",
18:"object--vehicle--bus",
19:"object--vehicle--car",
20:"object--vehicle--caravan",
21:"object--vehicle--motorcycle",
22:"object--vehicle--on-rails",
23:"object--vehicle--truck",
24:"construction--flat--pedestrian-area",
25:"construction--structure--tunnel",
26:"nature--wasteland",
}

NAME_TO_ID = {}
for k,v in ID_TO_NAME.items():
    NAME_TO_ID[v] = k

NAME_TO_MAPILLARY_NAME= {
"construction--flat--road":"construction--flat--road",
"construction--flat--sidewalk":"construction--flat--sidewalk",
"object--street-light":"object--street-light",
"construction--structure--bridge":"construction--structure--bridge",
"construction--structure--building":"construction--structure--building",
"human":"human--person--individual",
"object--support--pole":"object--support--pole",
"marking--continuous--dashed":"marking--continuous--dashed",
"marking--continuous--solid":"marking--continuous--solid",
"marking--discrete--crosswalk-zebra":"marking--discrete--crosswalk-zebra",
"nature--sand":"nature--sand",
"nature--sky":"nature--sky",
"nature--snow":"nature--snow",
"nature--terrain":"nature--terrain",
"nature--vegetation":"nature--vegetation",
"nature--water":"nature--water",
"object--vehicle--bicycle":"object--vehicle--bicycle",
"object--vehicle--boat":"object--vehicle--boat",
"object--vehicle--bus":"object--vehicle--bus",
"object--vehicle--car":"object--vehicle--car",
"object--vehicle--caravan":"object--vehicle--caravan",
"object--vehicle--motorcycle":"object--vehicle--motorcycle",
"object--vehicle--on-rails":"object--vehicle--on-rails",
"object--vehicle--truck":"object--vehicle--truck",
"construction--flat--pedestrian-area":"construction--flat--pedestrian-area",
"construction--structure--tunnel":"construction--structure--tunnel",
"nature--wasteland":"void--ground",
}
class Loader(BaseLoader):
    num_classes = len(ID_TO_NAME)
    ignore_label = 255
    trainid_to_name = ID_TO_NAME
    color_mapping = []

    def __init__(self, mode, quality='semantic', joint_transform_list=None,
                 img_transform=None, label_transform=None, eval_folder=None):

        super(Loader, self).__init__(quality=quality,
                                     mode=mode,
                                     joint_transform_list=joint_transform_list,
                                     img_transform=img_transform,
                                     label_transform=label_transform)
        root = cfg.DATASET.MAPILLARY_DIR
        config_fn = os.path.join(root, 'config.json')
        self.fill_colormap_and_names(config_fn)

        ######################################################################
        # Assemble image lists
        ######################################################################
        if mode == 'folder':
            self.all_imgs = make_dataset_folder(eval_folder)
        else:
            splits = {'train': 'training',
                      'val': 'validation',
                      'test': 'testing'}
            split_name = splits[mode]
            img_ext = 'jpg'
            mask_ext = 'png'
            img_root = os.path.join(root, split_name, 'images')
            mask_root = os.path.join(root, split_name, 'boe_labels')
            self.all_imgs = self.find_images(img_root, mask_root, img_ext,
                                             mask_ext)
            print(f"Find images1 {len(self.all_imgs)}")
            all_imgs_2 = self.find_images("/home/wj/ai/mldata/boesemantic/boe_labels",
                                          "/home/wj/ai/mldata/boesemantic/boe_labels",
                                          img_ext,
                                          mask_ext)
            print(f"Find images2 {len(all_imgs_2)}")
            all_imgs_2 = all_imgs_2*10
            if mode == "train":
                self.all_imgs = all_imgs_2 +self.all_imgs
            else:
                pass
            print(f"Find all images {len(self.all_imgs)}")

        logx.msg('all imgs {}'.format(len(self.all_imgs)))
        self.centroids = uniform.build_centroids(self.all_imgs,
                                                 self.num_classes,
                                                 self.train,
                                                 cv=cfg.DATASET.CV)
        self.build_epoch()

    def fill_colormap_and_names(self, config_fn):
        """
        Mapillary code for color map and class names

        Outputs
        -------
        self.trainid_to_name
        self.color_mapping
        """
        if not osp.exists(config_fn):
            print(f"Find file {config_fn} faild, use random color.")
            colormap = np.random.rand(255*3)*255
            colormap = colormap.astype(np.uint8)
            self.color_mapping = colormap
            return

        with open(config_fn) as config_file:
            config = json.load(config_file)
        config_labels = config['labels']

        mapillary_name2color = {}

        for i in range(0, len(config_labels)):
            label = config_labels[i]['name']
            color = config_labels[i]['color']
            mapillary_name2color[label] = color

        # calculate label color mapping
        colormap = []
        self.trainid_to_name = {}
        for i in range(0, len(ID_TO_NAME)):
            name = ID_TO_NAME[i]
            m_name = NAME_TO_MAPILLARY_NAME[name]
            color = mapillary_name2color[m_name]
            colormap = colormap + color
        colormap += [255,255,255]*(256-len(ID_TO_NAME))
        self.color_mapping = colormap