text:
#!/usr/bin/env python
"""
Data Augmentation wrappers and functions.
Functions to be used as "processing" arguments
are here.
"""

from __future__ import division

import neukrill_net.image_processing as image_processing
import numpy as np
import skimage.util
import itertools

def augmentation_wrapper(units='float64', **augment_settings):
    """
    Takes settings for augmentation as **kwargs
    and produces a "processing" function to make more
    training data.
    The returned function will return a list of images
    containing EVERY possible combinations of options
    as given in the augmentation settings.
    """
    # components of the processing pipeline
    components = []

    # Datatype unit conversion
    # We very probably want this to be float throughout
    if units == 'float64' or units == 'float':
        unitconvert = lambda images: [skimage.util.img_as_float(image)
                                        for image in images]

    elif units == 'uint8':
        unitconvert = lambda images: [skimage.util.img_as_ubyte(image)
                                        for image in images]

    elif units == None or units == 'auto':
        unitconvert = lambda images: images

    else:
        raise ValueError('Unrecognised output units: {}'.format(
                            augment_settings['units']))

    # Inversion
    if 'invert' in augment_settings and augment_settings['invert']:
        invert = lambda images: [image_processing.invert_image(image)
                                    for image in images]
    else:
        invert = lambda images: images

    # Shape-fixing without resizing
    if 'shape' in augment_settings:
        #
        shapefix = lambda images: [image_processing.shape_fix(
                                    image,
                                    augment_settings['shape'])
                                    for image in images]
    else:
        shapefix = lambda images: images

    # Resize
    if 'resize' in augment_settings:
        if 'resize_order' not in augment_settings:
            augment_settings['resize_order'] = 0.75
        # apply our resize
        resize = lambda images: [image_processing.resize_image(
                                    image,
                                    augment_settings['resize'],
                                    augment_settings['resize_order'])
                                    for image in images]
    else:
        resize = lambda images: images

    # Rescale
    if 'scale' in augment_settings and not augment_settings['scale']==None:
        scale = lambda images: [image_processing.scale_image(image, sf)
                                for sf in augment_settings['scale']
                                for image in images]
    else:
        scale = lambda images: images

    # Landscapise
    # Set to True if you want to ensure all the images are landscape
    if 'landscapise' in augment_settings and augment_settings['landscapise']:
        landscapise = lambda images: [image_processing.landscapise_image(image)
                                        for image in images]
    else:
        landscapise = lambda images: images

    # Rotate
    if 'rotate' in augment_settings and not augment_settings['rotate']==None:
        # apply rotate with options
        # will require a function here to call it a number of times
        rotate = lambda images: [rotatedImage for image in images
                                    for rotatedImage in
                                    rotations(image,
                                        augment_settings['rotate'],
                                        augment_settings['rotate_is_resizable']
                                    )
                                ]
    else:
        rotate = lambda images: images

    # Shear
    if 'shear' in augment_settings and not augment_settings['shear']==None:
        shear = lambda images: [image_processing.shear_image(image, this_shear)
                                for this_shear in augment_settings['shear']
                                for image in images]
    else:
        shear = lambda images: images

    # Flip (always horizontally)
    # All other relfections can be acheived by coupling with an appropriate reflection
    # Flip setting should either be True or False in settings
    if 'flip' in augment_settings and augment_settings['flip']:
        flip = lambda images: images + [image_processing.flip_image(image, True)
                                            for image in images]
    else:
        flip = lambda images: images

    # Crop (every side or not at all)
    if 'crop' in augment_settings and augment_settings['crop']:
        crop = lambda images: images + [croppedImage for image in images
                                            for croppedImage in
                                            allcrops(image)]
    else:
        crop = lambda images: images

    # Pad to translate within shape_fix window
    if 'translations' in augment_settings and not augment_settings['translations']==None:
        pad = lambda images: [paddedImage for image in images
                                            for paddedImage in
                                            allpads(image, augment_settings['translations'])]
    else:
        pad = lambda images: images

    # Add pixel noise
    if 'noise' in augment_settings and not augment_settings['noise']==None:
        noisify = lambda images: [image_processing.noisify_image(image, augment_settings['noise'])
                                    for image in images]
    else:
        noisify = lambda images: images

    # Stack all our functions together
    # Order matters here:
    # - Landscapise first because we don't want to undo rotations
    # - Rotate before anything else because then it has higher accuracy
    #   (might want to move in pixels which would otherwise be cropped
    #    don't want the loss of resolution caused by resize)
    # - Flip
    # - Crop
    # - Pad
    # - Align shape window
    # - Resize last because it is lossy
    # - Add noise independent for all output pixels
    processing=lambda image:noisify(
                                resize(
                                    shapefix(
                                        pad(
                                            crop(
                                                shear(
                                                    flip(
                                                        rotate(
                                                            landscapise(
                                                                unitconvert(
                                                                    invert(
                                                                        [ image ]
                                                                    )
                                                                )
                                                            )
                                                        )
                                                    )
                                                )
                                            )
                                        )
                                    )
                                )
                            )

    return processing

def rotations(image, num_rotations, resizable):
    """
    Returns a number of rotations depending on the settings.
    """
    return [image_processing.rotate_image(image, angle, resizable)
        for angle in np.linspace(0, 360, num_rotations, endpoint=False)]

def allcrops(image):
    """
    Returns a number of cropped copies of an image
    """
    return [image_processing.crop_image(image,side_id)
        for side_id in range(4)]

def allpads(image, pad_amounts):
    """
    Returns a list of padded images, with centre shift amounts
    specified by pad_amounts
    """
    # Can go in all four directions, so include +ve and -ve
    pad_amounts = np.array(pad_amounts)
    pad_amounts = np.union1d(-pad_amounts, pad_amounts)

    # Make permutations of all x and y shifts
    return [image_processing.padshift_image(image, shift)
        for shift in itertools.product(pad_amounts, repeat=2)]

class RandomAugment(object):
    """
    A class for performing random augmentations.
    When you call it with an image it gives
    you back a randomly manipulated version of it.
    """

    def __init__(self, random_seed=42, units='float64', rotate_is_resizable=False, **kwargs):
        """
        Inputs: random_seed
                keyword assignment of augmentation settings.
        """
        # Initialise the random number generator
        self.rng = np.random.RandomState(seed=random_seed)
        # Let 'shapefix' and 'shape' be aliases of the same property
        if 'shapefix' in kwargs:
            kwargs['shape'] = kwargs['shapefix']
        # Store settings as a dictionary
        self.settings = kwargs
        # Add kwargs with defaults into dictionary
        self.settings['units'] = units
        # Don't even need this anymore
        #self.settings['rotate_is_resizable'] = rotate_is_resizable

    def get_augs(self):
        """
        Computes the augmentation options
        """
        # Flip (always horizontally)
        # All other relfections can be acheived by coupling with an appropriate reflection
        # Flip setting should either be True or False in settings
        if 'flip' in self.settings and self.settings['flip']:
            # Flip if coin-toss says so
            flip = self.rng.binomial(1, 0.5)
        else:
            flip = False

        ### Affine Transformation
        # Rescale
        if ('scale' not in self.settings or
                self.settings['scale']==None or
                self.settings['scale']==0.0 ):
            scalefactor = 1.0
        else:
            if isinstance(self.settings['scale'],(list,tuple)):
                sf_index = self.rng.randint(0, len(self.settings['scale']))
                scalefactor = self.settings['scale'][sf_index]
            else:
                scalefactor = self.rng.normal(loc=1.0, scale=self.settings['scale'])
                # Ignore extreme events
                if (scalefactor-1.0) > 3*self.settings['scale']:
                    scalefactor = 1.0
                elif (scalefactor-1.0) < -3*self.settings['scale']:
                    scalefactor = 1.0

        # Asymmetric x and y scaling
        if ('scale_asym' not in self.settings
                or self.settings['scale_asym']==None
                or self.settings['scale_asym']==0.0 ):
            # x and y scales are the same
            scalefactor = (scalefactor,scalefactor)
        else:
            # Pick a new x and y scale near the overall scale
            x_sf = self.rng.normal(scale=self.settings['scale_asym'])
            y_sf = self.rng.normal(scale=self.settings['scale_asym'])
            # Ignore extreme events
            if x_sf > 3*self.settings['scale_asym'] or x_sf < -3*self.settings['scale_asym']:
                x_sf = 0.0
            if y_sf > 3*self.settings['scale_asym'] or y_sf < -3*self.settings['scale_asym']:
                y_sf = 0.0
            # Combine x and y scales as a tuple
            scalefactor = (scalefactor+x_sf, scalefactor+y_sf)

        # Define translation (shunt)
        if ('shunt' not in self.settings
                or self.settings['shunt']==None
                or self.settings['shunt']==0.0 ):
            shuntlist = (0,0)
        else:
            shuntlist = []
            for i in range(2):
                if isinstance(self.settings['shunt'],(list,tuple)):
                    # Pick from list
                    shunt_index = self.rng.randint(0, len(self.settings['shunt']))
                    shunt = self.settings['shunt'][shunt_index]
                    # Ensure symmetry
                    if self.rng.binomial(1, 0.5):
                        shunt = -shunt
                else:
                    # Chose from gaussian
                    shunt = self.rng.normal(loc=0.0, scale=self.settings['shunt'])
                    # Ignore extreme events
                    if shunt<-3*self.settings['shunt'] or shunt>3*self.settings['shunt']:
                        shunt = 0
                # Add to the list
                shuntlist += [shunt]
            # Convert list to tuple
            shuntlist = tuple(shuntlist)

        # Rotate
        if ('rotate' not in self.settings
                or self.settings['rotate']==None
                or self.settings['rotate']==0.0 ):
            rot_angle = 0.0
        else:
            if isinstance(self.settings['rotate'],(list,tuple)):
                # Pick from list of potential rotations
                rot_index = self.rng.randint(0, len(self.settings['rotate']))
                rot_angle = self.settings['rotate'][rot_index]

            elif self.settings['rotate']==-1:
                # Uniform distribution of rotations
                rot_angle = self.rng.uniform(low=0.0, high=360.0)

            else:
                # Select from implied list
                rot_index = self.rng.randint(0, self.settings['rotate'])
                rot_angle = 360.0 * rot_index / self.settings['rotate']

        # Shear
        if ('shear' not in self.settings
                or self.settings['shear']==None
                or self.settings['shear']==0.0 ):
            shear_angle = 0.0
        else:
            if isinstance(self.settings['shear'],(list,tuple)):
                shear_index = self.rng.randint(0, len(self.settings['shear']))
                shear_angle = self.settings['shear'][shear_index]
            else:
                # Chose from gaussian
                shear_angle = self.rng.normal(loc=0.0, scale=self.settings['shear'])

        # Order of affine transformation
        if 'transform_order' not in self.settings or self.settings['transform_order']==None:
            transform_order = 0.5
        else:
            transform_order = self.settings['transform_order']

        aug_dic = {'flip': flip,
                    'scalefactor': scalefactor, 'rot_angle': rot_angle,
                    'shear_angle': shear_angle, 'shuntlist': shuntlist,
                    'transform_order': transform_order}
        return aug_dic

    def __call__(self, image):
        """
        Basically a wrapper function for augment_and_process
        """

        # Landscapise
        # Set to True if you want to ensure all the images are landscape
        if 'landscapise' in self.settings and self.settings['landscapise']:
            image = image_processing.landscapise_image(image)

        # Get the augmentation properties
        aug_dic = self.get_augs()

        # Augment and preprocess
        return self.augment_and_process(image, aug_dic, self.settings)

    def augment_and_process(self, image, aug_dic, processing_settings):
        """
        Maps raw image to augmented image.
        """

        # Note down the original type
        original_dtype = image.dtype

        # Convert to float while we process it
        image = skimage.util.img_as_float(image)

        #####################################################
        # Pre-augmentation processing

        if 'invert' in processing_settings and processing_settings['invert']:
            image = image_processing.invert_image(image)

        # Rotate to nearest 90 degrees & update the angle to rotate by
        image, aug_dic['rot_angle'] = image_processing.rotate_nearest90_image(image, aug_dic['rot_angle'])

        # Ensure image is square now if we are going to resize
        if 'resize' in processing_settings:
            image = image_processing.pad_to_square(image)

        # Shape-fixing without resizing
        if 'shape' in processing_settings:
            if 'dynamic_shapefix' not in processing_settings or processing_settings['dynamic_shapefix']:
                # Do a dynamic shapefix where we pan to a random location of those viable
                pos_x = self.rng.uniform(low=0.0, high=1.0)
                pos_y = self.rng.uniform(low=0.0, high=1.0)
                image = image_processing.dynamic_shape_fix(image, processing_settings['shape'],
                            (pos_x,pos_y), do_crop=False, do_pad=True)

            else:
                image = image_processing.shape_fix(image, processing_settings['shape'],
                            do_crop=False, do_pad=True)

        #####################################################
        # Random augmentation

        if aug_dic['flip']:
            image = image_processing.flip_image(image, True)

        # Perform affine transformation
        image = image_processing.custom_transform_nice_units(image,
                    scale=aug_dic['scalefactor'], rotation=aug_dic['rot_angle'],
                    shear=aug_dic['shear_angle'], translation=aug_dic['shuntlist'],
                    order=aug_dic['transform_order'])

        #####################################################
        # Post-augmentation processing

        # Shape-fixing without resizing
        if 'shape' in processing_settings:
            if 'dynamic_shapefix' not in processing_settings or processing_settings['dynamic_shapefix']:
                # Do a dynamic shapefix where we pan to a random location of those viable
                image = image_processing.dynamic_shape_fix(image, processing_settings['shape'],
                            (pos_x,pos_y))

            else:
                image = image_processing.shape_fix(image, processing_settings['shape'])

        # Resize
        if not 'resize_order' in processing_settings or processing_settings['resize_order']==None:
            resize_order = 0.75
        if 'resize' in processing_settings:
            image = image_processing.resize_image(image, processing_settings['resize'], resize_order)

        #####################################################

        # Add pixel noise
        if ('noise' in self.settings
                and not self.settings['noise']==None
                and not self.settings['noise']==0.0 ):
            noise_seed = self.rng.randint(2**32)
            image = image_processing.noisify_image(image, self.settings['noise'], noise_seed)

        #####################################################

        if 'normalise' in processing_settings:
            if processing_settings['normalise']['global_or_pixel'] == 'global':
                mu = processing_settings['normalise']['mu']
                sigma = processing_settings['normalise']['sigma']
                image = (image - mu)/sigma
            elif processing_settings['normalise']['global_or_pixel'] == 'pixel':
                for i in range(image.shape[0]):
                    for j in range(image.shape[1]):
                        mu = processing_settings['normalise']['mu'][str((i,j))]
                        sigma = processing_settings['normalise']['sigma'][str((i,j))]
                        image[i,j] = (image[i,j] - mu)/sigma
            else:
                raise ValueError("Invalid option for global_or_pixel, should be "
                                     "one of global or pixel.")

        #####################################################

        # Datatype unit conversion
        if ('units' not in self.settings or
                self.settings['units'] == None or self.settings['units'] == 'auto'):
            # Preserve the datatype
            # Ensure output matches input
            image = image_processing.img_as_dtype(image, original_dtype)

        elif self.settings['units'] == 'float64' or self.settings['units'] == 'float':
            image = skimage.util.img_as_float(image)

        elif self.settings['units'] == 'uint8':
            image = skimage.util.img_as_ubyte(image)

        else:
            raise ValueError('Unrecognised output units: {}'.format(
                                self.settings['units']))

        return image

class ParallelRandomAugment(RandomAugment):
    """
    Random augmentation, but more than one at a time!
    """
    def __init__(self, preproc_list, **kwargs):

        # Assign preprocessing options to attributes
        self.preproc_list = preproc_list

        for i in range(len(self.preproc_list)):
            # Let 'shapefix' and 'shape' be aliases of the same property
            if 'shapefix' in self.preproc_list[i]:
                self.preproc_list[i]['shape'] = self.preproc_list[i]['shapefix']

        # Call superclass
        RandomAugment.__init__(self, **kwargs)

    def __call__(self, image):
        """
        Wraps two augment_and_process functions and returns
        two results as a tuple
        """

        # Landscapise
        # Set to True if you want to ensure all the images are landscape
        if 'landscapise' in self.settings and self.settings['landscapise']:
            image = image_processing.landscapise_image(image)

        # Get the augmentation properties
        aug_dic = self.get_augs()

        # Augment and preprocess
        images = []
        for preproc in self.preproc_list:
            images += [self.augment_and_process(image, aug_dic, preproc)]

        return images

max_stars_repo_name: Neuroglycerin/neukrill-net-tools

max_issues_repo_name: Neuroglycerin/neukrill-net-tools

max_forks_repo_name: Neuroglycerin/neukrill-net-tools


text:
# coding=utf-8
# Copyright 2021 The Uncertainty Baselines Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""BatchEnsemble Vision Transformer."""

import copy
import functools
import itertools
import multiprocessing
import os
import time

from absl import app
from absl import flags
from absl import logging
from clu import metric_writers
from clu import preprocess_spec
import flax
import flax.jax_utils
import flax.struct
import jax
import jax.config
import jax.nn
import jax.numpy as jnp
import ml_collections
import numpy as np
import robustness_metrics as rm

import tensorflow as tf
from tensorflow.io import gfile
import uncertainty_baselines as ub
import batchensemble_utils  # local file import
import checkpoint_utils  # local file import
import input_utils  # local file import
import preprocess_utils  # local file import
import train_utils  # local file import

# TODO(dusenberrymw): Open-source remaining imports.
ensemble = None
train = None
xprof = None
core = None
xm = None
xm_api = None
BIG_VISION_DIR = None

ml_collections.config_flags.DEFINE_config_file(
    'config', None, 'Training configuration.', lock_config=True)
flags.DEFINE_string('output_dir', default=None, help='Work unit directory.')
flags.DEFINE_integer(
    'num_cores', default=None, help='Unused. How many devices being used.')
flags.DEFINE_boolean(
    'use_gpu', default=None, help='Unused. Whether or not running on GPU.')
flags.DEFINE_string('tpu', None,
                    'Unused. Name of the TPU. Only used if use_gpu is False.')

# Adds jax flags to the program.
jax.config.parse_flags_with_absl()

def main(_):
  config = flags.FLAGS.config
  output_dir = flags.FLAGS.output_dir
  tf.io.gfile.makedirs(output_dir)

  seed = config.get('seed', 0)
  rng = jax.random.PRNGKey(seed)
  tf.random.set_seed(seed)

  save_checkpoint_path = None
  if config.get('checkpoint_steps'):
    gfile.makedirs(output_dir)
    save_checkpoint_path = os.path.join(output_dir, 'checkpoint.npz')

  # Create an asynchronous multi-metric writer.
  writer = metric_writers.create_default_writer(
      output_dir, just_logging=jax.process_index() > 0)
  checkpoint_writer = None

  # Loss to apply.
  loss_to_apply = getattr(core, config.get('loss_to_apply', 'softmax_xent'))
  compute_ece = config.get('compute_ece', False)
  is_sigmoid = config.get('loss_to_apply', 'softmax_xent') == 'sigmoid_xent'
  if compute_ece and is_sigmoid:
    error_msg = 'Inconsistent config: ECE can only be used with "softmax_xent".'
    raise ValueError(error_msg)

  ens_size = config.get('model.transformer.ens_size', 1)

  # The pool is used to perform misc operations such as logging in async way.
  pool = multiprocessing.pool.ThreadPool()

  # Ideally, this should make code deterministic, but for many reasons we are
  # not there yet. For instance, tf.data.map is not determisntic.
  rng_generator = tf.random.Generator.from_seed(config.get('seed', 0))
  tf.random.set_global_generator(
      rng_generator.split(jax.process_count())[jax.process_index()])

  logging.info('Number of devices: %s  (process index: %s)', jax.device_count(),
               jax.process_index())
  logging.info('Config:\n%s', str(config))

  def write_note(note):
    if jax.host_id() == 0:
      logging.info('NOTE: %s', note)

  write_note('Initializing...')

  # Verify settings to make sure no checkpoints are accidentally missed.
  if config.get('keep_checkpoint_steps'):
    assert config.get('checkpoint_steps'), 'Specify `checkpoint_steps`.'
    assert config.keep_checkpoint_steps % config.checkpoint_steps == 0, (
        f'`keep_checkpoint_steps` ({config.checkpoint_steps}) should be'
        f'divisible by `checkpoint_steps ({config.checkpoint_steps}).`')

  if (config.batch_size % jax.device_count() != 0 or
      config.batch_size_eval % jax.device_count() != 0):
    raise ValueError(f'Batch sizes ({config.batch_size} and '
                     f'{config.batch_size_eval}) must be divisible by '
                     f'the number of devices ({jax.device_count()})')

  batch_size_per_host = config.batch_size // jax.process_count()
  batch_size_per_core = config.batch_size // jax.device_count()
  batch_size_per_host_eval = config.batch_size_eval // jax.process_count()

  rng, train_ds_rng = jax.random.split(rng)
  train_ds_rng = jax.random.fold_in(train_ds_rng, jax.process_index())

  train_ds = input_utils.get_data(
      dataset=config.dataset,
      split=config.train_split,
      rng=train_ds_rng,
      host_batch_size=batch_size_per_host,
      preprocess_fn=preprocess_spec.parse(
          spec=config.pp_train, available_ops=preprocess_utils.all_ops()),
      shuffle_buffer_size=config.shuffle_buffer_size,
      prefetch_size=config.get('prefetch_to_host', 2),
      data_dir=config.get('dataset_dir'))
  train_iter = input_utils.start_input_pipeline(
      train_ds, config.get('prefetch_to_device', 1))

  ntrain_img = input_utils.get_num_examples(
      config.dataset,
      split=config.train_split,
      host_batch_size=batch_size_per_host,
      data_dir=config.get('dataset_dir'))
  steps_per_epoch = ntrain_img / config.batch_size

  if config.get('num_epochs'):
    total_steps = int(config.num_epochs * steps_per_epoch)
    assert not config.get('total_steps'), 'Set either num_epochs or total_steps'
  else:
    total_steps = config.total_steps

  logging.info(
      'Running for %d steps, that means %f epochs and %f steps per epoch',
      total_steps, total_steps * config.batch_size / ntrain_img,
      steps_per_epoch)

  def _get_val_split(dataset, split, pp_eval, data_dir=None):
    # We do ceil rounding such that we include the last incomplete batch.
    nval_img = input_utils.get_num_examples(
        dataset,
        split=split,
        host_batch_size=batch_size_per_host_eval,
        drop_remainder=False,
        data_dir=data_dir)
    val_steps = int(np.ceil(nval_img / config.batch_size_eval))
    logging.info('Running validation for %d steps for %s, %s', val_steps,
                 dataset, split)

    if isinstance(pp_eval, str):
      pp_eval = preprocess_spec.parse(
          spec=pp_eval, available_ops=preprocess_utils.all_ops())

    val_ds = input_utils.get_data(
        dataset=dataset,
        split=split,
        rng=None,
        host_batch_size=batch_size_per_host_eval,
        preprocess_fn=pp_eval,
        cache=config.get('val_cache', 'batched'),
        repeat_after_batching=True,
        shuffle=False,
        prefetch_size=config.get('prefetch_to_host', 2),
        drop_remainder=False,
        data_dir=data_dir)
    val_iter = input_utils.start_input_pipeline(
        val_ds, config.get('prefetch_to_device', 1))

    return (val_iter, val_steps)

  val_iter_splits = {
      'val':
          _get_val_split(
              config.dataset,
              split=config.val_split,
              pp_eval=config.pp_eval,
              data_dir=config.get('data_dir'))
  }

  # Note: we return the train loss and val loss for use in reproducibility unit
  # tests.
  train_loss = -jnp.inf
  val_loss = {val_name: -jnp.inf for val_name, _ in val_iter_splits.items()}
  # TODO(zmariet): Add fewshot evaluation.
  fewshot_results = {'dummy': {(0, 1): -jnp.inf}}

  opt_def = train.get_optimizer_from_config(config, f'{BIG_VISION_DIR}.optims')
  eval_config = copy.deepcopy(config)
  if config.get('eval_overrides'):
    with eval_config.unlocked():
      eval_config.update(config.eval_overrides)
  model = getattr(ub.models, config.model_name)
  model_train = model(
      num_classes=config.num_classes, train=True, **config.model)
  model_eval = model(
      num_classes=config.num_classes, train=False, **config.model)

  image_size = tuple(train_ds.element_spec['image'].shape[2:])
  logging.info('Model initialization: Starting.')
  opt_init, rngs = train.model_and_optim_init(
      model_train.init, opt_def, (batch_size_per_core,) + image_size,
      config.get('init_head_bias'), config.get('seed', 0),
      config.get('extra_rngs', []))
  logging.info('Model initialization: Done.')
  # TODO(jpuigcerver): Support logging parameter count with new sharding.

  weight_decay_fn = train.get_weight_decay_function_from_config(config)
  batch_loss_fn = ensemble.wrap_ensemble_module_with_auxiliary_loss_fn(
      module=model_train,
      loss_fn=loss_to_apply,
      auxiliary_loss_weight=config.get('auxiliary_loss_weight', 0.0),
      ens_size=ens_size)

  update_fn = functools.partial(
      batchensemble_utils.update_fn_be,
      weight_decay_fn=weight_decay_fn,
      plot_grad_norm_name_fn=None,
      plot_grads_nan_inf=config.get('plot_grads_nan_inf', True),
      max_grad_norm_global=config.get('clip_grad_norm', None),
      frozen_vars_patterns=config.get('frozen_var_patterns', None),
      fast_weight_lr_multiplier=config.get('fast_weight_lr_multiplier', None))
  pmap_update_fn = core.pmap_sorted(
      update_fn, axis_name='batch', donate_argnums=(0, 1),
      static_broadcasted_argnums=(5,))

  @functools.partial(jax.pmap, axis_name='batch')
  def evaluation_fn(params, images, labels, mask):
    # Ignore the entries with all zero labels for evaluation.
    mask *= labels.max(axis=1)
    tiled_logits, _ = model_eval.apply({'params': flax.core.freeze(params)},
                                       images)
    ens_logits = jnp.asarray(jnp.split(tiled_logits, ens_size))

    losses = ensemble.ensemble_softmax_xent(
        logits=ens_logits, labels=labels)
    loss = jax.lax.psum(losses * mask, axis_name='batch')

    ncorrect = ensemble.ensemble_softmax_correct_multilabel(
        ens_logits, labels, mask, psum_axis_name='batch')
    n = jax.lax.psum(mask, axis_name='batch')
    logits = jnp.log(jnp.mean(jax.nn.softmax(ens_logits), axis=0))
    metric_args = jax.lax.all_gather([logits, labels, mask],
                                     axis_name='batch')
    return ncorrect, loss, n, metric_args

  checkpoint_data = checkpoint_utils.maybe_load_checkpoint(
      train_loop_rngs=rng,
      save_checkpoint_path=save_checkpoint_path,
      init_optimizer=opt_init,
      init_params=opt_init.target,
      init_fixed_model_states=None,
      default_reinit_params=['head/bias', 'head/kernel'],
      config=config)

  opt = checkpoint_data.optimizer
  opt = opt.replace(target=flax.core.freeze(opt.target))
  opt = flax.jax_utils.replicate(opt)
  accumulated_train_time = checkpoint_data.accumulated_train_time

  train_loop_rngs = jax.tree_map(
      train.rng_jax_fold_host_if_needed_and_shard, rngs)

  first_step = int(opt.state.step)
  start_time = time.time()

  # Prepare the learning-rate and pre-fetch it to device to avoid delays.
  lr_fn = train_utils.create_learning_rate_schedule(total_steps,
                                                    **config.get('lr', {}))
  lr_iter = train_utils.prefetch_scalar(
      map(lr_fn, range(first_step, total_steps)),
      config.get('prefetch_to_device', 1))

  if first_step > 0:
    write_note('Advancing iterators after resuming from a checkpoint...')
    lr_iter = itertools.islice(lr_iter, first_step, None)
    train_iter = itertools.islice(train_iter, first_step, None)
    # NOTE: Validation eval is only run on certain steps, so determine how many
    # times it was run previously.
    num_val_runs = sum(
        map(
            lambda i: train_utils.itstime(i, config.log_eval_steps, total_steps
                                         ), range(1, first_step + 1)))
    for val_name, (val_iter, val_steps) in val_iter_splits.items():
      val_iter = itertools.islice(val_iter, num_val_runs * val_steps, None)
      val_iter_splits[val_name] = (val_iter, val_steps)

  log_training_first_n_steps = config.get('log_training_first_n_steps', -1)
  with metric_writers.ensure_flushes(writer):
    callback_fn = lambda x: x  # Do nothing.
    xprof_session = xprof.MultiStepXprofSession(
        profile_steps=20,    # For how many steps to profile after warmup.
        warmup_steps=170,    # For how many steps to wait before profiling.
        stop_callback_fn=callback_fn)
    for step, train_batch, lr_repl in zip(
        range(first_step + 1, total_steps + 1), train_iter, lr_iter):
      with xprof_session:
        with jax.profiler.StepTraceAnnotation(name='train', step_num=step):
          opt, train_loop_rngs, loss_value, _ = pmap_update_fn(
              opt,
              train_loop_rngs,
              lr_repl,
              train_batch['image'],
              train_batch['labels'],
              batch_loss_fn)

      # Checkpoint saving.
      if train_utils.itstime(
          step=step,
          every_n_steps=config.get('checkpoint_steps'),
          total_steps=total_steps, host=0, first=False):
        train_utils.checkpointing_timeout(checkpoint_writer,
                                          config.get('checkpoint_timeout', 1))
        train.sync_all_hosts()

        time_since_last_start = float(time.time() - start_time)
        accumulated_train_time = accumulated_train_time + time_since_last_start
        opt_cpu = jax.tree_util.tree_map(lambda x: np.array(x[0]), opt)

        copy_step = None
        if train_utils.itstime(step, config.get('keep_checkpoint_steps'),
                               total_steps):
          write_note('Keeping a checkpoint copy...')
          copy_step = step

        checkpoint_data = checkpoint_utils.CheckpointData(
            train_loop_rngs=train_loop_rngs,
            optimizer=opt_cpu,
            accumulated_train_time=accumulated_train_time)
        checkpoint_writer = pool.apply_async(
            checkpoint_utils.checkpoint_trained_model,
            (checkpoint_data, save_checkpoint_path, copy_step))

      # Report training progress
      if (jax.process_index() == 0 and config.log_training_every_n_steps > 0 and
          (step % config.log_training_every_n_steps == 0 or
           step == total_steps or step < log_training_first_n_steps)):
        train_loss = loss_value[0]
        time_elapsed = time.time() - start_time + accumulated_train_time
        img_sec_core = (
            config.batch_size * step / time_elapsed / jax.device_count())
        writer.write_scalars(step, {'learning_rate': lr_repl[0],
                                    'training_loss': np.mean(loss_value),
                                    'img/sec/core': img_sec_core,
                                    'epoch': step / steps_per_epoch})

      # Evaluate the model.
      if train_utils.itstime(step, config.log_eval_steps, total_steps):
        for val_name, (val_iter, val_steps) in val_iter_splits.items():
          # Sets up evaluation metrics.
          ece_num_bins = config.get('ece_num_bins', 15)
          auc_num_bins = config.get('auc_num_bins', 1000)
          ece = rm.metrics.ExpectedCalibrationError(num_bins=ece_num_bins)
          calib_auc = rm.metrics.CalibrationAUC(correct_pred_as_pos_label=False)
          oc_auc_0_5 = rm.metrics.OracleCollaborativeAUC(oracle_fraction=0.005,
                                                         num_bins=auc_num_bins)
          oc_auc_1 = rm.metrics.OracleCollaborativeAUC(oracle_fraction=0.01,
                                                       num_bins=auc_num_bins)
          oc_auc_2 = rm.metrics.OracleCollaborativeAUC(oracle_fraction=0.02,
                                                       num_bins=auc_num_bins)
          oc_auc_5 = rm.metrics.OracleCollaborativeAUC(oracle_fraction=0.05,
                                                       num_bins=auc_num_bins)

          # Runs evaluation loop.
          ncorrect, loss, nseen = 0, 0, 0
          for _, batch in zip(range(val_steps), val_iter):
            batch_ncorrect, batch_losses, batch_n, batch_metric_args = (
                evaluation_fn(opt.target, batch['image'],
                              batch['labels'], batch['mask']))
            # All results are a replicated array shaped as follows:
            # (local_devices, per_device_batch_size, elem_shape...)
            # with each local device's entry being identical as they got psum'd.
            # So let's just take the first one to the host as numpy.
            ncorrect += np.sum(np.array(batch_ncorrect[0]))
            loss += np.sum(np.array(batch_losses[0]))
            nseen += np.sum(np.array(batch_n[0]))
            # Here we parse batch_metric_args to compute uncertainty metrics.
            # (e.g., ECE or Calibration AUC).
            logits, labels, masks = batch_metric_args
            masks = np.array(masks[0], dtype=np.bool)
            logits = np.array(logits[0])
            probs = jax.nn.softmax(logits)
            # From one-hot to integer labels, as required by ECE.
            int_labels = np.argmax(np.array(labels[0]), axis=-1)
            int_preds = np.argmax(logits, axis=-1)
            confidence = np.max(probs, axis=-1)
            for p, c, l, d, m in zip(probs, confidence, int_labels,
                                     int_preds, masks):
              ece.add_batch(p[m, :], label=l[m])
              calib_auc.add_batch(d[m], label=l[m], confidence=c[m])
              # TODO(jereliu): Extend to support soft multi-class probabilities.
              oc_auc_0_5.add_batch(d[m], label=l[m], custom_binning_score=c[m])
              oc_auc_1.add_batch(d[m], label=l[m], custom_binning_score=c[m])
              oc_auc_2.add_batch(d[m], label=l[m], custom_binning_score=c[m])
              oc_auc_5.add_batch(d[m], label=l[m], custom_binning_score=c[m])

          val_loss[val_name] = loss / nseen  # Keep for reproducibility tests.
          val_measurements = {
              f'{val_name}_prec@1': ncorrect / nseen,
              f'{val_name}_loss': val_loss[val_name],
              f'{val_name}_ece': ece.result()['ece'],
              f'{val_name}_calib_auc': calib_auc.result()['calibration_auc'],
              f'{val_name}_oc_auc_0.5%': oc_auc_0_5.result()[
                  'collaborative_auc'],
              f'{val_name}_oc_auc_1%': oc_auc_1.result()['collaborative_auc'],
              f'{val_name}_oc_auc_2%': oc_auc_2.result()['collaborative_auc'],
              f'{val_name}_oc_auc_5%': oc_auc_5.result()['collaborative_auc'],
          }
          writer.write_scalars(step, val_measurements)

  pool.close()
  pool.join()

  # Return final training loss, validation loss, and fewshot results for
  # reproducibility test cases.
  return train_loss, val_loss, fewshot_results

if __name__ == '__main__':
  app.run(main)

ext: py

max_stars_repo_name: pyun-ram/uncertainty-baselines

max_issues_repo_name: pyun-ram/uncertainty-baselines

max_forks_repo_name: pyun-ram/uncertainty-baselines

reason: import numpy,import jax


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

df = pd.read_csv("worlddev.csv", index_col = 0)

# Dataframe Variables
pop = df[['population']]
cont = df[['continent']]
life_exp = df[['life_exp']]
gdp = df[['gdp_cap']]

# Scatter Plot Setup
plt.scatter(gdp, life_exp, s = pop / 1000000)
plt.xscale('log')

# Graph Labels
xlab = 'GDP per Capita (USD)'
ylab = 'Life Expectancy (Years)'
title = 'World Development in 2007'
plt.xlabel(xlab)
plt.ylabel(ylab)
plt.title(title)

# Graph Intervals
tick_val = [1000, 10000, 100000]
tick_lab = ['1k', '10k', '100k']
plt.xticks(tick_val, tick_lab)

# Print Scatter Plot
plt.show

# Subsetting Dataframe Categorically
lifeexp_series = df['life_exp']
gdp_series = df['gdp_cap']
low_dev = np.logical_and(gdp_series < 12536, lifeexp_series < 70)
high_dev = np.logical_and(gdp_series >= 12536, lifeexp_series >= 70)

reason: import numpy


text:
"""GFW Colormaps

Official GFW Colormaps
----------------------

Colormaps for dark style:

dark.reception
dark.fishing
dark.presence

Colormaps for light style:

light.reception
light.fishing
light.presence

Unofficial GFW Colormaps
------------------------

Colormaps made by Juan Carlos:

unofficial.jc_presence
unofficial.jc_reception
unofficial.jc_linearorange
unofficial.jc_linearblue
unofficial.jc_linearpink

BivariateColormaps
------------------

## Factory functions to create bivariate colormaps
bivariate.TransparencyBivariateColormap
bivariate.SummedBivariateColormap

## Colormaps suitable for use with TransparencyBivariateColormap
bivariate.orange_blue
bivariate.blue_orange

"""
from matplotlib import colors as _colors
import numpy as _np

def _hex2cmap(name, hex_colors, flip=False):
    """Convert sequence of hex colors to matplotlib cmap

    Parameters
    ----------
    name : str
    hex_colors : sequence of str
        Must be in '#RRGGBB' format

    Returns
    -------
    LinearSegmentedColormap
    """
    if flip:
        hex_colors = hex_colors[::-1]

    def ramp(i0, i1):
        float_colors = [(int(x[i0:i1], 16) / 256.0) for x in hex_colors]
        ramp = zip(
            _np.linspace(0, 1, len(hex_colors), endpoint=True),
            float_colors,
            float_colors,
        )
        return tuple(ramp)

    cdict = {"red": ramp(1, 3), "green": ramp(3, 5), "blue": ramp(5, 7)}

    cmap = _colors.LinearSegmentedColormap(name, cdict, 256)
    cmap.set_bad(alpha=0.0)
    return cmap

class _Dark:
    reception = _hex2cmap("reception", ("#ff4573", "#7b2e8d", "#093b76", "#0c276c"))
    fishing = _hex2cmap("fishing", ("#0c276c", "#3b9088", "#eeff00", "#ffffff"))
    presence = _hex2cmap("presence", ("#0c276c", "#114685", "#00ffc3", "#ffffff"))

dark = _Dark()

class _Light:
    reception = _hex2cmap(
        "reception", ("#ff4573", "#7b2e8d", "#093b76", "#0c276c"), flip=True
    )
    fishing = _hex2cmap(
        "fishing", ("#0c276c", "#3b9088", "#eeff00", "#ffffff"), flip=True
    )
    presence = _hex2cmap(
        "presence", ("#0c276c", "#114685", "#00ffc3", "#ffffff"), flip=True
    )

light = _Light()

class _Bivariate:
    blue_orange = _colors.LinearSegmentedColormap.from_list(
        "blue_orange",
        [
            "#FF7E00",
            "#F18352",
            "#E08979",
            "#CC8F99",
            "#B496B4",
            "#999CCC",
            "#79A1E0",
            "#52A6F1",
            "#00AAFF",
        ][::-1],
    )
    orange_blue = blue_orange.reversed()
    bright_blue_orange = _colors.LinearSegmentedColormap.from_list(
        "bright_blue_orange",
        [
            "#0062F0",
            "#405FF2",
            "#645BE1",
            "#855ECF",
            "#A463BB",
            "#BB63A3",
            "#D06286",
            "#E16063",
            "#F2613B",
            "#F56A00",
        ],
    )
    bright_blue_orange = bright_blue_orange.reversed()
    aqua = _colors.LinearSegmentedColormap.from_list(
        "aqua", [(0.95, 0.95, 0.95), (0.39, 0.53, 1.0)]
    )
    orange = _colors.LinearSegmentedColormap.from_list(
        "orange", [(0.95, 0.95, 0.95), (1.0, 0.25, 0.0)]
    )
    lime = _colors.LinearSegmentedColormap.from_list(
        "lime", [(0.95, 0.95, 0.95), (0.39, 1.0, 0.53)]
    )
    pink = _colors.LinearSegmentedColormap.from_list(
        "pink", [(0.95, 0.95, 0.95), (1.0, 0.48, 0.66)]
    )
    from .maps.bivariate import TransparencyBivariateColormap
    from .maps.bivariate import MinBivariateColormap
    from .maps.bivariate import MaxBivariateColormap

bivariate = _Bivariate()

class _Misc:
    blue_orange = bivariate.blue_orange
    orange_blue = bivariate.orange_blue
    jc_presence = _hex2cmap("jc_presence", ("#3359A8", "#16A3A4", "#00FFC3", "#ffffff"))
    jc_reception = _hex2cmap(
        "jc_reception", ("#5E0C20", "#2927A8", "#41FFA7", "#ffffff")
    )
    jc_linearorange = _hex2cmap(
        "jc_linearorange", ("#0C276C", "#824158", "#FF9500", "#ffffff")
    )
    jc_linearblue = _hex2cmap(
        "jc_linearblue", ("#0C276C", "#1D5780", "#00FFC3", "#ffffff")
    )
    jc_linearpink = _hex2cmap(
        "jc_linearpink", ("#0C276C", "#4E289B", "#F74559", "#ffffff")
    )

misc = unofficial = _Misc()

max_stars_repo_name: GlobalFishingWatch/pyseas

max_issues_repo_name: GlobalFishingWatch/pyseas

max_forks_repo_name: GlobalFishingWatch/pyseas


text:
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
import copy

##@package graphs
#Convert Output csv file from the main model to pandas DataFrame
#Draw graph of values through time and scatter graph of several pairs of variables

##@method read_file
#read a csv file from given location
#input
#n : file number to read, location : location of files in string
#output
#data_people, data_relation : DataFrame
def read_file(n,location):
    data_people = pd.read_csv("{}People_{}.csv".format(location,n),names=["GroupID", "PersonID","Selflessness","Selfishness","Money","Strength","Power"])
    data_relation = pd.read_csv("{}Relationship_{}.csv".format(location,n), names = ["Subjet", "Target", "relation", "probability"])
    return data_people, data_relation

##@method read_entire_data
#read all required csv files from given location
#input
#N : number of files to read respectively for people and relationship
#location : location of files in string
#output
#people,relationship : list of DataFrame
def read_entire_data(N,location):
    people = []
    relationship = []
    for i in range(0, N):
        d1, d2 = read_file(i,location)
        people.append(d1)
        relationship.append(d2)
    return people, relationship

##@method entiregrouping
#make all people in one group and count several aspects
#input
#data : DataFrame with datas for people
#output
#DataFrame with number of people, number of people who are alive, average of Money, Power, Selfishness, Selflessness for living people
def entiregrouping(data):
    average = dict()
    for i in range(len(data.index)):
        group = 0
        if group not in average.keys():
            average[group] = [0,0,0,0,0,0]
        average[group][0] += 1
        if data["Money"][i] > 0:
            average[group][1] += 1
            average[group][2] += int(data["Money"][i])
            average[group][3] += int(data["Power"][i])
            average[group][4] += int(data["Selfishness"][i])
            average[group][5] += int(data["Selflessness"][i])
    for i in average:
        average[i][2] = average[i][2]/average[i][1]
        average[i][3] = average[i][3]/average[i][1]
        average[i][4] = average[i][4]/average[i][1]
        average[i][5] = average[i][5]/average[i][1]
    return pd.DataFrame(average,index = ["Total", "Alive", "Money", "Power","Selfishness","Selflessness"]).transpose()

##@method time_based
#collect data based on time
#input
#data_set : list of DataFrame, name : label of data to collect
#output
#Dataframe with data collected based on time
def time_based(data_set, name):
    collection = {}
    for i in data_set[0].index:
        collection[i] = []
        for j in data_set:
            collection[i].append(j[name][i])
    return pd.DataFrame(collection)

##@method relation_average
#compute average of relationship among all people for each person
#input
#data_relation : DataFrame with datas of relationship
#output
#DataFrame with average of relationship for each person
def relation_average(data_relation):
    average = dict()
    for i in range(len(data_relation.index)):
        person = data_relation.iloc[i,0]
        if person not in average.keys():
            average[person] = 0
        average[person] += float(data_relation.iloc[i,2])
    for i in average.keys():
        average[i] = average[i]/(len(average)-1)
    return pd.DataFrame(average,index = ["average"]).transpose()

##@method power_relation
#make DataFrame of power and average of relation for each person
#input
#data_people : DataFrame of datas of people, data_relation : DataFrame of datas of relationship
#output
#DataFrame of power and average of relation for each person
def power_relation(data_people, data_relation):
    relation=relation_average(data_relation)
    power = data_people[["Power"]]
    data = pd.concat([power,relation],axis=1)
    return data

def main(N, location):
    N = int(N)
    people, relationship = read_entire_data(N, location+"/")
    pow_rel = [power_relation(people[i],relationship[i]) for i in range(N)]
    people_group = []
    for i in people:
        people_group.append(entiregrouping(i))
    time_based_alive = time_based(people_group,"Alive")
    time_based_money = time_based(people_group,"Money")
    time_based_power = time_based(people_group,"Power")

    plt.figure(1, figsize = (9,12))
    plt.title("Time-based Features")
    plt.subplot(311)
    plt.plot(time_based_alive)
    plt.ylabel("Alive")
    plt.subplot(312)
    plt.plot(time_based_money)
    plt.ylabel("Money")
    plt.subplot(313)
    plt.plot(time_based_power)
    plt.xlabel("Time")
    plt.ylabel("Power")

    fig1, axs1 = plt.subplots(2,(N//5+1)//2,figsize=(18,18),sharex=True,sharey=True)
    fig1.suptitle("Selfishness-Selflessness and Money")
    for i in range(0,N//5):
        j = 5*i+4
        axs1[i%2, i//2].scatter(people[j]["Selfishness"]-people[j]["Selflessness"],people[j]["Money"])
        axs1[i%2, i//2].set_title("epoch {}".format((i+1)*5))
        axs1[i%2, i//2].set(xlabel = "Selfishness-Selflessness",ylabel = "Money")
        axs1[i%2, i//2].label_outer()

    fig2, axs2 = plt.subplots(2,(N//5+1)//2,figsize=(18,18),sharex=True,sharey=True)
    fig2.suptitle("Power and Relation")
    for i in range(0,N//5):
        j = 5*i+4
        axs2[i%2, i//2].scatter(pow_rel[j]["Power"],pow_rel[j]["average"])
        axs2[i%2, i//2].set_title("epoch {}".format((i+1)*5))
        axs2[i%2, i//2].set(xlabel = "Power",ylabel = "Average of Relation")
        axs2[i%2, i//2].label_outer()

    fig3, axs3 = plt.subplots(2,(N//5+1)//2,figsize=(18,18),sharex=True,sharey=True)
    fig3.suptitle("Strength and Power")
    for i in range(0,N//5):
        j = 5*i+4
        axs3[i%2, i//2].scatter(people[j]["Strength"],people[j]["Power"])
        axs3[i%2, i//2].set_title("epoch {}".format((i+1)*5))
        axs3[i%2, i//2].set(xlabel = "Strength",ylabel = "Power")
        axs3[i%2, i//2].label_outer()

    fig4, axs4 = plt.subplots(2,(N//5+1)//2,figsize=(18,18),sharex=True,sharey=True)
    fig4.suptitle("Power and Money")
    for i in range(0,N//5):
        j = 5*i+4
        axs4[i%2, i//2].scatter(people[j]["Power"],people[j]["Money"])
        axs4[i%2, i//2].set_title("epoch {}".format((i+1)*5))
        axs4[i%2, i//2].set(xlabel = "Power",ylabel = "Money")
        axs4[i%2, i//2].label_outer()

    fig5, axs5 = plt.subplots(2,(N//5+1)//2,figsize=(18,18),sharex=True,sharey=True)
    fig5.suptitle("Selfishness-Selflessness and Relation")
    for i in range(0,N//5):
        j = 5*i+4
        axs5[i%2, i//2].scatter(people[j]["Selfishness"]-people[j]["Selflessness"],pow_rel[j]["average"])
        axs5[i%2, i//2].set_title("epoch {}".format((i+1)*5))
        axs5[i%2, i//2].set(xlabel = "Selfishness-Selflessness",ylabel = "Average of Relation")
        axs5[i%2, i//2].label_outer()

    fig6, axs6 = plt.subplots(2,(N//5+1)//2,figsize=(18,18),sharex=True,sharey=True)
    fig6.suptitle("Relation and Money")
    for i in range(0,N//5):
        j = 5*i+4
        axs6[i%2, i//2].scatter(pow_rel[j]["average"], people[i]["Money"])
        axs6[i%2, i//2].set_title("epoch {}".format((i+1)*5))
        axs6[i%2, i//2].set(xlabel = "Average of Relation",ylabel = "Money")
        axs6[i%2, i//2].label_outer()
    plt.show()

if __name__=="__main__":
    N = input("Number of meeting cycle : ")
    location = input("Location of data : ")
    main(N, location)


text:
"""
An extended version of the trainRegression.py example with pretraining and
some graphs at the end to visualize the output of the BNN.
"""

import os
import math
import warnings
import time

import numpy as np
import random as rn
import tensorflow as tf
import pylab as plt

from tensorBNN.activationFunctions import Tanh
from tensorBNN.layer import GaussianDenseLayer
from tensorBNN.networkFinal import network
from tensorBNN.likelihood import FixedGaussianLikelihood
from tensorBNN.metrics import SquaredError, PercentError
from tensorBNN.predictor import predictor

startTime = time.time()

# This supresses many deprecation warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)

# Set the GPU to use
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

os.environ["PYTHONHASHSEED"] = "0"
np.random.seed(42)
rn.seed(12345)
tf.random.set_seed(3)

def main():

    trainIn=np.linspace(-2,2,num=11)
    valIn=np.linspace(-2+2/30,2.0-2/30,num=30)
    trainOut = np.sin(trainIn*math.pi*2)*trainIn-np.cos(trainIn*math.pi)
    valOut = np.sin(valIn*math.pi*2)*valIn-np.cos(valIn*math.pi)

    data=[trainIn, trainOut, valIn, valOut]

    dtype=tf.float32

    inputDims=1
    outputDims=1
    width = 10 # perceptrons per layer
    hidden = 3 # number of hidden layers
    patience=20
    cycles=3
    epochs=100
    seed=1000

    normInfo=(0,1) # mean, sd

    #Peform pre-training to start the Markov Chain at a better spot
    model = tf.keras.Sequential()

    model.add(tf.keras.layers.Dense(width, kernel_initializer='glorot_uniform',
                                    input_shape=(inputDims, ),
                                    activation="tanh"))
    model.add(tf.keras.layers.ReLU())

    for n in range(hidden-1):
        model.add(tf.keras.layers.Dense(width,
                                        kernel_initializer='glorot_uniform',
                                        activation="tanh"))

    model.add(tf.keras.layers.Dense(outputDims,
                                    kernel_initializer='glorot_uniform'))

    callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss',
                                                patience=patience,
                                                restore_best_weights=True)

    #Train with decreasing learning rate
    for x in range(cycles):
        model.compile(optimizer=tf.keras.optimizers.Adam(0.01*(10**(-x)),
                                                         amsgrad=True),
                  loss='mean_squared_error',
                  metrics=['mean_absolute_error', 'mean_squared_error'])
        model.summary()
        model.fit(trainIn, trainOut.T, validation_data=(valIn, valOut.T),
                  epochs=epochs, batch_size=32, callbacks=[callback])

    #Save the backup
    model.save("backup")

    #Extract weights and biases
    weights=[]
    biases=[]
    activation=[]
    for layer in model.layers:
        weightBias=layer.get_weights()
        if(len(weightBias)==2):
            weights.append(weightBias[0].T)
            bias=weightBias[1]
            bias=np.reshape(bias, (len(bias),1))
            biases.append(bias)
        if(len(weightBias)==1):
            activation.append(weightBias[0])

    likelihood=FixedGaussianLikelihood(sd=0.1)
    metricList=[SquaredError(mean=normInfo[0], sd=normInfo[1]),
                PercentError(mean=normInfo[0], sd=normInfo[1])]

    neuralNet = network(
                dtype, # network datatype
                inputDims, # dimension of input vector
                data[0], # training input data
                data[1].T, # training output data
                data[2], # validation input data
                data[3].T) # validation output data)

    layer = GaussianDenseLayer( # Dense layer object
        inputDims, # Size of layer input vector
        width, # Size of layer output vector
        seed=seed, # Random seed
        dtype=dtype,
        weights=weights[0], biases=biases[0])
    neuralNet.add(layer) # Layer datatype
    neuralNet.add(Tanh()) # Tanh activation function
    seed += 1000 # Increment random seed
    for n in range(hidden - 1): # Add more hidden layers
        neuralNet.add(GaussianDenseLayer(width,
                                width,
                                seed=seed,
                                dtype=dtype,
                                weights=weights[n+1], biases=biases[n+1]))
        neuralNet.add(Tanh())
        seed += 1000

    neuralNet.add(GaussianDenseLayer(width,
                            outputDims,
                            seed=seed,
                            dtype=dtype,
                            weights=weights[-1], biases=biases[-1]))

    neuralNet.setupMCMC(
        stepSizeStart=1e-3,#0.0004 # starting stepsize
        stepSizeMin=1e-4, #0.0002 # minimum stepsize
        stepSizeMax=1e-2, # maximum stepsize
        stepSizeOptions=100, # number of stepsize options in stepsize adapter
        leapfrogStart=1000, # starting number of leapfrog steps
        leapfogMin=100, # minimum number of leapfrog steps
        leapFrogMax=10000, # maximum number of leapfrog steps
        leapfrogIncrement=10, # stepsize between leapfrog steps in leapfrog step adapter
        hyperStepSize=0.001, # hyper parameter stepsize
        hyperLeapfrog=100, # hyper parameter number of leapfrog steps
        burnin=1000, # number of burnin epochs
        averagingSteps=10) # number of averaging steps for param adapters)

    neuralNet.train(
        6001, # epochs to train for
        10, # increment between network saves
        likelihood,
        metricList=metricList,
        adjustHypers=True,
        folderName="TrigRegression", # Name of folder for saved networks
        networksPerFile=50) # Number of networks saved per file

    print("Total time elapsed (seconds):", time.time() - startTime)

    #Load predictor
    loadedNetwork = predictor("TrigRegression/", tf.float32)

    #Look at the predictions ins the space between the training data
    closeIn=np.linspace(-2,2,num=1000)
    closeOut = np.sin(closeIn*math.pi*2)*closeIn-np.cos(closeIn*math.pi)

    closePredictions = np.squeeze(np.array(loadedNetwork.predict(
                                                  np.array([closeIn]).T, n=1)))
    closePredictionsMean = np.mean(closePredictions, axis=0)
    closePredictionsStd = np.std(closePredictions, axis=0)
    plt.figure()

    plt.fill_between(closeIn, closePredictionsMean-2*closePredictionsStd,
                     closePredictionsMean-1*closePredictionsStd, color=(1,1,0),
                     label="2 sd")
    plt.fill_between(closeIn, closePredictionsMean-1*closePredictionsStd,
                     closePredictionsMean+1*closePredictionsStd, color=(0,1,0),
                     label="1 sd")
    plt.fill_between(closeIn, closePredictionsMean+1*closePredictionsStd,
                     closePredictionsMean+2*closePredictionsStd, color=(1,1,0))
    plt.plot(closeIn,closePredictionsMean, color="k", label="predicted mean")
    plt.plot(closeIn, closeOut, color="r", label="true")
    plt.scatter(trainIn, trainOut, color="b", label="training data")
    plt.legend()
    plt.show()

    #Look at the predictions away from the training data
    farIn=np.linspace(-4,4,num=2000)
    farOut = np.sin(farIn*math.pi*2)*farIn-np.cos(farIn*math.pi)

    farPredictions = np.squeeze(np.array(loadedNetwork.predict(
                                                    np.array([farIn]).T, n=1)))
    farPredictionsMean = np.mean(farPredictions, axis=0)
    farPredictionsStd = np.std(farPredictions, axis=0)

    plt.figure()
    plt.fill_between(farIn, farPredictionsMean-2*farPredictionsStd,
                     farPredictionsMean-1*farPredictionsStd, color=(1,1,0),
                     label="2 sd")
    plt.fill_between(farIn, farPredictionsMean-1*farPredictionsStd,
                     farPredictionsMean+1*farPredictionsStd, color=(0,1,0),
                     label="1 sd")
    plt.fill_between(farIn, farPredictionsMean+1*farPredictionsStd,
                     farPredictionsMean+2*farPredictionsStd, color=(1,1,0))
    plt.plot(farIn,farPredictionsMean, color="k", label="predicted mean")
    plt.plot(farIn, farOut, color="r", label="true")
    plt.scatter(trainIn, trainOut, color="b", label="training data")
    plt.legend()
    plt.show()

if(__name__ == "__main__"):
    main()


text:
#!/usr/bin/python

from brainart.utils import get_avg_rgb, get_color_lookup, generate_matching_df, get_packagedir
from brainart.template import get_template, sub_template, save_template
from glob import glob
import numpy as np
import webbrowser
import tempfile
import pandas
import json
import sys
import os

base = get_packagedir()

def generate(template,output_folder=None,color_lookup="white",image_base_path=None,bgcolor="black",top=20,sample=10):
    '''generate
    create a brainart image using a particular color lookup table
    :param template: the template image to generate brainart for! Only jpg has been tested.
    :param output_folder: output folder for the final file. If not specified, will use TMP
    :param color_lookup: one of "white" "black" "greywhite" or "greyblack" will default to using brainart database, and you do not need to worry about setting image_base_path. If you set this to a custom lookup table (should be a pandas dataframe) then you also need to set image_base_path to be the directory with images defined in your lookup table index. Default color lookup is white.
    :param image_base_path: Only specify if you provide a custom color lookup pandas. Default is None, as the base path is the github repo images.
    :param top: The number of top matches to sample from. Smaller values mean less variation in brains and color. Higher values means possibly more variation in color, and more brains.
    :param bgcolor: background color (string or hex) for the output html page. Default is black, and will be matched to color lookup if one is provided.
    :param sample: how many pixels to sample from. Default is every 5th pixel (5). Larger numbers mean generating images faster.
    '''
    if output_folder == None:
        output_folder = tempfile.mkdtemp()

    if isinstance(color_lookup,pandas.DataFrame):
        print("Custom color lookup found!")
        if image_base_path == None:
            print("Error: color lookup is set to custom file but image_base_path is set to 'None.' If you want to use images for a custom table you have created, you need to provide a base path to those images.")
            sys.exit()
    elif color_lookup not in ["white","black"]:
        print("Error: package color lookup tables include 'white','black','greywhite' and 'greyblack'")
        sys.exit()
    else:
        # Background color of html page should be black in all cases except for white lookup
        if color_lookup == "white":
            bgcolor = "white"
        color_lookup = pandas.read_pickle("%s/data/%s_lookup.pkl" %(base,color_lookup))
        image_base_path = "https://rawgithub.com/vsoch/brainart/master/png"

    print("Generating image... this can take a few minutes.")
    new_image = generate_matching_df(template,color_lookup,top=top,sample=sample)

    # Make paths relative to image_base_path
    new_image.png = ["%s/%s" %(image_base_path,x) for x in new_image.png]

    # Get image template
    html_template = get_template("index")
    image_name = os.path.splitext(os.path.basename(template))[0]

    # Save output to file
    new_image_dict = new_image.to_dict(orient="records")
    result = {"image":new_image_dict,"bgcolor":bgcolor}
    html_template = sub_template(html_template,"DATA",str(result))
    output_html = "%s/index.html" %output_folder
    save_template(output_html,html_template)

    print("Result files being saved to %s" %output_folder)
    webbrowser.open_new_tab(output_html)

max_stars_repo_name: vsoch/mosaic

max_issues_repo_name: vsoch/mosaic

max_forks_repo_name: vsoch/mosaic


text:
import tensorflow as tf
from deep_heatmaps_model_primary_valid import DeepHeatmapsModel
import os
import numpy as np

num_tests = 10
params = np.logspace(-8, -2, num_tests)
max_iter = 80000

output_dir = 'tests_lr_fusion'
data_dir = '../conventional_landmark_detection_dataset'

flags = tf.app.flags
flags.DEFINE_string('output_dir', output_dir, "directory for saving the log file")
flags.DEFINE_string('img_path', data_dir, "data directory")
FLAGS = flags.FLAGS

if not os.path.exists(FLAGS.output_dir):
    os.mkdir(FLAGS.output_dir)

for param in params:
    test_name = str(param)
    test_dir = os.path.join(FLAGS.output_dir,test_name)
    if not os.path.exists(test_dir):
        os.mkdir(test_dir)

    print '##### RUNNING TESTS ##### current directory:', test_dir

    save_model_path = os.path.join(test_dir, 'model')
    save_sample_path = os.path.join(test_dir, 'sample')
    save_log_path = os.path.join(test_dir, 'logs')

    # create directories if not exist
    if not os.path.exists(save_model_path):
        os.mkdir(save_model_path)
    if not os.path.exists(save_sample_path):
        os.mkdir(save_sample_path)
    if not os.path.exists(save_log_path):
        os.mkdir(save_log_path)

    tf.reset_default_graph()  # reset graph

    model = DeepHeatmapsModel(mode='TRAIN', train_iter=max_iter, learning_rate=param, momentum=0.95, step=80000,
                              gamma=0.1, batch_size=4, image_size=256, c_dim=3, num_landmarks=68,
                              augment_basic=True, basic_start=0, augment_texture=True, p_texture=0.5,
                              augment_geom=True, p_geom=0.5, artistic_start=0, artistic_step=10,
                              img_path=FLAGS.img_path, save_log_path=save_log_path, save_sample_path=save_sample_path,
                              save_model_path=save_model_path)

    model.train()

max_stars_repo_name: konatasick/face-of-art

max_issues_repo_name: TrueMatthewKirkham/face-of-art

max_forks_repo_name: TrueMatthewKirkham/face-of-art