text:
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
r"""A Beam job to add contextual memory banks to tf.Examples.

This tool groups images containing bounding boxes and embedded context features
by a key, either `image/location` or `image/seq_id`, and time horizon,
then uses these groups to build up a contextual memory bank from the embedded
context features from each image in the group and adds that context to the
output tf.Examples for each image in the group.

Steps to generate a dataset with context from one with bounding boxes and
embedded context features:
1. Use object/detection/export_inference_graph.py to get a `saved_model` for
  inference. The input node must accept a tf.Example proto.
2. Run this tool with `saved_model` from step 1 and a TFRecord of tf.Example
  protos containing images, bounding boxes, and embedded context features.
  The context features can be added to tf.Examples using
  generate_embedding_data.py.

Example Usage:
--------------
python add_context_to_examples.py \
  --input_tfrecord path/to/input_tfrecords* \
  --output_tfrecord path/to/output_tfrecords \
  --sequence_key image/location \
  --time_horizon month

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import datetime
import io
import itertools
import json
import os

from absl import app
from absl import flags
import apache_beam as beam
import numpy as np
import PIL.Image
import six
import tensorflow as tf

from apache_beam import runners

flags.DEFINE_string('input_tfrecord', None, 'TFRecord containing images in '
                    'tf.Example format for object detection, with bounding'
                    'boxes and contextual feature embeddings.')
flags.DEFINE_string('output_tfrecord', None,
                    'TFRecord containing images in tf.Example format, with '
                    'added contextual memory banks.')
flags.DEFINE_string('sequence_key', None, 'Key to use when grouping sequences: '
                    'so far supports `image/seq_id` and `image/location`.')
flags.DEFINE_string('time_horizon', None, 'What time horizon to use when '
                    'splitting the data, if any. Options are: `year`, `month`,'
                    ' `week`, `day `, `hour`, `minute`, `None`.')
flags.DEFINE_integer('subsample_context_features_rate', 0, 'Whether to '
                     'subsample the context_features, and if so how many to '
                     'sample. If the rate is set to X, it will sample context '
                     'from 1 out of every X images. Default is sampling from '
                     'every image, which is X=0.')
flags.DEFINE_boolean('reduce_image_size', True, 'downsamples images to'
                     'have longest side max_image_dimension, maintaining aspect'
                     ' ratio')
flags.DEFINE_integer('max_image_dimension', 1024, 'sets max image dimension')
flags.DEFINE_boolean('add_context_features', True, 'adds a memory bank of'
                     'embeddings to each clip')
flags.DEFINE_boolean('sorted_image_ids', True, 'whether the image source_ids '
                     'are sortable to deal with date_captured tie-breaks')
flags.DEFINE_string('image_ids_to_keep', 'All', 'path to .json list of image'
                    'ids to keep, used for ground truth eval creation')
flags.DEFINE_boolean('keep_context_features_image_id_list', False, 'Whether or '
                     'not to keep a list of the image_ids corresponding to the '
                     'memory bank')
flags.DEFINE_boolean('keep_only_positives', False, 'Whether or not to '
                     'keep only positive boxes based on score')
flags.DEFINE_boolean('keep_only_positives_gt', False, 'Whether or not to '
                     'keep only positive boxes based on gt class')
flags.DEFINE_float('context_features_score_threshold', 0.7, 'What score '
                   'threshold to use for boxes in context_features')
flags.DEFINE_integer('max_num_elements_in_context_features', 2000, 'Sets max '
                     'num elements per memory bank')
flags.DEFINE_integer('num_shards', 0, 'Number of output shards.')
flags.DEFINE_string('output_type', 'tf_sequence_example', 'Output type, one of '
                    '`tf_example`, `tf_sequence_example`')
flags.DEFINE_integer('max_clip_length', None, 'Max length for sequence '
                     'example outputs.')

FLAGS = flags.FLAGS

DEFAULT_FEATURE_LENGTH = 2057

class ReKeyDataFn(beam.DoFn):
  """Re-keys tfrecords by sequence_key.

  This Beam DoFn re-keys the tfrecords by a user-defined sequence_key
  """

  def __init__(self, sequence_key, time_horizon,
               reduce_image_size, max_image_dimension):
    """Initialization function.

    Args:
      sequence_key: A feature name to use as a key for grouping sequences.
        Must point to a key of type bytes_list
      time_horizon: What length of time to use to partition the data when
        building the memory banks. Options: `year`, `month`, `week`, `day `,
        `hour`, `minute`, None
      reduce_image_size: Whether to reduce the sizes of the stored images.
      max_image_dimension: maximum dimension of reduced images
    """
    self._sequence_key = sequence_key
    if time_horizon is None or time_horizon in {'year', 'month', 'week', 'day',
                                                'hour', 'minute'}:
      self._time_horizon = time_horizon
    else:
      raise ValueError('Time horizon not supported.')
    self._reduce_image_size = reduce_image_size
    self._max_image_dimension = max_image_dimension
    self._session = None
    self._num_examples_processed = beam.metrics.Metrics.counter(
        'data_rekey', 'num_tf_examples_processed')
    self._num_images_resized = beam.metrics.Metrics.counter(
        'data_rekey', 'num_images_resized')
    self._num_images_read = beam.metrics.Metrics.counter(
        'data_rekey', 'num_images_read')
    self._num_images_found = beam.metrics.Metrics.counter(
        'data_rekey', 'num_images_read')
    self._num_got_shape = beam.metrics.Metrics.counter(
        'data_rekey', 'num_images_got_shape')
    self._num_images_found_size = beam.metrics.Metrics.counter(
        'data_rekey', 'num_images_found_size')
    self._num_examples_cleared = beam.metrics.Metrics.counter(
        'data_rekey', 'num_examples_cleared')
    self._num_examples_updated = beam.metrics.Metrics.counter(
        'data_rekey', 'num_examples_updated')

  def process(self, tfrecord_entry):
    return self._rekey_examples(tfrecord_entry)

  def _largest_size_at_most(self, height, width, largest_side):
    """Computes new shape with the largest side equal to `largest_side`.

    Args:
      height: an int indicating the current height.
      width: an int indicating the current width.
      largest_side: A python integer indicating the size of
        the largest side after resize.
    Returns:
      new_height: an int indicating the new height.
      new_width: an int indicating the new width.
    """

    x_scale = float(largest_side) / float(width)
    y_scale = float(largest_side) / float(height)
    scale = min(x_scale, y_scale)

    new_width = int(width * scale)
    new_height = int(height * scale)

    return new_height, new_width

  def _resize_image(self, input_example):
    """Resizes the image within input_example and updates the height and width.

    Args:
      input_example: A tf.Example that we want to update to contain a resized
        image.
    Returns:
      input_example: Updated tf.Example.
    """

    original_image = copy.deepcopy(
        input_example.features.feature['image/encoded'].bytes_list.value[0])
    self._num_images_read.inc(1)

    height = copy.deepcopy(
        input_example.features.feature['image/height'].int64_list.value[0])

    width = copy.deepcopy(
        input_example.features.feature['image/width'].int64_list.value[0])

    self._num_got_shape.inc(1)

    new_height, new_width = self._largest_size_at_most(
        height, width, self._max_image_dimension)

    self._num_images_found_size.inc(1)

    encoded_jpg_io = io.BytesIO(original_image)
    image = PIL.Image.open(encoded_jpg_io)
    resized_image = image.resize((new_width, new_height))

    with io.BytesIO() as output:
      resized_image.save(output, format='JPEG')
      encoded_resized_image = output.getvalue()

    self._num_images_resized.inc(1)

    del input_example.features.feature['image/encoded'].bytes_list.value[:]
    del input_example.features.feature['image/height'].int64_list.value[:]
    del input_example.features.feature['image/width'].int64_list.value[:]

    self._num_examples_cleared.inc(1)

    input_example.features.feature['image/encoded'].bytes_list.value.extend(
        [encoded_resized_image])
    input_example.features.feature['image/height'].int64_list.value.extend(
        [new_height])
    input_example.features.feature['image/width'].int64_list.value.extend(
        [new_width])
    self._num_examples_updated.inc(1)

    return input_example

  def _rekey_examples(self, tfrecord_entry):
    serialized_example = copy.deepcopy(tfrecord_entry)

    input_example = tf.train.Example.FromString(serialized_example)

    self._num_images_found.inc(1)

    if self._reduce_image_size:
      input_example = self._resize_image(input_example)
      self._num_images_resized.inc(1)

    new_key = input_example.features.feature[
        self._sequence_key].bytes_list.value[0]

    if self._time_horizon:
      date_captured = datetime.datetime.strptime(
          six.ensure_str(input_example.features.feature[
              'image/date_captured'].bytes_list.value[0]), '%Y-%m-%d %H:%M:%S')
      year = date_captured.year
      month = date_captured.month
      day = date_captured.day
      week = np.floor(float(day) / float(7))
      hour = date_captured.hour
      minute = date_captured.minute

      if self._time_horizon == 'year':
        new_key = new_key + six.ensure_binary('/' + str(year))
      elif self._time_horizon == 'month':
        new_key = new_key + six.ensure_binary(
            '/' + str(year) + '/' + str(month))
      elif self._time_horizon == 'week':
        new_key = new_key + six.ensure_binary(
            '/' + str(year) + '/' + str(month) + '/' + str(week))
      elif self._time_horizon == 'day':
        new_key = new_key + six.ensure_binary(
            '/' + str(year) + '/' + str(month) + '/' + str(day))
      elif self._time_horizon == 'hour':
        new_key = new_key + six.ensure_binary(
            '/' + str(year) + '/' + str(month) + '/' + str(day) + '/' + (
                str(hour)))
      elif self._time_horizon == 'minute':
        new_key = new_key + six.ensure_binary(
            '/' + str(year) + '/' + str(month) + '/' + str(day) + '/' + (
                str(hour) + '/' + str(minute)))

    self._num_examples_processed.inc(1)

    return [(new_key, input_example)]

class SortGroupedDataFn(beam.DoFn):
  """Sorts data within a keyed group.

  This Beam DoFn sorts the grouped list of image examples by frame_num
  """

  def __init__(self, sequence_key, sorted_image_ids,
               max_num_elements_in_context_features):
    """Initialization function.

    Args:
      sequence_key: A feature name to use as a key for grouping sequences.
        Must point to a key of type bytes_list
      sorted_image_ids: Whether the image ids are sortable to use as sorting
        tie-breakers
      max_num_elements_in_context_features: The maximum number of elements
        allowed in the memory bank
    """
    self._session = None
    self._num_examples_processed = beam.metrics.Metrics.counter(
        'sort_group', 'num_groups_sorted')
    self._too_many_elements = beam.metrics.Metrics.counter(
        'sort_group', 'too_many_elements')
    self._split_elements = beam.metrics.Metrics.counter(
        'sort_group', 'split_elements')
    self._sequence_key = six.ensure_binary(sequence_key)
    self._sorted_image_ids = sorted_image_ids
    self._max_num_elements_in_context_features = (
        max_num_elements_in_context_features)

  def process(self, grouped_entry):
    return self._sort_image_examples(grouped_entry)

  def _sort_image_examples(self, grouped_entry):
    key, example_collection = grouped_entry
    example_list = list(example_collection)

    def get_frame_num(example):
      return example.features.feature['image/seq_frame_num'].int64_list.value[0]

    def get_date_captured(example):
      return datetime.datetime.strptime(
          six.ensure_str(
              example.features.feature[
                  'image/date_captured'].bytes_list.value[0]),
          '%Y-%m-%d %H:%M:%S')

    def get_image_id(example):
      return example.features.feature['image/source_id'].bytes_list.value[0]

    if self._sequence_key == six.ensure_binary('image/seq_id'):
      sorting_fn = get_frame_num
    elif self._sequence_key == six.ensure_binary('image/location'):
      if self._sorted_image_ids:
        sorting_fn = get_image_id
      else:
        sorting_fn = get_date_captured

    sorted_example_list = sorted(example_list, key=sorting_fn)

    self._num_examples_processed.inc(1)

    if len(sorted_example_list) > self._max_num_elements_in_context_features:
      leftovers = sorted_example_list
      output_list = []
      count = 0
      self._too_many_elements.inc(1)
      while len(leftovers) > self._max_num_elements_in_context_features:
        self._split_elements.inc(1)
        new_key = key + six.ensure_binary('_' + str(count))
        new_list = leftovers[:self._max_num_elements_in_context_features]
        output_list.append((new_key, new_list))
        leftovers = leftovers[:self._max_num_elements_in_context_features]
        count += 1
    else:
      output_list = [(key, sorted_example_list)]

    return output_list

def get_sliding_window(example_list, max_clip_length, stride_length):
  """Yields a sliding window over data from example_list.

  Sliding window has width max_clip_len (n) and stride stride_len (m).
     s -> (s0,s1,...s[n-1]), (s[m],s[m+1],...,s[m+n]), ...

  Args:
    example_list: A list of examples.
    max_clip_length: The maximum length of each clip.
    stride_length: The stride between each clip.

  Yields:
    A list of lists of examples, each with length <= max_clip_length
  """

  # check if the list is too short to slide over
  if len(example_list) < max_clip_length:
    yield example_list
  else:
    starting_values = [i*stride_length for i in
                       range(len(example_list)) if
                       len(example_list) > i*stride_length]
    for start in starting_values:
      result = tuple(itertools.islice(example_list, start,
                                      min(start + max_clip_length,
                                          len(example_list))))
      yield result

class GenerateContextFn(beam.DoFn):
  """Generates context data for camera trap images.

  This Beam DoFn builds up contextual memory banks from groups of images and
  stores them in the output tf.Example or tf.Sequence_example for each image.
  """

  def __init__(self, sequence_key, add_context_features, image_ids_to_keep,
               keep_context_features_image_id_list=False,
               subsample_context_features_rate=0,
               keep_only_positives=False,
               context_features_score_threshold=0.7,
               keep_only_positives_gt=False,
               max_num_elements_in_context_features=5000,
               pad_context_features=False,
               output_type='tf_example', max_clip_length=None):
    """Initialization function.

    Args:
      sequence_key: A feature name to use as a key for grouping sequences.
      add_context_features: Whether to keep and store the contextual memory
        bank.
      image_ids_to_keep: A list of image ids to save, to use to build data
        subsets for evaluation.
      keep_context_features_image_id_list: Whether to save an ordered list of
        the ids of the images in the contextual memory bank.
      subsample_context_features_rate: What rate to subsample images for the
        contextual memory bank.
      keep_only_positives: Whether to only keep high scoring
        (>context_features_score_threshold) features in the contextual memory
        bank.
      context_features_score_threshold: What threshold to use for keeping
        features.
      keep_only_positives_gt: Whether to only keep features from images that
        contain objects based on the ground truth (for training).
      max_num_elements_in_context_features: the maximum number of elements in
        the memory bank
      pad_context_features: Whether to pad the context features to a fixed size.
      output_type: What type of output, tf_example of tf_sequence_example
      max_clip_length: The maximum length of a sequence example, before
        splitting into multiple
    """
    self._session = None
    self._num_examples_processed = beam.metrics.Metrics.counter(
        'sequence_data_generation', 'num_seq_examples_processed')
    self._num_keys_processed = beam.metrics.Metrics.counter(
        'sequence_data_generation', 'num_keys_processed')
    self._sequence_key = sequence_key
    self._add_context_features = add_context_features
    self._pad_context_features = pad_context_features
    self._output_type = output_type
    self._max_clip_length = max_clip_length
    if six.ensure_str(image_ids_to_keep) == 'All':
      self._image_ids_to_keep = None
    else:
      with tf.io.gfile.GFile(image_ids_to_keep) as f:
        self._image_ids_to_keep = json.load(f)
    self._keep_context_features_image_id_list = (
        keep_context_features_image_id_list)
    self._subsample_context_features_rate = subsample_context_features_rate
    self._keep_only_positives = keep_only_positives
    self._keep_only_positives_gt = keep_only_positives_gt
    self._context_features_score_threshold = context_features_score_threshold
    self._max_num_elements_in_context_features = (
        max_num_elements_in_context_features)

    self._images_kept = beam.metrics.Metrics.counter(
        'sequence_data_generation', 'images_kept')
    self._images_loaded = beam.metrics.Metrics.counter(
        'sequence_data_generation', 'images_loaded')

  def process(self, grouped_entry):
    return self._add_context_to_example(copy.deepcopy(grouped_entry))

  def _build_context_features(self, example_list):
    context_features = []
    context_features_image_id_list = []
    count = 0
    example_embedding = []

    for idx, example in enumerate(example_list):
      if self._subsample_context_features_rate > 0:
        if (idx % self._subsample_context_features_rate) != 0:
          example.features.feature[
              'context_features_idx'].int64_list.value.append(
                  self._max_num_elements_in_context_features + 1)
          continue
      if self._keep_only_positives:
        if example.features.feature[
            'image/embedding_score'
            ].float_list.value[0] < self._context_features_score_threshold:
          example.features.feature[
              'context_features_idx'].int64_list.value.append(
                  self._max_num_elements_in_context_features + 1)
          continue
      if self._keep_only_positives_gt:
        if len(example.features.feature[
            'image/object/bbox/xmin'
            ].float_list.value) < 1:
          example.features.feature[
              'context_features_idx'].int64_list.value.append(
                  self._max_num_elements_in_context_features + 1)
          continue

      example_embedding = list(example.features.feature[
          'image/embedding'].float_list.value)
      context_features.extend(example_embedding)
      example.features.feature[
          'context_features_idx'].int64_list.value.append(count)
      count += 1
      example_image_id = example.features.feature[
          'image/source_id'].bytes_list.value[0]
      context_features_image_id_list.append(example_image_id)

    if not example_embedding:
      example_embedding.append(np.zeros(DEFAULT_FEATURE_LENGTH))

    feature_length = DEFAULT_FEATURE_LENGTH

    # If the example_list is not empty and image/embedding_length is in the
    # featture dict, feature_length will be assigned to that. Otherwise, it will
    # be kept as default.
    if example_list and (
        'image/embedding_length' in example_list[0].features.feature):
      feature_length = example_list[0].features.feature[
          'image/embedding_length'].int64_list.value[0]

    if self._pad_context_features:
      while len(context_features_image_id_list) < (
          self._max_num_elements_in_context_features):
        context_features_image_id_list.append('')

    return context_features, feature_length, context_features_image_id_list

  def _add_context_to_example(self, grouped_entry):
    key, example_collection = grouped_entry
    list_of_examples = []

    example_list = list(example_collection)

    if self._add_context_features:
      context_features, feature_length, context_features_image_id_list = (
          self._build_context_features(example_list))

    if self._image_ids_to_keep is not None:
      new_example_list = []
      for example in example_list:
        im_id = example.features.feature['image/source_id'].bytes_list.value[0]
        self._images_loaded.inc(1)
        if six.ensure_str(im_id) in self._image_ids_to_keep:
          self._images_kept.inc(1)
          new_example_list.append(example)
      if new_example_list:
        example_list = new_example_list
      else:
        return []

    if self._output_type == 'tf_sequence_example':
      if self._max_clip_length is not None:
        # For now, no overlap
        clips = get_sliding_window(
            example_list, self._max_clip_length, self._max_clip_length)
      else:
        clips = [example_list]

      for clip_num, clip_list in enumerate(clips):
        # initialize sequence example
        seq_example = tf.train.SequenceExample()
        video_id = six.ensure_str(key)+'_'+ str(clip_num)
        seq_example.context.feature['clip/media_id'].bytes_list.value.append(
            video_id.encode('utf8'))
        seq_example.context.feature['clip/frames'].int64_list.value.append(
            len(clip_list))

        seq_example.context.feature[
            'clip/start/timestamp'].int64_list.value.append(0)
        seq_example.context.feature[
            'clip/end/timestamp'].int64_list.value.append(len(clip_list))
        seq_example.context.feature['image/format'].bytes_list.value.append(
            six.ensure_binary('JPG'))
        seq_example.context.feature['image/channels'].int64_list.value.append(3)
        context_example = clip_list[0]
        seq_example.context.feature['image/height'].int64_list.value.append(
            context_example.features.feature[
                'image/height'].int64_list.value[0])
        seq_example.context.feature['image/width'].int64_list.value.append(
            context_example.features.feature['image/width'].int64_list.value[0])

        seq_example.context.feature[
            'image/context_feature_length'].int64_list.value.append(
                feature_length)
        seq_example.context.feature[
            'image/context_features'].float_list.value.extend(
                context_features)
        if self._keep_context_features_image_id_list:
          seq_example.context.feature[
              'image/context_features_image_id_list'].bytes_list.value.extend(
                  context_features_image_id_list)

        encoded_image_list = seq_example.feature_lists.feature_list[
            'image/encoded']
        timestamps_list = seq_example.feature_lists.feature_list[
            'image/timestamp']
        context_features_idx_list = seq_example.feature_lists.feature_list[
            'image/context_features_idx']
        date_captured_list = seq_example.feature_lists.feature_list[
            'image/date_captured']
        unix_time_list = seq_example.feature_lists.feature_list[
            'image/unix_time']
        location_list = seq_example.feature_lists.feature_list['image/location']
        image_ids_list = seq_example.feature_lists.feature_list[
            'image/source_id']
        gt_xmin_list = seq_example.feature_lists.feature_list[
            'region/bbox/xmin']
        gt_xmax_list = seq_example.feature_lists.feature_list[
            'region/bbox/xmax']
        gt_ymin_list = seq_example.feature_lists.feature_list[
            'region/bbox/ymin']
        gt_ymax_list = seq_example.feature_lists.feature_list[
            'region/bbox/ymax']
        gt_type_list = seq_example.feature_lists.feature_list[
            'region/label/index']
        gt_type_string_list = seq_example.feature_lists.feature_list[
            'region/label/string']
        gt_is_annotated_list = seq_example.feature_lists.feature_list[
            'region/is_annotated']

        for idx, example in enumerate(clip_list):

          encoded_image = encoded_image_list.feature.add()
          encoded_image.bytes_list.value.extend(
              example.features.feature['image/encoded'].bytes_list.value)

          image_id = image_ids_list.feature.add()
          image_id.bytes_list.value.append(
              example.features.feature['image/source_id'].bytes_list.value[0])

          timestamp = timestamps_list.feature.add()
          # Timestamp is currently order in the list.
          timestamp.int64_list.value.extend([idx])

          context_features_idx = context_features_idx_list.feature.add()
          context_features_idx.int64_list.value.extend(
              example.features.feature['context_features_idx'].int64_list.value)

          date_captured = date_captured_list.feature.add()
          date_captured.bytes_list.value.extend(
              example.features.feature['image/date_captured'].bytes_list.value)
          unix_time = unix_time_list.feature.add()
          unix_time.float_list.value.extend(
              example.features.feature['image/unix_time'].float_list.value)
          location = location_list.feature.add()
          location.bytes_list.value.extend(
              example.features.feature['image/location'].bytes_list.value)

          gt_xmin = gt_xmin_list.feature.add()
          gt_xmax = gt_xmax_list.feature.add()
          gt_ymin = gt_ymin_list.feature.add()
          gt_ymax = gt_ymax_list.feature.add()
          gt_type = gt_type_list.feature.add()
          gt_type_str = gt_type_string_list.feature.add()

          gt_is_annotated = gt_is_annotated_list.feature.add()
          gt_is_annotated.int64_list.value.append(1)

          gt_xmin.float_list.value.extend(
              example.features.feature[
                  'image/object/bbox/xmin'].float_list.value)
          gt_xmax.float_list.value.extend(
              example.features.feature[
                  'image/object/bbox/xmax'].float_list.value)
          gt_ymin.float_list.value.extend(
              example.features.feature[
                  'image/object/bbox/ymin'].float_list.value)
          gt_ymax.float_list.value.extend(
              example.features.feature[
                  'image/object/bbox/ymax'].float_list.value)

          gt_type.int64_list.value.extend(
              example.features.feature[
                  'image/object/class/label'].int64_list.value)
          gt_type_str.bytes_list.value.extend(
              example.features.feature[
                  'image/object/class/text'].bytes_list.value)

        self._num_examples_processed.inc(1)
        list_of_examples.append(seq_example)

    elif self._output_type == 'tf_example':

      for example in example_list:
        im_id = example.features.feature['image/source_id'].bytes_list.value[0]

        if self._add_context_features:
          example.features.feature[
              'image/context_features'].float_list.value.extend(
                  context_features)
          example.features.feature[
              'image/context_feature_length'].int64_list.value.append(
                  feature_length)

        if self._keep_context_features_image_id_list:
          example.features.feature[
              'image/context_features_image_id_list'].bytes_list.value.extend(
                  context_features_image_id_list)

        self._num_examples_processed.inc(1)
        list_of_examples.append(example)

    return list_of_examples

def construct_pipeline(input_tfrecord,
                       output_tfrecord,
                       sequence_key,
                       time_horizon=None,
                       subsample_context_features_rate=0,
                       reduce_image_size=True,
                       max_image_dimension=1024,
                       add_context_features=True,
                       sorted_image_ids=True,
                       image_ids_to_keep='All',
                       keep_context_features_image_id_list=False,
                       keep_only_positives=False,
                       context_features_score_threshold=0.7,
                       keep_only_positives_gt=False,
                       max_num_elements_in_context_features=5000,
                       num_shards=0,
                       output_type='tf_example',
                       max_clip_length=None):
  """Returns a beam pipeline to run object detection inference.

  Args:
    input_tfrecord: An TFRecord of tf.train.Example protos containing images.
    output_tfrecord: An TFRecord of tf.train.Example protos that contain images
      in the input TFRecord and the detections from the model.
    sequence_key: A feature name to use as a key for grouping sequences.
    time_horizon: What length of time to use to partition the data when building
      the memory banks. Options: `year`, `month`, `week`, `day `, `hour`,
      `minute`, None.
    subsample_context_features_rate: What rate to subsample images for the
        contextual memory bank.
    reduce_image_size: Whether to reduce the size of the stored images.
    max_image_dimension: The maximum image dimension to use for resizing.
    add_context_features: Whether to keep and store the contextual memory bank.
    sorted_image_ids: Whether the image ids are sortable, and can be used as
      datetime tie-breakers when building memory banks.
    image_ids_to_keep: A list of image ids to save, to use to build data subsets
      for evaluation.
    keep_context_features_image_id_list: Whether to save an ordered list of the
      ids of the images in the contextual memory bank.
    keep_only_positives: Whether to only keep high scoring
      (>context_features_score_threshold) features in the contextual memory
      bank.
    context_features_score_threshold: What threshold to use for keeping
      features.
    keep_only_positives_gt: Whether to only keep features from images that
      contain objects based on the ground truth (for training).
    max_num_elements_in_context_features: the maximum number of elements in the
      memory bank
    num_shards: The number of output shards.
    output_type: What type of output, tf_example of tf_sequence_example
    max_clip_length: The maximum length of a sequence example, before
      splitting into multiple
  """
  def pipeline(root):
    if output_type == 'tf_example':
      coder = beam.coders.ProtoCoder(tf.train.Example)
    elif output_type == 'tf_sequence_example':
      coder = beam.coders.ProtoCoder(tf.train.SequenceExample)
    else:
      raise ValueError('Unsupported output type.')
    input_collection = (
        root | 'ReadInputTFRecord' >> beam.io.tfrecordio.ReadFromTFRecord(
            input_tfrecord,
            coder=beam.coders.BytesCoder()))
    rekey_collection = input_collection | 'RekeyExamples' >> beam.ParDo(
        ReKeyDataFn(sequence_key, time_horizon,
                    reduce_image_size, max_image_dimension))
    grouped_collection = (
        rekey_collection | 'GroupBySequenceKey' >> beam.GroupByKey())
    grouped_collection = (
        grouped_collection | 'ReshuffleGroups' >> beam.Reshuffle())
    ordered_collection = (
        grouped_collection | 'OrderByFrameNumber' >> beam.ParDo(
            SortGroupedDataFn(sequence_key, sorted_image_ids,
                              max_num_elements_in_context_features)))
    ordered_collection = (
        ordered_collection | 'ReshuffleSortedGroups' >> beam.Reshuffle())
    output_collection = (
        ordered_collection | 'AddContextToExamples' >> beam.ParDo(
            GenerateContextFn(
                sequence_key, add_context_features, image_ids_to_keep,
                keep_context_features_image_id_list=(
                    keep_context_features_image_id_list),
                subsample_context_features_rate=subsample_context_features_rate,
                keep_only_positives=keep_only_positives,
                keep_only_positives_gt=keep_only_positives_gt,
                context_features_score_threshold=(
                    context_features_score_threshold),
                max_num_elements_in_context_features=(
                    max_num_elements_in_context_features),
                output_type=output_type,
                max_clip_length=max_clip_length)))

    output_collection = (
        output_collection | 'ReshuffleExamples' >> beam.Reshuffle())
    _ = output_collection | 'WritetoDisk' >> beam.io.tfrecordio.WriteToTFRecord(
        output_tfrecord,
        num_shards=num_shards,
        coder=coder)
  return pipeline

def main(_):
  """Runs the Beam pipeline that builds context features.

  Args:
    _: unused
  """
  # must create before flags are used
  runner = runners.DirectRunner()

  dirname = os.path.dirname(FLAGS.output_tfrecord)
  tf.io.gfile.makedirs(dirname)
  runner.run(
      construct_pipeline(FLAGS.input_tfrecord,
                         FLAGS.output_tfrecord,
                         FLAGS.sequence_key,
                         FLAGS.time_horizon,
                         FLAGS.subsample_context_features_rate,
                         FLAGS.reduce_image_size,
                         FLAGS.max_image_dimension,
                         FLAGS.add_context_features,
                         FLAGS.sorted_image_ids,
                         FLAGS.image_ids_to_keep,
                         FLAGS.keep_context_features_image_id_list,
                         FLAGS.keep_only_positives,
                         FLAGS.context_features_score_threshold,
                         FLAGS.keep_only_positives_gt,
                         FLAGS.max_num_elements_in_context_features,
                         FLAGS.num_shards,
                         FLAGS.output_type,
                         FLAGS.max_clip_length))

if __name__ == '__main__':
  flags.mark_flags_as_required([
      'input_tfrecord',
      'output_tfrecord'
  ])
  app.run(main)


text:
import cv2
import numpy as np

camera = cv2.VideoCapture(0)

number = 0
while True:
    ret, image = camera.read()
    if not ret:
        continue

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)  # Converts BGR image to HSV image for ease of use
    min_color = np.array([0, 0, 0])
    max_color = np.array([179, 239, 61])
    mask = cv2.inRange(hsv, min_color, max_color)  # Creates a mask of the image based on the min and max colors
    cv2.imshow("The Image", image)
    cv2.imshow("Mask", mask)

    key = cv2.waitKey(1) & 0xFF

    if key == ord('q'):
        break
    elif key == ord("s"):
        print(333)
        cv2.imwrite(f"Images/thresholding_images/image{number}.jpg", image)
        number += 1

camera.release()
cv2.destroyAllWindows()

ext: py

reason: import numpy


text:
#!/usr/bin/env python
'''
enrich_plot: contains a number of functions that use matplotlib to generate visualizations of enrich data.  This module is intended to be extensible, containing all enrich-related visuals
'''

__author__ = "Douglas M. Fowler"
__copyright__ = "Copyright 2011"
__credits__ = ["Douglas M Fowler", "Carlos L. Araya"]
__license__ = "FreeBSD"
__version__ = "0.2"
__maintainer__ = "Douglas M. Fowler"
__email__ = "dfowler@uw.edu"

import sys, os, time, math

try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pylab as pylab
    import matplotlib.cm as cm
    import matplotlib.pyplot as plt

except:
    print 'Error: Matplotlib not present. Cannot create visuals'

try:
    import numpy as np

except:
    print 'Error: Numpy not present. Cannot create visuals'

def freqplot(path, infile):
    '''
    freqplot: takes a mapunlink file and produces a heat map for visualizing diversity
    '''

    '''
    Basic I/O checks - making sure files exist, etc.
    '''
    try:
        # Check to make sure the input file exists
        f_infile = open(path + 'data/output/' + infile, 'U')
        f_infile.close()

    except:
        print 'Error: could not open input file'
        return 1

    try:
        plt_test = plt
        np_test = np.array

    except: # already warned the user about this above, so no message
        return 1

    f_infile = open(path + 'data/output/' + infile, 'U')
    header_line = f_infile.readline().split()

    # Make sure we can locate the position column
    try:
        pos_index = header_line.index('position')
    except:
        print 'Could not find the position index for plotting. Was the input file labeled properly?'
        return 1

    '''
    Reading in the data

    xs: the positions
    ys: the AA substitutions/PRO substitutions
    zs: frequencies
    values: a len(y) x len(x) matrix containing all the z values.
            Necessary for plotting the data as a heat map.
    '''

    xs = []
    ys = [index for index in xrange(len(header_line)-1)]
    zs = []

    # Grab all the data from the file and dump it into xs and zs
    for line in f_infile:
        line = line.split()
        xs.append(int(line[pos_index]))
        zs.extend([float(line[index]) for index in xrange(len(line)) if index is not pos_index])

    values = np.array(zs)
    values = values.reshape(len(xs), len(ys))
    values = values.transpose()

    '''
    Plotting the data

    Setting up graphical parameters for pyplot
    '''

    # Labels and label fonts

    xlab = 'Position'

    if 'PRO' in infile:
        ylab = 'Amino acid'
    elif 'DNA' in infile:
        ylab = 'Nucleotide'
    else:
        print 'Could not determine type of data present for plotting a heat map. Is the file name appropriately labeled?'
        return 1

    zlab = 'Frequency'

    title_font = {'fontsize': 18}

    # Tick marks - the padding is necessary so heat map boxes are not cut off

    plotys = ['']
    plotys.extend([header_line[index] for index in xrange(len(header_line)) if index != pos_index])
    plotys.append('')

    # Plotting the actual data

    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.matshow(values, cmap=cm.Blues)

    # Resizing the figure to be slightly larger

    DefaultSize = fig.get_size_inches()
    fig.set_size_inches( (DefaultSize[0]*1.5, DefaultSize[1]*1.5) )

    # Set the labels for the axes

    cax.axes.xaxis.set_label_text(xlab, **title_font)
    cax.axes.yaxis.set_label_text(ylab, **title_font)

    # Plotting the tick marks - the padding is necessary so boxes are not cut off
    len_y = [-1]
    len_y.extend(range(len(plotys)))
    len_y.append(len(plotys))

    pylab.yticks(len_y, plotys)

    # Drawing the color bar an appropriate size

    colorbar_args = {'shrink': 0.62}
    pylab.colorbar(cax, **colorbar_args)

    '''
    Saving

    Save the graph out as a pdf to disk
    '''
    fileout = infile + '_diversity.pdf'
    fileout_total = path + 'plots/' + fileout
    plt.title(fileout)

    pylab.savefig(fileout_total, dpi=100)

    return 0

def ratioplot(path, infile1, infile2):
    '''
    ratioplot: takes two mapunlink files, calculates and calculates the position-averaged ratio of the frequency of mutation between them
    '''

    '''
    Basic I/O checks - making sure files exist, etc.
    '''

    try: #check to make sure the input file exists
        f_infile = open(path + 'data/output/' + infile1, 'U')
        f_infile.close()
        f_infile = open(path + 'data/output/' + infile2, 'U')
        f_infile.close()

    except:
        print 'Error: could not open input file'
        return 1

    try:
        plt_test = plt
        np_test = np.array

    except: # already warned the user about this above, so no message
        return 1

    '''
    Reading the data: Calculate the ratio of frequencies for both files.
                      Assign the ratios appropriate colors based on whether the
                      ratio is above 1 or <= 1.

    file1_values/file2_values: dictionaries mapping {'position': frequency}
    '''
    f_infile = open(path + 'data/output/' + infile1, 'U')
    file1_values = read_position_file(f_infile)

    f_infile = open(path + 'data/output/' + infile2, 'U')
    file2_values = read_position_file(f_infile)

    ratios = []
    colors = []
    positions = []
    for position in file1_values:
        x = float(sum(file1_values[position]))
        y = float(sum(file2_values[position]))
        #rat = (x - position) / (y - position) WHY?
        if y != 0:
            rat = x / y

        else:
            rat = 1 #this is highly unsatisfactory, there should be some indication that this data is missing

        ratios.append(rat)
        positions.append(position)
        if (rat <= 1):
            colors.append('b')
        else:
            colors.append('r')

    '''
    Plotting the data: plot the ratios
    '''

    fig = plt.figure()
    fig.add_subplot(111)

    ax = plt.plot( [positions[index] for index in xrange(len(ratios)) if ratios[index] > 1],
             [ratios[index] for index in xrange(len(ratios)) if ratios[index] > 1], 'ro')
    ax = plt.plot( [positions[index] for index in xrange(len(ratios)) if ratios[index] <= 1],
             [ratios[index] for index in xrange(len(ratios)) if ratios[index] <= 1], 'bo')

    # Set the labels and other graphical parameters
    ylab = 'Position-averaged mutation frequency in selected/input'
    xlab = 'Position'

    plt.gca().set_xlim([min(positions) - 1, max(positions) + 1])
    plt.xlabel(xlab)
    plt.ylabel(ylab)

    # Draw the vertical lines that connect the baseline to each point
    plt.vlines(positions, [1] * len(positions), ratios, color=colors, linestyles='solid', lw=2)
    plt.axhline(y=1, color='gray')

    # Resize the image to be a little larger
    DefaultSize = fig.get_size_inches()
    fig.set_size_inches( (DefaultSize[0]*1.5, DefaultSize[1]*1.5) )

    # Draw the title manually because otherwise the default offset doesn't work
    fileout =  infile1 + "_" + infile2 + '_position_enrichment.pdf'
    fig.text(0.52, 0.97, fileout, ha="center", va="center", size="medium")

    '''
    Saving to disk
    '''
    fileout_total = path + 'plots/' + fileout

    pylab.savefig(fileout_total, dpi=100)

    return 0

def read_position_file(f_infile):
    '''
    reads a file into a dictionary mapping positions -> frequencies
    '''

    header = f_infile.readline().split()
    position_index = header.index('position')

    file_values = {}

    for line in f_infile:
        line = line.split()
        file_values[int(line[position_index])] = [float(line[index]) for index in xrange(len(line)) if index is not position_index]

    f_infile.close()
    return file_values

def read_substitution_file(f_infile):
    '''
    Reads in a file mapping
    '''
    header = f_infile.readline().split()
    position_index = header.index('position')

    file_values = {'positions' : []}

    for index in xrange(len(header)):
        if index is not position_index:
            file_values[header[index]] = {'values': []}

    for line in f_infile:
        line = line.split()
        [file_values['positions'].append(int(line[position_index]))]
        [file_values[header[index]]['values'].append(float(line[index])) for index in xrange(len(line)) if index is not position_index]

    f_infile.close()
    return file_values

def separate_data(substitution, positions, ratio, wt_data, axis_num, cols, rows):
    '''
    Another small helper function to follow DRY principles with the AA plot below.
    '''

    values = ratio[substitution]
    xr = xrange(len(values))

    # Separate out the values based on their category
    wt_x = [positions[index] for index in xr if (values[index] == 0 and wt_data['variable'][index] == substitution)]
    wt_y = [values[index] for index in xr if (values[index] == 0 and wt_data['variable'][index] == substitution)]
    nan_x = [positions[index] for index in xr if (values[index] == 0 and wt_data['variable'][index] != substitution)]
    nan_y = [values[index] for index in xr if (values[index] == 0 and wt_data['variable'][index] != substitution)]
    rest_x = [positions[index] for index in xr if values[index] != 0]
    rest_y = [values[index] for index in xr if values[index] != 0]

    ax = plt.subplot(rows, cols, axis_num)

    header_tick_values = range(0, len(values), (len(values)/4))

    all_residue_plot_data(wt_x, wt_y, nan_x, nan_y, rest_x, rest_y)
    adjust_axes(ax, str(substitution), header_tick_values, header_tick_values, axis_num, cols, 18)

def AAPlot(wt_data, ratio, axis_num, positions):

    '''
    Takes in ratios and maps how values should be plotted. Residues are
    plotted based on 3 categories: WT residues, residues with ratios of 1,
    and all other residues.

    Since this is a faux lattice plot, each substitution must be iterated through.
    '''

    rows = int(math.ceil(float(len(ratio.keys())) / 5.0))
    cols = 5

    keys = ratio.keys()
    keys = sorted(keys)

    for substitution in keys:
        if (substitution != '*'): # Save the stop index for last
            axis_num += 1
            # a : {0, 1, 2, 0, ... }
            separate_data(substitution, positions, ratio, wt_data, axis_num, cols, rows)

    '''
    Plotting the final substitution, the stop codon/index
    '''
    substitution = '*'
    axis_num += 1
    separate_data(substitution, positions, ratio, wt_data, axis_num, cols, rows)

def PosPlot(wt_data, ratio, axis_num, positions):

    '''
    Takes in ratios and maps how values should be plotted. Positions are
    plotted based on 3 categories: WT residues, residues with ratios of 1,
    and all other residues.

    Rather than re-read the data based on positions instead of substitutions, it makes sense to just iterate through the substitution dictionary and gather the necessary data.
    '''
    position_d = {'header' : []}

    '''
    Dealing with the stop position/codon, as well as mapping positions -> frequencies
    '''
    aa_sorted = ratio.keys()
    aa_sorted.sort()
    aa_sorted.pop(0)
    aa_sorted.append('*')

    for substitution in aa_sorted:
        if (substitution != '*'):
            position_d['header'].append(substitution)
            for index in xrange(len(ratio[substitution])):
                if position_d.get(index,None) == None:
                    position_d[index] = []
                position_d[index].append(ratio[substitution][index])

    position_d['header'].append('*')
    for index in xrange(len(ratio['*'])):
        if position_d.get(index, None) == None:
            position_d[index] = []
        position_d[index].append(ratio['*'][index])

    rows = int(math.ceil(float(len(position_d.keys())) / 5.0))
    cols = 5

    keys = [key for key in position_d if key != 'header']
    keys.sort()

    for position in keys:

        axis_num += 1
        # 1: {0.1, ...}

        values = position_d[position]
        xr = xrange(len(values))

        '''
        Separate categorizing from the AA plot above because we're indexing on position.
        '''
        wt_x = [positions[index] for index in xr if (values[index] == 0 and wt_data['variable'][position] == position_d['header'][index])]
        wt_y = [values[index] for index in xr if (values[index] == 0 and wt_data['variable'][position] == position_d['header'][index])]
        nan_x = [positions[index] for index in xr if (values[index] == 0 and wt_data['variable'][position] != position_d['header'][index])]
        nan_y = [values[index] for index in xr if (values[index] == 0 and wt_data['variable'][position] != position_d['header'][index])]
        rest_x = [positions[index] for index in xr if values[index] != 0]
        rest_y = [values[index] for index in xr if values[index] != 0]

        ax = plt.subplot(rows, cols, axis_num)

        all_residue_plot_data(wt_x, wt_y, nan_x, nan_y, rest_x, rest_y)
        adjust_axes(ax, str(position), range(len(position_d['header'])), position_d['header'], axis_num, cols, 13)

def adjust_axes(ax, title, header_ticks, header_list, axis_num, cols, axesfontsize):
    '''
    A small helper function that adjusts the axes on the plots
    '''

    frame = plt.gca()

    axes_font = {'fontsize':axesfontsize}
    title_font = {'fontsize':18}

    plt.text(0.1,3, title, **title_font)

    plt.axhline()

    if (axis_num <= cols) and (axis_num % 2 == 0):
        frame.xaxis.set_ticks_position('top')
        frame.xaxis.set_label_position('top')
        frame.xaxis.tick_top()
        ax.set_xticks(header_ticks)
        ax.set_xticklabels(header_list, **axes_font)
    else:
        frame.axes.get_xaxis().set_visible(False)

    if (((axis_num - 1) / cols) % 2 == 0) and axis_num % cols == 1:
        ax.set_yticks((-4,-2,0,2,4))
        ax.set_yticklabels([-4,-2,0,2,4], **axes_font)
    else:
        frame.axes.get_yaxis().set_visible(False)

def all_residue_plot_data(wt_x, wt_y, nan_x, nan_y, rest_x, rest_y):
    '''
    Another small helper function that assists with plotting the actual data.
    '''
    plt.plot(wt_x, wt_y, "rs")
    plt.plot(nan_x, nan_y, color="gray", marker="s")
    plt.plot(rest_x, rest_y, "bo")
    plt.ylim(-4, 4)
    xs = wt_x
    xs.extend(nan_x)
    xs.extend(rest_x)
    plt.xlim(min(xs) - 1, max(xs) + 1)

def all_residue_plot(path, infile1, infile2, wtseq, mode):
    '''
    all_residue_plot: takes two mapunlink files, calculates and calculates the ratio of the frequency of mutation between them for each mutation-position combination
    '''

    '''
    Basic I/O checks - check that files exist, etc.
    '''

    try:
        # Check to make sure the input file exists
        f_infile = open(path + 'data/output/' + infile1, 'U')
        f_infile.close()
        f_infile = open(path + 'data/output/' + infile2, 'U')
        f_infile.close()

    except:
        print 'Error: could not open input file'
        return 1

    try:
        plt_test = plt
        np_test = np.array

    except: # already warned the user about this above, so no message
        return 1
    '''
    Read in the data
    file1_values/file2_values: a dictionary maaping {'position': 'frequency'}
    wt_data: a dictionary containing the wildtype positions and AA/proteins
    '''

    f_infile = open(path + 'data/output/' + infile1, 'U')
    file1_values = read_substitution_file(f_infile)

    f_infile = open(path + 'data/output/' + infile2, 'U')
    file2_values = read_substitution_file(f_infile)

    # Obtain wildtype residues and set at the enrichment ratio for the wildtype sequence:
    wt_data = {'position': [], 'variable': []}
    for position in xrange(len(wtseq)):
        wt_data['position'].append(position)
        wt_data['variable'].append(wtseq[position])

    '''
    Calculating the log2 enrichment ratio for each mutation
    '''

    ratio = {}
    positions = file1_values['positions']

    for substitution in file1_values:
        if substitution != 'positions':
            f1p = file1_values[substitution] # Caching for faster lookup
            f2p = file2_values[substitution]

            if (ratio.get(substitution, None) == None):
                ratio[substitution] = []

            for index in xrange(len(f1p['values'])):
                if f2p['values'][index] != 0.0 and f1p['values'][index] != 0.0:
                    ratio[substitution].append(math.log(f1p['values'][index] / f2p['values'][index],2))
                else:
                    ratio[substitution].append(0)

    '''
    Setting up the graph to accomodate multiple sub plots and look similar
    to a lattice plot
    '''

    axis_num = 0
    fig = pylab.figure()
    fig.subplots_adjust(wspace=0.0001, hspace=0.0001)
    DefaultSize = fig.get_size_inches()
    fig.set_size_inches( (DefaultSize[0]*2.5, DefaultSize[1]*2.5) )

    '''
    Plot the data - see individual functions for more details
    '''
    ylab = ""
    if mode is 'AA':
        AAPlot(wt_data, ratio, axis_num, positions)
        ylab = "Position"
    elif mode is 'Pos':
        PosPlot(wt_data, ratio, axis_num, positions)
        ylab = "Amino Acid Substitution"

    # Drawing axis labels and other titles

    fileout = infile1 + "_" + infile2 + '_all_residue_by_' + mode + '.pdf'
    fileout_total = path + 'plots/' + fileout

    fig.text(0.1, 0.52, "log2(Mutation frequency in selected/input)", ha="right", va="center", size="xx-large", rotation="vertical")
    fig.text(0.52, 0.08, ylab, ha="center", va="bottom", size="xx-large")
    fig.text(0.52, 0.99, fileout, ha="center", va="top", size="xx-large")

    '''
    Saving to disk
    '''

    pylab.savefig(fileout_total, dpi=100)

    return 0


text:
import os
import numpy as np
from pyimagesearch_lrfinder import learningratefinder

def bestLR(model, trainData, startLR = 1e-10, endLR = 1e+1,
	batchSize=32, class_weight=None, epochs=5, sampleSize=2048,
	verbose=1, outfile=None):
	"""Returns the best LR based on the test from learningratefinder"""

	#save the model weights
	we = model.get_weights()

	#use pyImageSearch LR finder
	lrf = learningratefinder.LearningRateFinder(model)
	lrf.find(
		trainData = trainData,
		startLR = startLR, endLR = endLR,
		batchSize = batchSize,
		class_weight = class_weight,
		epochs = epochs,
		sampleSize = sampleSize,
		verbose = verbose)

	#restore initial weights
	model.set_weights(we)

	#plotting, if requested
	if outfile is not None:
		lrf.plot_loss(skipBegin=10, skipEnd=1, title="", outfile=outfile)

	#extracting and returning the best LR
	i = np.argmin(lrf.losses)
	return(lrf.lrs[i])


text:
import pycuda.autoinit
import pycuda.driver as drv
import pycuda.gpuarray as gpuarray
import pycuda.cumath
import skcuda.misc
import skcuda.linalg as linalg
from pycuda.compiler import SourceModule
import numpy as np
import matplotlib.pyplot as plt
import time
import scipy.io as sio
import pkg_resources
from NMF import *
from NMFJoint import *

def getResourceString(filename):
    fin = open(filename)
    s = fin.read()
    fin.close()
    return s.decode('utf8')

#Global handles to CUDA kernels
MatMulNaive_ = None
MatMulConv2D_ = None
ZerosToOnes_ = None
bitonicSortNonneg_ = None
TileWDenom_ = None
TileHDenom_ = None

def initParallelAlgorithms():
    """
    Compile all of the parallel algorithms
    """
    global MatMulNaive_
    global MatMulConv2D_
    s = getResourceString("MatMul.cu")
    mod = SourceModule(s)
    MatMulNaive_ = mod.get_function("MatMulNaive")
    MatMulConv2D_ = mod.get_function("MatMulConv2D")

    global ZerosToOnes_
    global TileWDenom_
    global TileHDenom_
    global bitonicSortNonneg_
    s = getResourceString("OtherUtils.cu")
    mod = SourceModule(s)
    ZerosToOnes_ = mod.get_function("ZerosToOnes")
    TileWDenom_ = mod.get_function("TileWDenom")
    TileHDenom_ = mod.get_function("TileHDenom")
    bitonicSortNonneg_ = mod.get_function("bitonicSortNonneg")

    linalg.init()
    skcuda.misc.init()

def plot3Diff(A, B, eps = None):
    """
    For debugging CPU vs GPU; plot a CPU result versus a GPU result
    along with a difference, with colorbars, to see if error is within
    numerical precision
    :param A: An MxN matrix from the CPU
    :param B: An MxN matrix from the GPU
    :param eps: Plot a black/white image with white pixels if the\
        difference is above eps
    """
    plt.subplot(131)
    plt.imshow(A, cmap = 'afmhot', interpolation='none', aspect='auto')
    plt.colorbar()
    plt.title("CPU")
    plt.subplot(132)
    plt.imshow(B, cmap = 'afmhot', interpolation='none', aspect='auto')
    plt.title("GPU")
    plt.colorbar()
    plt.subplot(133)
    diff = A-B
    if eps:
        diff2 = np.ones(diff.shape)
        diff2[np.abs(diff) < eps] = 0
        diff = diff2
    plt.imshow(diff, cmap = 'afmhot', interpolation='none', aspect='auto')
    plt.title("Diff")
    plt.colorbar()

def ZerosToOnes(V, eps = 0.0):
    """
    :param V: An MxN GPU array
    """
    blockdim = 32
    M = V.shape[0]
    N = V.shape[1]
    GridDimM = int(np.ceil(1.0*M/blockdim))
    GridDimN = int(np.ceil(1.0*N/blockdim))
    M = np.array(M, dtype = np.int32)
    N = np.array(N, dtype = np.int32)
    eps = np.array(eps, dtype = np.float32)
    ZerosToOnes_(V, M, N, eps, block=(blockdim, blockdim, 1), \
        grid=(GridDimM, GridDimN))

def bitonicSortNonneg(XG, blockdim = 32):
    M = np.int32(XG.shape[0])
    N = np.int32(XG.shape[1])
    NPow2 = np.int32(2**np.ceil(np.log2(N)))
    N2 = NPow2/2;
    #Pad with -1s as dummy values
    XGPow2 = -np.ones((M, NPow2), dtype = np.float32)
    XGPow2 = gpuarray.to_gpu(XGPow2)
    XGPow2[:, 0:N] = XG.copy()

    print("NPow2 = ", NPow2)
    GridDimM = int(np.ceil(1.0*M/blockdim))
    GridDimN = int(np.ceil(1.0*N2/blockdim))
    bitonicSortNonneg_(XG, XGPow2, M, N, NPow2, block=(blockdim, blockdim, 1), \
                grid=(GridDimM, GridDimN))

def getEuclideanErrorGPU(V, WH):
    """
    Return the Frobenius norm between V and W*H
    """
    diff = skcuda.misc.subtract(V, WH)
    diff = skcuda.misc.multiply(diff, diff)
    return skcuda.misc.sum(diff).get()

def getKLErrorGPU(V, WH, eps = 1e-10):
    """
    Return the Kullback-Liebler diverges between V and W*H
    """
    denom = WH.copy()
    ZerosToOnes(denom) #Prevent divide by zero
    arg = skcuda.misc.divide(V, denom)
    ZerosToOnes(arg, eps)
    res = skcuda.misc.multiply(V, pycuda.cumath.log(arg))
    res = skcuda.misc.subtract(res, V)
    res = skcuda.misc.add(res, WH)
    res = skcuda.misc.sum(res).get()
    print(res)
    return res

def multiplyConv2DWGradGPU(W, H, V, VLam, doDivision = True):
    """
    Compute the 2D convolutional multiplicative update for W using skcuda
    :param W: A TxNxK GPU array of K sources over spatiotemporal spans NxT\
    :param H: A FxKxM GPU array of source activations for each submatrix of W\
            over F transpositions over M time
    :param V: An MxN GPU target array
    :param VLam: An MxN GPU estimate array
    :param doDivision: If true, return the factor Numerator/Denomenator\
        otherwise, return (Numerator, Denomenator)
    """
    thisV = V.copy()
    thisVLam = VLam.copy()
    thisH = H.copy()
    WNums = gpuarray.zeros(W.shape, np.float32)
    WDenoms = gpuarray.zeros(W.shape, np.float32)
    for f in range(H.shape[0]):
        if f > 0:
            z = gpuarray.zeros((f, V.shape[1]), np.float32)
            thisV[0:-f, :] = V[f::, :]
            thisV[-f::, :] = z
            thisVLam[0:-f, :] = VLam[f::, :]
            thisVLam[-f::, :] = z
        for t in range(W.shape[0]):
            if t > 0:
                thisH[f, :, t::] = H[f, :, 0:-t]
                thisH[f, :, 0:t] = gpuarray.zeros((H.shape[1], t), np.float32)
            linalg.add_dot(thisV, thisH[f, :, :], WNums[t, :, :], transb='T')
            linalg.add_dot(thisVLam, thisH[f, :, :], WDenoms[t, :, :], transb='T')
    if doDivision:
        return skcuda.misc.divide(WNums, WDenoms)
    else:
        return (WNums, WDenoms)

def multiplyConv2DHGradGPU(W, H, V, VLam, doDivision = True):
    """
    Compute the 2D convolutional multiplicative update for H using skcuda
    :param W: A TxNxK GPU array of K sources over spatiotemporal spans NxT\
    :param H: A FxKxM GPU array of source activations for each submatrix of W\
            over F transpositions over M time
    :param V: An MxN GPU target array
    :param VLam: An MxN GPU estimate array
    :param doDivision: If true, return the factor Numerator/Denomenator\
        otherwise, return (Numerator, Denomenator)
    """
    thisV = V.copy()
    thisVLam = VLam.copy()
    thisW = W.copy()
    HNums = gpuarray.zeros(H.shape, np.float32)
    HDenoms = gpuarray.zeros(H.shape, np.float32)
    for t in range(W.shape[0]):
        if t > 0:
            #thisV = shiftMatLRUD(V, dj=-t)
            z = gpuarray.zeros((V.shape[0], t), np.float32)
            thisV[:, 0:-t] = V[:, t::]
            thisV[:, -t::] = z
            thisVLam[:, 0:-t] = VLam[:, t::]
            thisVLam[:, -t::] = z
        for f in range(H.shape[0]):
            if f > 0:
                #thisW = shiftMatLRUD(W[t, :, :], di=f)
                thisW[t, f::, :] = W[t, 0:-f, :]
                thisW[t, 0:f, :] = gpuarray.zeros((f, W.shape[2]), np.float32)
            linalg.add_dot(thisW[t, :, :], thisV, HNums[f, :, :], transa='T')
            linalg.add_dot(thisW[t, :, :], thisVLam, HDenoms[f, :, :], transa='T')
    if doDivision:
        return skcuda.misc.divide(HNums, HDenoms)
    else:
        return (HNums, HDenoms)

def TileWDenom(WDenomIn, M):
    """
    :param WDenomIn: A T x 1 x K array that needs to be tiled
    :param M: Dimension of tile axis
    :returns: WDenomOut: A T x M x K tiled array
    """
    blockdim = 32
    T = WDenomIn.shape[0]
    K = WDenomIn.shape[2]
    GridDimT = int(np.ceil(1.0*T/blockdim))
    GridDimK = int(np.ceil(1.0*K/blockdim))
    T = np.array(T, dtype = np.int32)
    M = np.array(M, dtype = np.int32)
    K = np.array(K, dtype = np.int32)
    WDenomOut = gpuarray.zeros((T, M, K), np.float32)
    TileWDenom_(WDenomIn, WDenomOut, T, M, K, block=(blockdim, blockdim, 1), \
        grid=(GridDimT, GridDimK))
    return WDenomOut

def multiplyConv2DWGradKLGPU(W, H, V, VLam, doDivision = True):
    """
    Compute the 2D convolutional multiplicative update for W under the
    Kullback-Liebler divergence, using skcuda
    :param W: A TxNxK matrix of K sources over spatiotemporal spans NxT\
    :param H: A FxKxM matrix of source activations for each submatrix of W\
            over F transpositions over M time
    :param VLam: Convolutional WH multiplication
    :param doDivision: If true, return the factor Numerator/Denomenator\
        otherwise, return (Numerator, Denomenator)
    :returns Ratio: A TxNkX matrix of multiplicative updates for W\
        or (RatioNum, RatioDenom) if doDivision = False
    """
    WNums = gpuarray.zeros(W.shape, np.float32)
    WDenoms = gpuarray.zeros((W.shape[0], W.shape[2]), np.float32)
    thisVLam = VLam.copy()
    ZerosToOnes(thisVLam)
    VLamQuot = skcuda.misc.divide(V, thisVLam)
    thisVLamQuot = VLamQuot.copy()
    thisH = H.copy()
    for f in range(H.shape[0]):
        if f > 0:
            thisVLamQuot[0:-f, :] = VLamQuot[f::, :]
            thisVLamQuot[-f::, :] = gpuarray.zeros((f, V.shape[1]), np.float32)
        for t in range(W.shape[0]):
            if t > 0:
                thisH[f, :, t::] = H[f, :, 0:-t]
                thisH[f, :, 0:t] = gpuarray.zeros((H.shape[1], t), np.float32)
            linalg.add_dot(thisVLamQuot, thisH[f, :, :], WNums[t, :, :], transb='T')
            WDenoms[t, :] = skcuda.misc.add(WDenoms[t, :], skcuda.misc.sum(thisH[f, :, :], 1))
    WDenoms = TileWDenom(WDenoms[:, None, :], W.shape[1])
    if doDivision:
        return skcuda.misc.divide(WNums, WDenoms)
    else:
        return (WNums, WDenoms)

def TileHDenom(HDenomIn, N):
    """
    :param HDenomIn: A F x K x 1 array that needs to be tiled
    :param N: Dimension of tile axis
    :returns: HDenomOut: A F x K x N tiled array
    """
    blockdim = 32
    F = HDenomIn.shape[0]
    K = HDenomIn.shape[1]
    GridDimF = int(np.ceil(1.0*F/blockdim))
    GridDimK = int(np.ceil(1.0*K/blockdim))
    F = np.array(F, dtype = np.int32)
    K = np.array(K, dtype = np.int32)
    N = np.array(N, dtype = np.int32)
    HDenomOut = gpuarray.zeros((F, K, N), np.float32)
    TileHDenom_(HDenomIn, HDenomOut, F, K, N, block=(blockdim, blockdim, 1), \
        grid=(GridDimF, GridDimK))
    return HDenomOut

def multiplyConv2DHGradKLGPU(W, H, V, VLam, doDivision = True):
    """
    Compute the 2D convolutional multiplicative update for H
    under the Kullback-Liebler divergence, using skcuda
    :param W: A TxNxK matrix of K sources over spatiotemporal spans NxT\
    :param H: A FxKxM matrix of source activations for each submatrix of W\
            over F transpositions over M time
    :param VLam: Convolutional WH multiplication
    :param doDivision: If true, return the factor Numerator/Denomenator\
        otherwise, return (Numerator, Denomenator)
    :returns Ratio: A FxKxM matrix of multiplicative updates for H\
        or (RatioNum, RatioDenom) if doDivision = False
    """
    HNums = gpuarray.zeros(H.shape, np.float32)
    HDenoms = gpuarray.zeros((H.shape[0], H.shape[1]), np.float32)
    thisVLam = VLam.copy()
    ZerosToOnes(thisVLam)
    VLamQuot = skcuda.misc.divide(V, thisVLam)
    thisVLamQuot = VLamQuot.copy()
    thisW = W.copy()
    for t in range(W.shape[0]):
        if t > 0:
            z = gpuarray.zeros((V.shape[0], t), np.float32)
            thisVLamQuot[:, 0:-t] = VLamQuot[:, t::]
            thisVLamQuot[:, -t::] = z
        for f in range(H.shape[0]):
            if f > 0:
                thisW[t, f::, :] = W[t, 0:-f, :]
                thisW[t, 0:f, :] = gpuarray.zeros((f, W.shape[2]), np.float32)
            linalg.add_dot(thisW[t, :, :], thisVLamQuot, HNums[f, :, :], transa='T')
            HDenoms[f, :] = skcuda.misc.add(HDenoms[f, :], skcuda.misc.sum(thisW[t, :, :], 0))
    HDenoms = TileHDenom(HDenoms[:, :, None], H.shape[2])
    if doDivision:
        return skcuda.misc.divide(HNums, HDenoms)
    else:
        return (HNums, HDenoms)

def multiplyConv2DGPU(W, H, blockdim = 32, Verbose = False):
    M = W.shape[1]
    N = H.shape[2]
    K = W.shape[2]
    T = W.shape[0]
    F = H.shape[0]
    sharedmem = 4*((F+blockdim)*T+(T+blockdim)*F)
    if Verbose:
        print("Shared Memory: %g kB"%(sharedmem/1024.0))

    #Figure out how to loop when copying over memory
    TBlockRound = blockdim*np.ceil(T/float(blockdim))
    FBlockRound = blockdim*np.ceil(F/float(blockdim))
    TBlocks = np.array(TBlockRound/blockdim, dtype=np.int32)
    FBlocks = np.array(FBlockRound/blockdim, dtype=np.int32)
    if Verbose:
        print("TBlocks = %i, FBlocks = %i"%(TBlocks, FBlocks))

    M = np.array(M, dtype=np.int32)
    N = np.array(N, dtype=np.int32)
    K = np.array(K, dtype=np.int32)
    T = np.array(T, dtype=np.int32)
    F = np.array(F, dtype=np.int32)

    GridDimM = int(np.ceil(1.0*M/blockdim))
    GridDimN = int(np.ceil(1.0*N/blockdim))
    Lam = gpuarray.zeros((M, N), np.float32)

    MatMulConv2D_(W, H, Lam, M, N, K, T, F, TBlocks, FBlocks, \
        block=(blockdim, blockdim, 1), \
        grid=(GridDimM, GridDimN, 1), shared=sharedmem )
    return Lam

def makeRandomGPUArray(M, N, L):
    return gpuarray.to_gpu(np.array(np.random.rand(M, N, L), dtype=np.float32))

def doNMF2DConvGPU(V, K, T, F, L, W = np.array([]), doKL = False, plotfn = None, \
    plotInterval = 60, plotFirst = False):
    """
    Implementing the Euclidean 2D NMF technique described in
    "Nonnegative Matrix Factor 2-D Deconvolution
        for Blind Single Channel Source Separation"
    This is the GPU version, so matrices are converted to GPU arrays at
    the beginning of the function
    :param V: An N x M target matrix
    :param K: Number of latent factors
    :param T: Time extent of W matrices
    :param F: Frequency extent of H matrices
    :param L: Number of iterations
    :param doKL: Whether to do Kullback-Leibler divergence.  If false, do Euclidean
    :param plotfn: A function used to plot each iteration, which should\
        take the arguments (V, W, H, iter) (assumed to take CPU arrays, \
        so conversion will be done in the function)
    :param plotInterval: At what interval to save plots of progress
    :returns (W, H): \
        W is an TxNxK CPU matrix of K sources over spatiotemporal spans NxT\
        H is a FxKxM CPU matrix of source activations for each submatrix of W\
            over F transpositions over M time
    """
    N = V.shape[0]
    M = V.shape[1]
    V = gpuarray.to_gpu(np.array(V, dtype=np.float32))
    WFixed = False
    if W.size == 0:
        W = makeRandomGPUArray(T, N, K)
    else:
        WFixed = True
        thisW = np.zeros(W.shape, dtype=np.float32)
        thisW[:, :, :] = W
        W = gpuarray.to_gpu(np.array(thisW, dtype=np.float32))
    H = makeRandomGPUArray(F, K, M)
    WH = multiplyConv2DGPU(W, H)
    errfn = getEuclideanErrorGPU
    WGradfn = multiplyConv2DWGradGPU
    HGradfn = multiplyConv2DHGradGPU
    if doKL:
        errfn = getKLErrorGPU
        WGradfn = multiplyConv2DWGradKLGPU
        HGradfn = multiplyConv2DHGradKLGPU
    errs = [errfn(V, WH)]
    if plotfn:
        res=4
        pK = K
        plt.figure(figsize=((4+pK)*res, res))
        if plotFirst:
            plotfn(V.get(), W.get(), H.get(), 0, errs)
            plt.savefig("NMF2DConv_%i.png"%0, bbox_inches = 'tight')
    for l in range(L):
        print("NMF iteration %i of %i"%(l+1, L))
        tic = time.time()
        #Step 1: Update Ws
        if not WFixed:
            VLam = multiplyConv2DGPU(W, H)
            WFac = WGradfn(W, H, V, VLam)
            W = skcuda.misc.multiply(W, WFac)

        #Step 2: Update Hs
        VLam = multiplyConv2DGPU(W, H)
        HFac = HGradfn(W, H, V, VLam)
        H = skcuda.misc.multiply(H, HFac)
        WH = multiplyConv2DGPU(W, H)
        errs.append(errfn(V, WH))
        if plotfn and ((l+1) == L or (l+1)%plotInterval == 0):
            plt.clf()
            plotfn(V.get(), W.get(), H.get(), l+1, errs)
            plt.savefig("NMF2DConv_%i.png"%(l+1), bbox_inches = 'tight')
        print("Elapsed Time: %.3g"%(time.time()-tic))
    return (W.get(), H.get())

def doNMF2DConvJointGPU(A, Ap, K, T, F, L, W1 = np.array([]), H = np.array([]), \
    doKL = False, plotfn = None, plotInterval = 60, plotFirst = False, \
    foldername = ".", errfn = getEuclideanErrorGPU):
    """
    Do a joint version of 2DNMF solving for W1, W2, and H, where
    A ~= W1*H and Ap ~= W2*H
    This is the GPU version, so matrices are converted to GPU arrays at
    the beginning of the function
    :param A: An N x M target matrix for dataset 1
    :param Ap: An N x M target matrix for dataset 2
    :param K: Number of latent factors
    :param T: Time extent of W matrices
    :param F: Frequency extent of H matrices
    :param L: Number of iterations
    :param W1: If this matrix is specified, hold it fixed
    :param H: If this matrix is specified, hold it fixed
    :param doKL: Whether to do Kullback-Leibler divergence.  If false, do Euclidean
    :param plotfn: A function used to plot each iteration, which should\
        take the arguments (V, W, H, iter)
    :returns (W1, W2, H): \
        W1 is an TxNxK matrix of K A sources over spatiotemporal spans NxT\
        W2 is an TxNxK matrix of K Ap sources over spatiotemporal spans NxT\
        H is a FxKxM matrix of source activations for each submatrix of W\
            over F transpositions over M time
    """
    N = A.shape[0]
    M = A.shape[1]
    A = gpuarray.to_gpu(np.array(A, dtype=np.float32))
    Ap = gpuarray.to_gpu(np.array(Ap, dtype=np.float32))
    W1Fixed = False
    HFixed = False
    if W1.size > 0:
        W1Fixed = True
        W1 = gpuarray.to_gpu(np.array(W1, dtype=np.float32))
    else:
        W1 = makeRandomGPUArray(T, N, K)
    W2 = makeRandomGPUArray(T, N, K)
    if H.size > 0:
        HFixed = True
        H = gpuarray.to_gpu(np.array(H, dtype=np.float32))
    else:
        H = makeRandomGPUArray(F, K, M)

    errfn = getEuclideanErrorGPU
    WGradfn = multiplyConv2DWGradGPU
    HGradfn = multiplyConv2DHGradGPU
    if doKL:
        errfn = getKLErrorGPU
        WGradfn = multiplyConv2DWGradKLGPU
        HGradfn = multiplyConv2DHGradKLGPU

    errs = [[errfn(A, multiplyConv2DGPU(W1, H)), errfn(Ap, multiplyConv2DGPU(W2, H))]]
    if plotfn:
        res=4
        plt.figure(figsize=((2+K)*res, 2*res))
        if plotFirst:
            plotfn(A.get(), Ap.get(), W1.get(), W2.get(), H.get(), 0, errs, foldername)
            plt.savefig("%s/NMF2DConvJoint_%i.png"%(foldername, 0), bbox_inches = 'tight')
    for l in range(L):
        print("Joint 2DNMF iteration %i of %i"%(l+1, L))
        tic = time.time()
        #Step 1: Update Ws
        if not W1Fixed:
            ALam = multiplyConv2DGPU(W1, H)
            WFac = WGradfn(W1, H, A, ALam)
            W1 = skcuda.misc.multiply(W1, WFac)
        ApLam = multiplyConv2DGPU(W2, H)
        WFac = WGradfn(W2, H, Ap, ApLam)
        W2 = skcuda.misc.multiply(W2, WFac)

        #Step 2: Update Hs
        if not HFixed:
            ALam = multiplyConv2DGPU(W1, H)
            ApLam = multiplyConv2DGPU(W2, H)
            (HNums1, HDenoms1)= HGradfn(W1, H, A, ALam, doDivision = False)
            (HNums2, HDenoms2)= HGradfn(W2, H, Ap, ApLam, doDivision = False)
            HNums = skcuda.misc.add(HNums1, HNums2)
            HDenoms = skcuda.misc.add(HDenoms1, HDenoms2)
            HFac = skcuda.misc.divide(HNums, HDenoms)
            H = skcuda.misc.multiply(H, HFac)

        #Plot results
        errs.append([errfn(A, multiplyConv2DGPU(W1, H)), errfn(Ap, multiplyConv2DGPU(W2, H))])
        if plotfn and ((l+1) == L or (l+1)%plotInterval == 0):
            plt.clf()
            plotfn(A.get(), Ap.get(), W1.get(), W2.get(), H.get(), l+1, errs, foldername)
            plt.savefig("%s/NMF2DConvJoint_%i.png"%(foldername, l+1), bbox_inches = 'tight')
        print("Elapsed Time: %.3g"%(time.time()-tic))
    return (W1.get(), W2.get(), H.get())

def doNMF2DConvJoint3WayGPU(A, Ap, B, K, T, F, L, doKL = False, plotfn = None, \
         foldername = ".", plotInterval = 60, plotFirst = False):
    """
    A version of 2DNMF that jointly solves for W1, H1, W2, and H2 that
    satisfy W1*H1 ~= A, W2*H1 ~= Ap, and W2*H1 ~= B
    :param A: An M x N1 matrix for song A
    :param Ap: An M x N1 matrix for song A'
    :param B: An M x N2 matrix for song B
    :param K: Number of latent factors
    :param T: Time extent of W matrices
    :param F: Frequency extent of H matrices
    :param L: Number of iterations
    :param doKL: Whether to do Kullback-Leibler divergence.  If false, do Euclidean
    :param plotfn: A function used to plot each iteration, which should\
        take the arguments (V, W, H, iter)
    :returns (W1, W2, H1, H2): \
        W1, W2 are an MxKxT matrices of K sources over spatiotemporal spans MxT\
        H1 is a KxN1xF matrix of source activations for the A, A' pair
        H2 is a KxN2xF matrix of source activations for B
    """
    if not (A.shape[1] == Ap.shape[1]):
        print("Error: A and A' should have same number of frames")
        return None
    M = A.shape[0]
    N1 = A.shape[1]
    N2 = B.shape[1]
    W1 = makeRandomGPUArray(T, M, K)
    W2 = makeRandomGPUArray(T, M, K)
    H1 = makeRandomGPUArray(F, K, N1)
    H2 = makeRandomGPUArray(F, K, N2)

    #Convert to GPU Arrays
    A = gpuarray.to_gpu(np.array(A, dtype=np.float32))
    Ap = gpuarray.to_gpu(np.array(Ap, dtype=np.float32))
    B = gpuarray.to_gpu(np.array(B, dtype=np.float32))

    errfn = getEuclideanErrorGPU
    WGradfn = multiplyConv2DWGradGPU
    HGradfn = multiplyConv2DHGradGPU
    if doKL:
        errfn = getKLErrorGPU
        WGradfn = multiplyConv2DWGradKLGPU
        HGradfn = multiplyConv2DHGradKLGPU

    errs = [getJoint3WayError(A, Ap, B, W1, W2, H1, H2, errfn, multiplyConv2DGPU)]
    if plotfn:
        res=4
        plt.figure(figsize=((8+2*K)*res*1.2, 2*res))
        if plotFirst:
            plotfn(A.get(), Ap.get(), B.get(), W1.get(), W2.get(), H1.get(), H2.get(), \
                    0, errs, foldername)
            plt.savefig("%s/NMF2DConvJoint3Way_%i.png"%(foldername, 0), bbox_inches = 'tight')
    for l in range(L):
        print("Joint 2DNMF 3 way iteration %i of %i"%(l+1, L))
        tic = time.time()
        #Step 1: Update Ws
        Lam11 = multiplyConv2DGPU(W1, H1)
        Lam12 = multiplyConv2DGPU(W1, H2)
        Lam21 = multiplyConv2DGPU(W2, H1)

        #Update W1
        (N11, D11) = WGradfn(W1, H1, A, Lam11, doDivision = False)
        (N12, D12) = WGradfn(W1, H2, B, Lam12, doDivision = False)
        Num = skcuda.misc.add(N11, N12)
        Denom = skcuda.misc.add(D11, D12)
        Fac = skcuda.misc.divide(Num, Denom)
        W1 = skcuda.misc.multiply(W1, Fac)
        #Update W2
        Fac = WGradfn(W2, H1, Ap, Lam21, doDivision = True)
        W2 = skcuda.misc.multiply(W2, Fac)

        #Step 2: Update Hs
        #Update H1
        Lam11 = multiplyConv2DGPU(W1, H1)
        Lam12 = multiplyConv2DGPU(W1, H2)
        Lam21 = multiplyConv2DGPU(W2, H1)
        (N11, D11) = HGradfn(W1, H1, A, Lam11, doDivision = False)
        (N12, D12) = HGradfn(W2, H1, Ap, Lam21, doDivision = False)
        Num = skcuda.misc.add(N11, N12)
        Denom = skcuda.misc.add(D11, D12)
        Fac = skcuda.misc.divide(Num, Denom)
        H1 = skcuda.misc.multiply(H1, Fac)
        #Update H2
        Fac = HGradfn(W1, H2, B, Lam12, doDivision = True)
        H2 = skcuda.misc.multiply(H2, Fac)

        errs.append(getJoint3WayError(A, Ap, B, W1, W2, H1, H2, errfn, multiplyConv2DGPU))
        print("Elapsed Time: %.3g"%(time.time()-tic))
        if plotfn and ((l+1) == L or (l+1)%plotInterval == 0):
            plt.clf()
            plotfn(A.get(), Ap.get(), B.get(), W1.get(), W2.get(), H1.get(), H2.get(), \
                    l+1, errs, foldername)
            plt.savefig("%s/NMF2DConvJoint3Way_%i.png"%(foldername, l+1), bbox_inches = 'tight')
    return (W1.get(), W2.get(), H1.get(), H2.get())

def doNMFDriedgerGPU(V, W, L, r = 7, p = 10, c = 3, plotfn = None):
    ####TODO: This GPU translation is incomplete at the moment
    """
    Implement the technique from "Let It Bee-Towards NMF-Inspired
    Audio Mosaicing" on the GPU
    :param V: M x N target matrix
    :param W: An M x K matrix of template sounds in some time order\
        along the second axis
    :param L: Number of iterations
    :param r: Width of the repeated activation filter
    :param p: Degree of polyphony; i.e. number of values in each column\
        of H which should be un-shrunken
    :param c: Half length of time-continuous activation filter
    :return H: A KxM array of activations
    """
    import scipy.ndimage
    N = V.shape[0]
    M = V.shape[1]
    K = W.shape[1]
    tic = time.time()
    V = gpuarray.to_gpu(np.array(V, dtype=np.float32))
    W = gpuarray.to_gpu(np.array(W, dtype=np.float32))
    H = gpuarray.to_gpu(np.array(np.random.rand(K, M), dtype=np.float32))
    errs = [getKLErrorGPU(V, linalg.dot(W, H))]
    if plotfn:
        res=4
        plt.figure(figsize=(res*5, res))
        plotfn(V.get(), W.get(), H.get(), 0, errs)
        plt.savefig("NMFDriedger_%i.png"%0, bbox_inches = 'tight')

    for l in range(L):
        print("NMF Driedger iteration %i of %i"%(l+1, L))
        iterfac = 1-float(l+1)/L
        #Step 1: Avoid repeated activations
        print("Doing Repeated Activations...")
        MuH = scipy.ndimage.filters.maximum_filter(H, size=(1, r))
        H[H<MuH] = H[H<MuH]*iterfac
        #Step 2: Restrict number of simultaneous activations
        print("Restricting simultaneous activations...")
        colCutoff = -np.sort(-H, 0)[p, :]
        H[H < colCutoff[None, :]] = H[H < colCutoff[None, :]]*iterfac
        #Step 3: Supporting time-continuous activations
        if c > 0:
            print("Supporting time-continuous activations...")
            for k in range(-H.shape[0]+1, H.shape[1]):
                z = np.cumsum(np.concatenate((np.zeros(c), np.diag(H, k), np.zeros(c))))
                x2 = z[2*c::] - z[0:-2*c]
                H[np.diag(I, k), np.diag(J, k)] = x2
        #KL Divergence Version
        tic = time.time()
        WH = W.dot(H)
        WH[WH == 0] = 1
        VLam = V/WH
        print("VLam Elapsed Time: %.3g"%(time.time() - tic))
        tic = time.time()
        WDenom = np.sum(W, 0)
        WDenom[WDenom == 0] = 1
        H = H*((W.T).dot(VLam)/WDenom[:, None])
        print("Elapsed Time H Update %.3g"%(time.time() - tic))
        errs.append(getKLError(V, W.dot(H)))
        if plotfn and ((l+1)==L):# or (l+1)%20 == 0):
            plt.clf()
            plotfn(V, W, H, l+1, errs)
            plt.savefig("NMDriedger_%i.png"%(l+1), bbox_inches = 'tight')
    return H

def testNMF2DMultiplyGPU():
    initParallelAlgorithms()
    np.random.seed(100)
    M = 1025
    K = 10
    T = 12
    F = 40
    N = 20000
    W = np.random.rand(T, M, K)
    H = np.random.rand(F, K, N)

    tic = time.time()
    LamGT = multiplyConv2D(W, H)
    cputime = time.time()-tic
    print("Elapsed Time CPU: %.3g"%cputime)

    tic = time.time()
    WGPU = gpuarray.to_gpu(np.array(W, dtype=np.float32))
    HGPU = gpuarray.to_gpu(np.array(H, dtype=np.float32))
    Lam = multiplyConv2DGPU(WGPU, HGPU, blockdim = 32, Verbose = True)
    Lam = Lam.get()
    print("AllClose: ", np.allclose(Lam, LamGT))
    gputime = time.time()-tic
    print("Elapsed Time GPU: %.3g"%gputime)
    print("Speedup Ratio: %.3g"%(cputime/gputime))
    plt.figure(figsize=(16, 4))
    plot3Diff(LamGT, Lam)
    #plt.savefig("fig.png", bbox_inches='tight')
    plt.show()

def testNMF2DWGradientGPU():
    initParallelAlgorithms()
    np.random.seed(100)
    blockdim = 32
    M = 1025
    K = 10
    T = 20
    F = 40
    N = 10000
    V = np.random.rand(M, N)
    W = np.random.rand(T, M, K)
    H = np.random.rand(F, K, N)
    VLam = multiplyConv2D(W, H)

    sharedmem = 4*((blockdim+F)*blockdim*2+(T+blockdim)*F)
    print("Shared Memory: %g kB"%(sharedmem/1024.0))

    tic = time.time()
    HFacGT = multiplyConv2DHGrad(W, H, V, VLam)
    cputime = time.time()-tic
    print("Elapsed Time CPU: %.3g"%cputime)

    tic = time.time()
    WGPU = gpuarray.to_gpu(np.array(W, dtype=np.float32))
    HGPU = gpuarray.to_gpu(np.array(H, dtype=np.float32))
    VGPU = gpuarray.to_gpu(np.array(V, dtype=np.float32))
    VLamGPU = gpuarray.to_gpu(np.array(VLam, dtype=np.float32))

    HFac = multiplyConv2DHGradGPU(WGPU, HGPU, VGPU, VLamGPU)
    HFac = HFac.get()
    gputime = time.time()-tic
    print("Elapsed Time GPU: %.3g"%gputime)
    print("Speedup Ratio: %.3g"%(cputime/gputime))
    plt.figure(figsize=(16, 4))
    print(HFac.shape)
    print("AllClose: ", np.allclose(HFacGT, HFac))
    plot3Diff(HFacGT[:, 0, :], HFac[:, 0, :])
    plt.show()

def testKLGradient():
    initParallelAlgorithms()
    M = 1025
    K = 10
    T = 20
    F = 40
    N = 10000
    V = np.random.rand(M, N)
    W = np.random.rand(T, M, K)
    H = np.random.rand(F, K, N)

    WGPU = gpuarray.to_gpu(np.array(W, dtype=np.float32))
    HGPU = gpuarray.to_gpu(np.array(H, dtype=np.float32))
    VGPU = gpuarray.to_gpu(np.array(V, dtype=np.float32))
    VLamGPU = multiplyConv2DGPU(WGPU, HGPU)
    VLam = VLamGPU.get()

    ##Step 1: Test W Gradient
    #Test separating out numerators and denomenators
    tic = time.time()
    (WNums, WDenoms) = multiplyConv2DWGradKL(W, H, V, VLam, doDivision = False)
    print("Elapsed Time CPU: %.3g"%(time.time() - tic))
    tic = time.time()
    (WNumsGPU, WDenomsGPU) = multiplyConv2DWGradKLGPU(WGPU, HGPU, VGPU, VLamGPU, doDivision = False)
    print("Elapsed Time GPU: %.3g"%(time.time() - tic))
    print("All Close W Nums: ", np.allclose(WNums, WNumsGPU.get()))

    #Test the ratio
    tic = time.time()
    WRatio = multiplyConv2DWGradKL(W, H, V, VLam)
    print("Elapsed Time CPU: %.3g"%(time.time() - tic))
    tic = time.time()
    WRatioGPU = multiplyConv2DWGradKLGPU(WGPU, HGPU, VGPU, VLamGPU)
    print("Elapsed Time GPU: %.3g"%(time.time() - tic))
    print("All Close W Ratio: ", np.allclose(WRatio, WRatioGPU.get()))

    ##Step 2: Test H Gradient
    #Test separating out numerators and denomenators
    tic = time.time()
    (HNums, HDenoms) = multiplyConv2DHGradKL(W, H, V, VLam, doDivision = False)
    print("Elapsed Time CPU: %.3g"%(time.time() - tic))
    tic = time.time()
    (HNumsGPU, HDenomsGPU) = multiplyConv2DHGradKLGPU(WGPU, HGPU, VGPU, VLamGPU, doDivision = False)
    print("Elapsed Time GPU: %.3g"%(time.time() - tic))
    print("All Close H Nums: ", np.allclose(HNums, HNumsGPU.get()))

    #Test the ratio
    tic = time.time()
    HRatio = multiplyConv2DHGradKL(W, H, V, VLam)
    print("Elapsed Time CPU: %.3g"%(time.time() - tic))
    tic = time.time()
    HRatioGPU = multiplyConv2DHGradKLGPU(WGPU, HGPU, VGPU, VLamGPU)
    print("Elapsed Time GPU: %.3g"%(time.time() - tic))
    print("All Close H Ratio: ", np.allclose(HRatio, HRatioGPU.get()))

def testDriedgerComponents():
    initParallelAlgorithms()
    np.random.seed(100)
    X = np.array(np.random.randn(1723, 22519), dtype=np.float32)
    XGPU = gpuarray.to_gpu(X)
    tic = time.time()
    XSorted = np.sort(X, 1)
    print("Elapsed Time CPU: %.3g"%(time.time()-tic))
    tic = time.time()
    bitonicSortNonneg(XGPU)
    print("Elapsed Time GPU: %.3g"%(time.time()-tic))
    plot3Diff(XSorted, XGPU.get())
    plt.show()

if __name__ == '__main__':
    #testNMF2DMultiplyGPU()
    #testNMF2DWGradientGPU()
    #testKLGradient()
    testDriedgerComponents()

max_stars_repo_name: ctralie/CoverSongSynthesis

max_issues_repo_name: ctralie/CoverSongSynthesis

max_forks_repo_name: ctralie/CoverSongSynthesis

reason: import numpy,import scipy,import pycuda,from pycuda


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

from typing import Any, Sized, Union

import numpy as np

Array = Union['np.ndarray[Any, np.dtype[Any]]']

class FakeGenericMeta(type):
    def __getitem__(cls, _item):
        return cls

_zip_strict_sentinel = object()

def _zip_strict_iters(*iterables):
    if not iterables:
        return
    iterators = tuple(map(iter, iterables))
    del iterables
    while True:
        pops = tuple(next(it, _zip_strict_sentinel) for it in iterators)
        if any(pop is _zip_strict_sentinel for pop in pops):
            break
        yield pops
    not_exhausted = tuple(i for i, pop in enumerate(pops) if pop is not _zip_strict_sentinel)
    if not_exhausted:
        raise ValueError('Some zip_strict arguments not exhausted: {}'.format(not_exhausted))

# Guess who's getting impatient waiting for Python 3.10?
def zip_strict(*iterables):
    if all(isinstance(it, Sized) for it in iterables):
        if len(set(map(len, iterables))) > 1:
            raise ValueError('Arguments to zip_strict have inconsistent lengths: {}'.format(tuple(map(len, iterables))))
        return zip(*iterables)
    return _zip_strict_iters(*iterables)

def zipstar(it):
    return tuple(zip(*it))

def zipstar_strict(it):
    return tuple(zip_strict(*it))

max_stars_repo_name: Cebtenzzre/cebnn

max_issues_repo_name: Cebtenzzre/cebnn

max_forks_repo_name: Cebtenzzre/cebnn