text:
from Evaluate import Evaluate
from data.Data import Data
from helpers.Runs import Runs
import tensorflow as tf
from models.Lstm import Lstm
from models.Cnn import Cnn
import numpy as np

class Model:
    def __init__(self, params):
        self.params = params
        self.cnn = Cnn(params)

        # helper class for storing run details
        self.runs = Runs()

    def start(self):
        # tf Graph
        self.x = tf.placeholder("float", [None, self.params["post_padding_size"], self.params["comment_padding_size"],
                                          self.params["word2vec_dim"]], name="input_x")
        self.y_sentiment = tf.placeholder("float",
                                          [None, self.params["post_padding_size"], self.params["n_classes_sentiment"]],
                                          name="input_y_sentiment")
        self.y_topics = tf.placeholder("float",
                                       [None, self.params["post_padding_size"], self.params["n_classes_topics"]],
                                       name="input_y_topics")
        self.y_emotion = tf.placeholder("float",
                                        [None, self.params["post_padding_size"], self.params["n_classes_emotion"]],
                                        name="input_y_emotion")
        self.y_speech_acts = tf.placeholder("float",
                                            [None, self.params["post_padding_size"],
                                             self.params["n_classes_speech_acts"]],
                                            name="input_y_speech_acts")

        self.keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
        self.sequence_length = tf.placeholder(tf.int32, [None])

        fully_connected_params = {
            "out_sentiment_w": tf.Variable(
                tf.random_normal([self.params["n_hidden"], self.params["n_classes_sentiment"]])),
            "out_topics_w": tf.Variable(tf.random_normal([self.params["n_hidden"], self.params["n_classes_topics"]])),
            "out_emotion_w": tf.Variable(tf.random_normal([self.params["n_hidden"], self.params["n_classes_emotion"]])),
            "out_speech_acts_w": tf.Variable(
                tf.random_normal([self.params["n_hidden"], self.params["n_classes_speech_acts"]])),

            "out_sentiment_b": tf.Variable(tf.random_normal([self.params["n_classes_sentiment"]])),
            "out_topics_b": tf.Variable(tf.random_normal([self.params["n_classes_topics"]])),
            "out_emotion_b": tf.Variable(tf.random_normal([self.params["n_classes_emotion"]])),
            "out_speech_acts_b": tf.Variable(tf.random_normal([self.params["n_classes_speech_acts"]]))
        }

        self.lstm = Lstm(params=self.params, fully_connected_params=fully_connected_params)

        # get predictions
        self.predictions = self.lstm.model(x=self.cnn.model(self.x, self.keep_prob), sequence_length=self.sequence_length,
                                      keep_prob=self.keep_prob)

        # define loss
        with tf.name_scope("loss_sentiment"):
            self.cost_sentiment = tf.reduce_mean(
                tf.nn.softmax_cross_entropy_with_logits(logits=self.predictions["prediction_sentiment"],
                                                        labels=self.y_sentiment))

        with tf.name_scope("loss_topics"):
            self.cost_topics = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(logits=self.predictions["prediction_topics"], labels=self.y_topics))

        with tf.name_scope("loss_emotions"):
            self.cost_emotions = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(logits=self.predictions["prediction_emotion"],
                                                        labels=self.y_emotion))

        with tf.name_scope("loss_speech_acts"):
            self.cost_speech_acts = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(logits=self.predictions["prediction_speech_acts"],
                                                        labels=self.y_speech_acts))

        # define optimizer
        self.optimizer = tf.train.AdamOptimizer(learning_rate=self.params["learning_rate"]).minimize(
            self.cost_sentiment + self.cost_topics + self.cost_emotions + self.cost_speech_acts)

        # evaluate model
        with tf.name_scope("accuracy_sentiment"):
            correct_pred_sentiment = tf.equal(tf.argmax(self.predictions["prediction_sentiment"], 1),
                                              tf.argmax(self.y_sentiment, 1))
            self.accuracy_sentiment = tf.reduce_mean(tf.cast(correct_pred_sentiment, tf.float32))

        with tf.name_scope("accuracy_topics"):
            correct_pred_topics = tf.equal(tf.round(tf.nn.sigmoid(self.predictions["prediction_topics"])),
                                           tf.round(self.y_topics))
            self.accuracy_topics = tf.reduce_mean(tf.cast(correct_pred_topics, tf.float32))

        with tf.name_scope("accuracy_emotion"):
            correct_pred_emotion = tf.equal(tf.round(tf.nn.sigmoid(self.predictions["prediction_emotion"])),
                                            tf.round(self.y_emotion))
            self.accuracy_emotion = tf.reduce_mean(tf.cast(correct_pred_emotion, tf.float32))

        with tf.name_scope("accuracy_speech_acts"):
            correct_pred_speech_acts = tf.equal(tf.round(tf.nn.sigmoid(self.predictions["prediction_speech_acts"])),
                                                tf.round(self.y_speech_acts))
            self.accuracy_speech_acts = tf.reduce_mean(tf.cast(correct_pred_speech_acts, tf.float32))

        # initializing the variables
        self.init = tf.global_variables_initializer()
        # 'Saver' op to save and restore all the variables
        self.saver = tf.train.Saver()

        # get data object
        self.data = Data(filename='data/word2vec/wiki.hr.vec',
                         comment_padding_size=self.params["comment_padding_size"],
                         post_padding_size=self.params["post_padding_size"],
                         word2vec_dim=self.params["word2vec_dim"],
                         binary_sentiment=self.params["binary_sentiment"])

        self.runs.create_run()
        # START LEARNING!!!
        self.learn()

    def learn(self):
        with tf.Session() as sess:
            self.sess = sess
            # initialize session
            sess.run(self.init)

            step = 1
            counter = 0
            for epoch in range(0, self.params["max_epoch"]):
                if epoch % self.params["evaluate_every"] == 0:
                    self.evaluate()
                    self.runs.save_model(sess=self.sess, saver=self.saver)
                batch_x = []
                batch_seq_length = []
                batch_y_sentiment = []
                batch_y_topics = []
                batch_y_emotions = []
                batch_y_speech_acts = []

                with open('data/threads/splits/split-0/train.txt', encoding="UTF-8") as f:
                    for line in f:
                        x, sequence_length_next, y_sentiment_next, y_topics_next, y_emotion_next, y_speech_acts_next = \
                            self.data.get_next(line)
                        batch_x.append(x)
                        batch_seq_length.append(sequence_length_next)
                        batch_y_sentiment.append(y_sentiment_next)
                        batch_y_topics.append(y_topics_next)
                        batch_y_emotions.append(y_emotion_next)
                        batch_y_speech_acts.append(y_speech_acts_next)
                        counter += 1
                        if len(batch_x) == self.params["batch_size"]:
                            # turn input to np.array
                            batch_x = np.array(batch_x)
                            batch_y_sentiment = np.array(batch_y_sentiment)
                            batch_seq_length = np.array(batch_seq_length)
                            # reshape input
                            batch_x = batch_x.reshape(
                                (self.params["batch_size"], self.params["post_padding_size"],
                                 self.params["comment_padding_size"], self.params["word2vec_dim"]))
                            batch_y_sentiment = batch_y_sentiment.reshape(
                                (self.params["batch_size"], self.params["post_padding_size"],
                                 self.params["n_classes_sentiment"]))
                            # TRAIN HERE
                            sess.run(self.optimizer,
                                     feed_dict={self.x: batch_x,
                                                self.y_sentiment: batch_y_sentiment,
                                                self.y_topics: batch_y_topics,
                                                self.y_emotion: batch_y_emotions,
                                                self.y_speech_acts: batch_y_speech_acts,
                                                self.sequence_length: batch_seq_length,
                                                self.keep_prob: self.params["keep_prob_global_train"]})

                            step += 1

                            if step % self.params["display_step"] == 0:
                                # SENTIMENT
                                acc_sentiment = sess.run(self.accuracy_sentiment,
                                                         feed_dict={self.x: batch_x,
                                                                    self.y_sentiment: batch_y_sentiment,
                                                                    self.sequence_length: batch_seq_length,
                                                                    self.keep_prob: self.params[
                                                                        "keep_prob_global_train"]})
                                loss_sentiment = sess.run(self.cost_sentiment,
                                                          feed_dict={self.x: batch_x,
                                                                     self.y_sentiment: batch_y_sentiment,
                                                                     self.sequence_length: batch_seq_length,
                                                                     self.keep_prob: self.params[
                                                                         "keep_prob_global_train"]})
                                print(
                                    "Epoch: " + str(epoch + 1) + " Iteration: " + str(step * self.params["batch_size"]))
                                print("[SENTIMENT]  " + " Minibatch Loss= ""{:.4f}".format(
                                    loss_sentiment) + ", Minibatch Accuracy= ""{:.4f}".format(acc_sentiment))

                                # EMOTIONS
                                acc_emotion = sess.run(self.accuracy_emotion,
                                                       feed_dict={self.x: batch_x, self.y_emotion: batch_y_emotions,
                                                                  self.sequence_length: batch_seq_length,
                                                                  self.keep_prob: self.params[
                                                                      "keep_prob_global_train"]})
                                loss_emotion = sess.run(self.cost_emotions,
                                                        feed_dict={self.x: batch_x, self.y_emotion: batch_y_emotions,
                                                                   self.sequence_length: batch_seq_length,
                                                                   self.keep_prob: self.params[
                                                                       "keep_prob_global_train"]})
                                print("[EMOTION]     " + " Minibatch Loss= ""{:.4f}".format(
                                    loss_emotion) + ", Minibatch Accuracy= ""{:.4f}".format(acc_emotion))

                                # TOPICS
                                acc_topics = sess.run(self.accuracy_topics,
                                                      feed_dict={self.x: batch_x, self.y_topics: batch_y_topics,
                                                                 self.sequence_length: batch_seq_length,
                                                                 self.keep_prob: self.params["keep_prob_global_train"]})
                                loss_topics = sess.run(self.cost_topics,
                                                       feed_dict={self.x: batch_x, self.y_topics: batch_y_topics,
                                                                  self.sequence_length: batch_seq_length,
                                                                  self.keep_prob: self.params[
                                                                      "keep_prob_global_train"]})
                                print("[TOPICS]     " + " Minibatch Loss= ""{:.4f}".format(
                                    loss_topics) + ", Minibatch Accuracy= ""{:.4f}".format(acc_topics))

                                # SPEECH ACTS
                                acc_speech_acts = sess.run(self.accuracy_speech_acts,
                                                           feed_dict={self.x: batch_x,
                                                                      self.y_speech_acts: batch_y_speech_acts,
                                                                      self.sequence_length: batch_seq_length,
                                                                      self.keep_prob: self.params[
                                                                          "keep_prob_global_train"]})
                                loss_speech_acts = sess.run(self.cost_speech_acts,
                                                            feed_dict={self.x: batch_x,
                                                                       self.y_speech_acts: batch_y_speech_acts,
                                                                       self.sequence_length: batch_seq_length,
                                                                       self.keep_prob: self.params[
                                                                           "keep_prob_global_train"]})
                                print("[SPEECH ACTS] " + "Minibatch Loss= ""{:.4f}".format(
                                    loss_speech_acts) + ", Minibatch Accuracy= ""{:.4f}".format(acc_speech_acts))

                                print("")

                            # RESET BATCH
                            batch_x = []
                            batch_seq_length = []
                            batch_y_sentiment = []
                            batch_y_topics = []
                            batch_y_emotions = []
                            batch_y_speech_acts = []
            self.evaluate()

    def evaluate(self):
        evaluate = Evaluate(self.data, self.params, self.predictions, self.sess, self.x, self.sequence_length,
                            self.keep_prob, self.y_sentiment, self.y_emotion, self.y_topics, self.y_speech_acts)
        evaluate.execute_evaluation("TRAIN", "train.txt")
        evaluate.execute_evaluation("TEST", "test.txt")


text:
import numpy as np
import os
import time
import emcee
from scipy.optimize import fmin

# kepler_grids
from . import burstfit
from . import mcmc_versions
from . import mcmc_tools
from pyburst.misc.pyprint import check_params_length

GRIDS_PATH = os.environ['KEPLER_GRIDS']
MODELS_PATH = os.environ['KEPLER_MODELS']

def setup_sampler(source, version, pos=None, n_walkers=None, n_threads=1,
                  **kwargs):
    """Initialises and returns EnsembleSampler object

    NOTE: Only uses pos to get n_walkers and n_dimensions
    """
    if pos is None:
        if n_walkers is None:
            print('ERROR: must provide either pos or n_walkers')
        pos = setup_positions(source=source, version=version, n_walkers=n_walkers)

    n_walkers = len(pos)
    n_dimensions = len(pos[0])

    bfit = burstfit.BurstFit(source=source, version=version, verbose=False,
                             re_interp=False, **kwargs)

    sampler = emcee.EnsembleSampler(n_walkers, n_dimensions, bfit.lhood,
                                    threads=n_threads)
    return sampler

def setup_positions(source, version, n_walkers, params0=None, mag=1e-3):
    """Sets up and returns posititons of walkers

    Parameters
    ----------
    source : str
    version : int
    n_walkers: int, number of mcmc walkers to use
    params0: array, initial guess (mdot1, x, z, qb, g, redshift, d, inc)
    mag: flt, magnitude of random seeds to use for initial mcmc 'ball'
    """
    if type(params0) == type(None):
        mcmc_version = mcmc_versions.McmcVersion(source=source, version=version)
        params0 = mcmc_version.initial_position

    n_dimensions = len(params0)
    pos = [params0 * (1 + mag * np.random.randn(n_dimensions)) for i in range(n_walkers)]
    return np.array(pos)

def run_sampler(sampler, pos, n_steps, verbose=True):
    """Runs emcee chain for n_steps
    """
    t0 = time.time()
    result = None
    # TODO: update to use pool
    for i, result in enumerate(sampler.sample(pos, iterations=n_steps, progress=True)):
        pass

    t1 = time.time()
    dtime = t1 - t0
    time_per_step = dtime / n_steps

    n_walkers = pos.shape[0]
    n_samples = n_walkers * n_steps
    time_per_sample = dtime / n_samples

    if verbose:
        print(f'Compute time: {dtime:.1f} s ({dtime/3600:.2f} hr)')
        print(f'Time per step: {time_per_step:.1f} s')
        print(f'Time per sample: {time_per_sample:.4f} s')
    return result

def get_acceptance_fraction(source=None, version=None, n_walkers=None,
                            n_steps=None, sampler=None):
    """Returns acceptance fraction averaged over all walkers for all steps

    Must provide either:
        1. a sampler object (as returned from load_sampler_state)
        2. source, version, n_walkers, and n_steps
    """
    if sampler is None:
        if None in (source, version, n_walkers, n_steps):
            raise ValueError('Must provide source, version, n_steps, '
                             + 'and n_walkers (or directly provide sampler object)')
        sampler = mcmc_tools.load_sampler_state(source=source, version=version,
                                                n_steps=n_steps, n_walkers=n_walkers)
    else:
        n_steps = sampler['iterations']

    return np.average(sampler['naccepted'] / n_steps)

def optimise(source, version, x0=None):
    """Optimise the starting position of the mcmc walkers using standard
        minimisation methods

    Parameters
    ----------
    x0 : ndarray
        Initial guess of parameters
    """
    bfit = burstfit.BurstFit(source=source, version=version, verbose=False,
                             lhood_factor=-1, zero_lhood=-1e9)
    if x0 is None:
        x0 = bfit.mcmc_version.initial_position

    return fmin(bfit.lhood, x0=x0, maxfun=10000)

def convert_params(params, source, version):
    """Converts params from dict to raw list format, and vice versa (ensures order)
    """
    pkeys = mcmc_versions.get_parameter(source, version, parameter='param_keys')
    check_params_length(params=params, n=len(pkeys))
    ptype = type(params)

    if ptype == dict:
        params_out = []
        for key in pkeys:
            params_out += [params[key]]

    elif (ptype == list) or (ptype == tuple) or (ptype == np.ndarray):
        params_out = {}
        for i, key in enumerate(pkeys):
            params_out[key] = params[i]
    else:
        raise TypeError('type(params) must be dict or array-like')

    return params_out

ext: py

max_stars_repo_name: zacjohnston/pyburst

max_issues_repo_name: zacjohnston/pyburst

max_forks_repo_name: zacjohnston/pyburst

reason: import numpy,from scipy


text:
"""Classes representing batches of normalized or unnormalized data."""
from __future__ import print_function, division, absolute_import

import collections

import numpy as np

from .. import imgaug as ia
from . import normalization as nlib
from . import utils

DEFAULT = "DEFAULT"

_AUGMENTABLE_NAMES = [
    "images", "heatmaps", "segmentation_maps", "keypoints",
    "bounding_boxes", "polygons", "line_strings"]

_AugmentableColumn = collections.namedtuple(
    "_AugmentableColumn",
    ["name", "value", "attr_name"])

def _get_column_names(batch, postfix):
    return [column.name
            for column
            in _get_columns(batch, postfix)]

def _get_columns(batch, postfix):
    result = []
    for name in _AUGMENTABLE_NAMES:
        attr_name = name + postfix
        value = getattr(batch, name + postfix)
        # Every data item is either an array or a list. If there are no
        # items in the array/list, there are also no shapes to change
        # as shape-changes are imagewise. Hence, we can afford to check
        # len() here.
        if value is not None and len(value) > 0:
            result.append(_AugmentableColumn(name, value, attr_name))
    return result

# TODO also support (H,W,C) for heatmaps of len(images) == 1
# TODO also support (H,W) for segmaps of len(images) == 1
class UnnormalizedBatch(object):
    """
    Class for batches of unnormalized data before and after augmentation.

    Parameters
    ----------
    images : None or (N,H,W,C) ndarray or (N,H,W) ndarray or iterable of (H,W,C) ndarray or iterable of (H,W) ndarray
        The images to augment.

    heatmaps : None or (N,H,W,C) ndarray or imgaug.augmentables.heatmaps.HeatmapsOnImage or iterable of (H,W,C) ndarray or iterable of imgaug.augmentables.heatmaps.HeatmapsOnImage
        The heatmaps to augment.
        If anything else than ``HeatmapsOnImage``, then the number of heatmaps
        must match the number of images provided via parameter `images`.
        The number is contained either in ``N`` or the first iterable's size.

    segmentation_maps : None or (N,H,W) ndarray or imgaug.augmentables.segmaps.SegmentationMapsOnImage or iterable of (H,W) ndarray or iterable of imgaug.augmentables.segmaps.SegmentationMapsOnImage
        The segmentation maps to augment.
        If anything else than ``SegmentationMapsOnImage``, then the number of
        segmaps must match the number of images provided via parameter
        `images`. The number is contained either in ``N`` or the first
        iterable's size.

    keypoints : None or list of (N,K,2) ndarray or tuple of number or imgaug.augmentables.kps.Keypoint or iterable of (K,2) ndarray or iterable of tuple of number or iterable of imgaug.augmentables.kps.Keypoint or iterable of imgaug.augmentables.kps.KeypointOnImage or iterable of iterable of tuple of number or iterable of iterable of imgaug.augmentables.kps.Keypoint
        The keypoints to augment.
        If a tuple (or iterable(s) of tuple), then iterpreted as (x,y)
        coordinates and must hence contain two numbers.
        A single tuple represents a single coordinate on one image, an
        iterable of tuples the coordinates on one image and an iterable of
        iterable of tuples the coordinates on several images. Analogous if
        ``Keypoint`` objects are used instead of tuples.
        If an ndarray, then ``N`` denotes the number of images and ``K`` the
        number of keypoints on each image.
        If anything else than ``KeypointsOnImage`` is provided, then the
        number of keypoint groups must match the number of images provided
        via parameter `images`. The number is contained e.g. in ``N`` or
        in case of "iterable of iterable of tuples" in the first iterable's
        size.

    bounding_boxes : None or (N,B,4) ndarray or tuple of number or imgaug.augmentables.bbs.BoundingBox or imgaug.augmentables.bbs.BoundingBoxesOnImage or iterable of (B,4) ndarray or iterable of tuple of number or iterable of imgaug.augmentables.bbs.BoundingBox or iterable of imgaug.augmentables.bbs.BoundingBoxesOnImage or iterable of iterable of tuple of number or iterable of iterable imgaug.augmentables.bbs.BoundingBox
        The bounding boxes to augment.
        This is analogous to the `keypoints` parameter. However, each
        tuple -- and also the last index in case of arrays -- has size 4,
        denoting the bounding box coordinates ``x1``, ``y1``, ``x2`` and ``y2``.

    polygons : None  or (N,#polys,#points,2) ndarray or imgaug.augmentables.polys.Polygon or imgaug.augmentables.polys.PolygonsOnImage or iterable of (#polys,#points,2) ndarray or iterable of tuple of number or iterable of imgaug.augmentables.kps.Keypoint or iterable of imgaug.augmentables.polys.Polygon or iterable of imgaug.augmentables.polys.PolygonsOnImage or iterable of iterable of (#points,2) ndarray or iterable of iterable of tuple of number or iterable of iterable of imgaug.augmentables.kps.Keypoint or iterable of iterable of imgaug.augmentables.polys.Polygon or iterable of iterable of iterable of tuple of number or iterable of iterable of iterable of tuple of imgaug.augmentables.kps.Keypoint
        The polygons to augment.
        This is similar to the `keypoints` parameter. However, each polygon
        may be made up of several ``(x,y)`` coordinates (three or more are
        required for valid polygons).
        The following datatypes will be interpreted as a single polygon on a
        single image:

          * ``imgaug.augmentables.polys.Polygon``
          * ``iterable of tuple of number``
          * ``iterable of imgaug.augmentables.kps.Keypoint``

        The following datatypes will be interpreted as multiple polygons on a
        single image:

          * ``imgaug.augmentables.polys.PolygonsOnImage``
          * ``iterable of imgaug.augmentables.polys.Polygon``
          * ``iterable of iterable of tuple of number``
          * ``iterable of iterable of imgaug.augmentables.kps.Keypoint``
          * ``iterable of iterable of imgaug.augmentables.polys.Polygon``

        The following datatypes will be interpreted as multiple polygons on
        multiple images:

          * ``(N,#polys,#points,2) ndarray``
          * ``iterable of (#polys,#points,2) ndarray``
          * ``iterable of iterable of (#points,2) ndarray``
          * ``iterable of iterable of iterable of tuple of number``
          * ``iterable of iterable of iterable of tuple of imgaug.augmentables.kps.Keypoint``

    line_strings : None or (N,#lines,#points,2) ndarray or imgaug.augmentables.lines.LineString or imgaug.augmentables.lines.LineStringOnImage or iterable of (#lines,#points,2) ndarray or iterable of tuple of number or iterable of imgaug.augmentables.kps.Keypoint or iterable of imgaug.augmentables.lines.LineString or iterable of imgaug.augmentables.lines.LineStringOnImage or iterable of iterable of (#points,2) ndarray or iterable of iterable of tuple of number or iterable of iterable of imgaug.augmentables.kps.Keypoint or iterable of iterable of imgaug.augmentables.polys.LineString or iterable of iterable of iterable of tuple of number or iterable of iterable of iterable of tuple of imgaug.augmentables.kps.Keypoint
        The line strings to augment.
        See `polygons` for more details as polygons follow a similar
        structure to line strings.

    data
        Additional data that is saved in the batch and may be read out
        after augmentation. This could e.g. contain filepaths to each image
        in `images`. As this object is usually used for background
        augmentation with multiple processes, the augmented Batch objects might
        not be returned in the original order, making this information useful.

    """

    def __init__(self, images=None, heatmaps=None, segmentation_maps=None,
                 keypoints=None, bounding_boxes=None, polygons=None,
                 line_strings=None, data=None):
        """Construct a new :class:`UnnormalizedBatch` instance."""
        self.images_unaug = images
        self.images_aug = None
        self.heatmaps_unaug = heatmaps
        self.heatmaps_aug = None
        self.segmentation_maps_unaug = segmentation_maps
        self.segmentation_maps_aug = None
        self.keypoints_unaug = keypoints
        self.keypoints_aug = None
        self.bounding_boxes_unaug = bounding_boxes
        self.bounding_boxes_aug = None
        self.polygons_unaug = polygons
        self.polygons_aug = None
        self.line_strings_unaug = line_strings
        self.line_strings_aug = None
        self.data = data

    def get_column_names(self):
        """Get the names of types of augmentables that contain data.

        This method is intended for situations where one wants to know which
        data is contained in the batch that has to be augmented, visualized
        or something similar.

        Added in 0.4.0.

        Returns
        -------
        list of str
            Names of types of augmentables. E.g. ``["images", "polygons"]``.

        """
        return _get_column_names(self, "_unaug")

    def to_normalized_batch(self):
        """Convert this unnormalized batch to an instance of Batch.

        As this method is intended to be called before augmentation, it
        assumes that none of the ``*_aug`` attributes is yet set.
        It will produce an AssertionError otherwise.

        The newly created Batch's ``*_unaug`` attributes will match the ones
        in this batch, just in normalized form.

        Returns
        -------
        imgaug.augmentables.batches.Batch
            The batch, with ``*_unaug`` attributes being normalized.

        """
        contains_no_augmented_data_yet = all([
            attr is None
            for attr_name, attr
            in self.__dict__.items()
            if attr_name.endswith("_aug")])
        assert contains_no_augmented_data_yet, (
            "Expected UnnormalizedBatch to not contain any augmented data "
            "before normalization, but at least one '*_aug' attribute was "
            "already set.")

        images_unaug = nlib.normalize_images(self.images_unaug)
        shapes = None
        if images_unaug is not None:
            shapes = [image.shape for image in images_unaug]

        return Batch(
            images=images_unaug,
            heatmaps=nlib.normalize_heatmaps(
                self.heatmaps_unaug, shapes),
            segmentation_maps=nlib.normalize_segmentation_maps(
                self.segmentation_maps_unaug, shapes),
            keypoints=nlib.normalize_keypoints(
                self.keypoints_unaug, shapes),
            bounding_boxes=nlib.normalize_bounding_boxes(
                self.bounding_boxes_unaug, shapes),
            polygons=nlib.normalize_polygons(
                self.polygons_unaug, shapes),
            line_strings=nlib.normalize_line_strings(
                self.line_strings_unaug, shapes),
            data=self.data
        )

    def fill_from_augmented_normalized_batch_(self, batch_aug_norm):
        """
        Fill this batch with (normalized) augmentation results in-place.

        This method receives a (normalized) Batch instance, takes all
        ``*_aug`` attributes out if it and assigns them to this
        batch *in unnormalized form*. Hence, the datatypes of all ``*_aug``
        attributes will match the datatypes of the ``*_unaug`` attributes.

        Added in 0.4.0.

        Parameters
        ----------
        batch_aug_norm: imgaug.augmentables.batches.Batch
            Batch after normalization and augmentation.

        Returns
        -------
        imgaug.augmentables.batches.UnnormalizedBatch
            This instance itself.
            All ``*_unaug`` attributes are unchanged.
            All ``*_aug`` attributes are taken from `batch_normalized`,
            converted to unnormalized form.

        """
        self.images_aug = nlib.invert_normalize_images(
            batch_aug_norm.images_aug, self.images_unaug)
        self.heatmaps_aug = nlib.invert_normalize_heatmaps(
            batch_aug_norm.heatmaps_aug, self.heatmaps_unaug)
        self.segmentation_maps_aug = nlib.invert_normalize_segmentation_maps(
            batch_aug_norm.segmentation_maps_aug, self.segmentation_maps_unaug)
        self.keypoints_aug = nlib.invert_normalize_keypoints(
            batch_aug_norm.keypoints_aug, self.keypoints_unaug)
        self.bounding_boxes_aug = nlib.invert_normalize_bounding_boxes(
            batch_aug_norm.bounding_boxes_aug, self.bounding_boxes_unaug)
        self.polygons_aug = nlib.invert_normalize_polygons(
            batch_aug_norm.polygons_aug, self.polygons_unaug)
        self.line_strings_aug = nlib.invert_normalize_line_strings(
            batch_aug_norm.line_strings_aug, self.line_strings_unaug)
        return self

    def fill_from_augmented_normalized_batch(self, batch_aug_norm):
        """
        Fill this batch with (normalized) augmentation results.

        This method receives a (normalized) Batch instance, takes all
        ``*_aug`` attributes out if it and assigns them to this
        batch *in unnormalized form*. Hence, the datatypes of all ``*_aug``
        attributes will match the datatypes of the ``*_unaug`` attributes.

        Parameters
        ----------
        batch_aug_norm: imgaug.augmentables.batches.Batch
            Batch after normalization and augmentation.

        Returns
        -------
        imgaug.augmentables.batches.UnnormalizedBatch
            New UnnormalizedBatch instance. All ``*_unaug`` attributes are
            taken from the old UnnormalizedBatch (without deepcopying them)
            and all ``*_aug`` attributes are taken from `batch_normalized`,
            converted to unnormalized form.

        """
        # we take here the .data from the normalized batch instead of from
        # self for the rare case where one has decided to somehow change it
        # during augmentation
        batch = UnnormalizedBatch(
            images=self.images_unaug,
            heatmaps=self.heatmaps_unaug,
            segmentation_maps=self.segmentation_maps_unaug,
            keypoints=self.keypoints_unaug,
            bounding_boxes=self.bounding_boxes_unaug,
            polygons=self.polygons_unaug,
            line_strings=self.line_strings_unaug,
            data=batch_aug_norm.data
        )

        batch.images_aug = nlib.invert_normalize_images(
            batch_aug_norm.images_aug, self.images_unaug)
        batch.heatmaps_aug = nlib.invert_normalize_heatmaps(
            batch_aug_norm.heatmaps_aug, self.heatmaps_unaug)
        batch.segmentation_maps_aug = nlib.invert_normalize_segmentation_maps(
            batch_aug_norm.segmentation_maps_aug, self.segmentation_maps_unaug)
        batch.keypoints_aug = nlib.invert_normalize_keypoints(
            batch_aug_norm.keypoints_aug, self.keypoints_unaug)
        batch.bounding_boxes_aug = nlib.invert_normalize_bounding_boxes(
            batch_aug_norm.bounding_boxes_aug, self.bounding_boxes_unaug)
        batch.polygons_aug = nlib.invert_normalize_polygons(
            batch_aug_norm.polygons_aug, self.polygons_unaug)
        batch.line_strings_aug = nlib.invert_normalize_line_strings(
            batch_aug_norm.line_strings_aug, self.line_strings_unaug)

        return batch

class Batch(object):
    """
    Class encapsulating a batch before and after augmentation.

    Parameters
    ----------
    images : None or (N,H,W,C) ndarray or list of (H,W,C) ndarray
        The images to augment.

    heatmaps : None or list of imgaug.augmentables.heatmaps.HeatmapsOnImage
        The heatmaps to augment.

    segmentation_maps : None or list of imgaug.augmentables.segmaps.SegmentationMapsOnImage
        The segmentation maps to augment.

    keypoints : None or list of imgaug.augmentables.kps.KeypointOnImage
        The keypoints to augment.

    bounding_boxes : None or list of imgaug.augmentables.bbs.BoundingBoxesOnImage
        The bounding boxes to augment.

    polygons : None or list of imgaug.augmentables.polys.PolygonsOnImage
        The polygons to augment.

    line_strings : None or list of imgaug.augmentables.lines.LineStringsOnImage
        The line strings to augment.

    data
        Additional data that is saved in the batch and may be read out
        after augmentation. This could e.g. contain filepaths to each image
        in `images`. As this object is usually used for background
        augmentation with multiple processes, the augmented Batch objects might
        not be returned in the original order, making this information useful.

    """

    def __init__(self, images=None, heatmaps=None, segmentation_maps=None,
                 keypoints=None, bounding_boxes=None, polygons=None,
                 line_strings=None, data=None):
        """Construct a new :class:`Batch` instance."""
        self.images_unaug = images
        self.images_aug = None
        self.heatmaps_unaug = heatmaps
        self.heatmaps_aug = None
        self.segmentation_maps_unaug = segmentation_maps
        self.segmentation_maps_aug = None
        self.keypoints_unaug = keypoints
        self.keypoints_aug = None
        self.bounding_boxes_unaug = bounding_boxes
        self.bounding_boxes_aug = None
        self.polygons_unaug = polygons
        self.polygons_aug = None
        self.line_strings_unaug = line_strings
        self.line_strings_aug = None
        self.data = data

    @property
    @ia.deprecated("Batch.images_unaug")
    def images(self):
        """Get unaugmented images."""
        return self.images_unaug

    @property
    @ia.deprecated("Batch.heatmaps_unaug")
    def heatmaps(self):
        """Get unaugmented heatmaps."""
        return self.heatmaps_unaug

    @property
    @ia.deprecated("Batch.segmentation_maps_unaug")
    def segmentation_maps(self):
        """Get unaugmented segmentation maps."""
        return self.segmentation_maps_unaug

    @property
    @ia.deprecated("Batch.keypoints_unaug")
    def keypoints(self):
        """Get unaugmented keypoints."""
        return self.keypoints_unaug

    @property
    @ia.deprecated("Batch.bounding_boxes_unaug")
    def bounding_boxes(self):
        """Get unaugmented bounding boxes."""
        return self.bounding_boxes_unaug

    def get_column_names(self):
        """Get the names of types of augmentables that contain data.

        This method is intended for situations where one wants to know which
        data is contained in the batch that has to be augmented, visualized
        or something similar.

        Added in 0.4.0.

        Returns
        -------
        list of str
            Names of types of augmentables. E.g. ``["images", "polygons"]``.

        """
        return _get_column_names(self, "_unaug")

    def to_normalized_batch(self):
        """Return this batch.

        This method does nothing and only exists to simplify interfaces
        that accept both :class:`UnnormalizedBatch` and :class:`Batch`.

        Added in 0.4.0.

        Returns
        -------
        imgaug.augmentables.batches.Batch
            This batch (not copied).

        """
        return self

    def to_batch_in_augmentation(self):
        """Convert this batch to a :class:`_BatchInAugmentation` instance.

        Added in 0.4.0.

        Returns
        -------
        imgaug.augmentables.batches._BatchInAugmentation
            The converted batch.

        """
        def _copy(var):
            # TODO first check here if _aug is set and if it is then use that?
            if var is not None:
                return utils.copy_augmentables(var)
            return var

        return _BatchInAugmentation(
            images=_copy(self.images_unaug),
            heatmaps=_copy(self.heatmaps_unaug),
            segmentation_maps=_copy(self.segmentation_maps_unaug),
            keypoints=_copy(self.keypoints_unaug),
            bounding_boxes=_copy(self.bounding_boxes_unaug),
            polygons=_copy(self.polygons_unaug),
            line_strings=_copy(self.line_strings_unaug)
        )

    def fill_from_batch_in_augmentation_(self, batch_in_augmentation):
        """Set the columns in this batch to the column values of another batch.

        This method works in-place.

        Added in 0.4.0.

        Parameters
        ----------
        batch_in_augmentation : _BatchInAugmentation
            Batch of which to use the column values.
            The values are *not* copied. Only their references are used.

        Returns
        -------
        Batch
            The updated batch. (Modified in-place.)

        """
        self.images_aug = batch_in_augmentation.images
        self.heatmaps_aug = batch_in_augmentation.heatmaps
        self.segmentation_maps_aug = batch_in_augmentation.segmentation_maps
        self.keypoints_aug = batch_in_augmentation.keypoints
        self.bounding_boxes_aug = batch_in_augmentation.bounding_boxes
        self.polygons_aug = batch_in_augmentation.polygons
        self.line_strings_aug = batch_in_augmentation.line_strings
        return self

    def deepcopy(self,
                 images_unaug=DEFAULT,
                 images_aug=DEFAULT,
                 heatmaps_unaug=DEFAULT,
                 heatmaps_aug=DEFAULT,
                 segmentation_maps_unaug=DEFAULT,
                 segmentation_maps_aug=DEFAULT,
                 keypoints_unaug=DEFAULT,
                 keypoints_aug=DEFAULT,
                 bounding_boxes_unaug=DEFAULT,
                 bounding_boxes_aug=DEFAULT,
                 polygons_unaug=DEFAULT,
                 polygons_aug=DEFAULT,
                 line_strings_unaug=DEFAULT,
                 line_strings_aug=DEFAULT):
        """Copy this batch and all of its column values.

        Parameters
        ----------
        images_unaug : imgaug.augmentables.batches.DEFAULT or None or (N,H,W,C) ndarray or list of (H,W,C) ndarray
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        images_aug : imgaug.augmentables.batches.DEFAULT or None or (N,H,W,C) ndarray or list of (H,W,C) ndarray
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        heatmaps_unaug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.heatmaps.HeatmapsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        heatmaps_aug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.heatmaps.HeatmapsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        segmentation_maps_unaug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.segmaps.SegmentationMapsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        segmentation_maps_aug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.segmaps.SegmentationMapsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        keypoints_unaug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.kps.KeypointOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        keypoints_aug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.kps.KeypointOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        bounding_boxes_unaug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.bbs.BoundingBoxesOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        bounding_boxes_aug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.bbs.BoundingBoxesOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        polygons_unaug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.polys.PolygonsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        polygons_aug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.polys.PolygonsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        line_strings_unaug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.lines.LineStringsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        line_strings_aug : imgaug.augmentables.batches.DEFAULT or None or list of imgaug.augmentables.lines.LineStringsOnImage
            Copies the current attribute value without changes if set to
            ``imgaug.augmentables.batches.DEFAULT``.
            Otherwise same as in :func:`Batch.__init__`.

        Returns
        -------
        Batch
            Deep copy of the batch, optionally with new attributes.

        """
        def _copy_optional(self_attr, arg):
            return utils.deepcopy_fast(arg if arg is not DEFAULT else self_attr)

        batch = Batch(
            images=_copy_optional(self.images_unaug, images_unaug),
            heatmaps=_copy_optional(self.heatmaps_unaug, heatmaps_unaug),
            segmentation_maps=_copy_optional(self.segmentation_maps_unaug,
                                             segmentation_maps_unaug),
            keypoints=_copy_optional(self.keypoints_unaug, keypoints_unaug),
            bounding_boxes=_copy_optional(self.bounding_boxes_unaug,
                                          bounding_boxes_unaug),
            polygons=_copy_optional(self.polygons_unaug, polygons_unaug),
            line_strings=_copy_optional(self.line_strings_unaug,
                                        line_strings_unaug),
            data=utils.deepcopy_fast(self.data)
        )
        batch.images_aug = _copy_optional(self.images_aug, images_aug)
        batch.heatmaps_aug = _copy_optional(self.heatmaps_aug, heatmaps_aug)
        batch.segmentation_maps_aug = _copy_optional(self.segmentation_maps_aug,
                                                     segmentation_maps_aug)
        batch.keypoints_aug = _copy_optional(self.keypoints_aug, keypoints_aug)
        batch.bounding_boxes_aug = _copy_optional(self.bounding_boxes_aug,
                                                  bounding_boxes_aug)
        batch.polygons_aug = _copy_optional(self.polygons_aug, polygons_aug)
        batch.line_strings_aug = _copy_optional(self.line_strings_aug,
                                                line_strings_aug)

        return batch

# Added in 0.4.0.
class _BatchInAugmentationPropagationContext(object):
    def __init__(self, batch, augmenter, hooks, parents):
        self.batch = batch
        self.augmenter = augmenter
        self.hooks = hooks
        self.parents = parents
        self.noned_info = None

    def __enter__(self):
        if self.hooks is not None:
            self.noned_info = self.batch.apply_propagation_hooks_(
                self.augmenter, self.hooks, self.parents)
        return self.batch

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.noned_info is not None:
            self.batch = \
                self.batch.invert_apply_propagation_hooks_(self.noned_info)

class _BatchInAugmentation(object):
    """
    Class encapsulating a batch during the augmentation process.

    Data within the batch is already verified and normalized, similar to
    :class:`Batch`. Data within the batch may be changed in-place. No initial
    copy is needed.

    Added in 0.4.0.

    Parameters
    ----------
    images : None or (N,H,W,C) ndarray or list of (H,W,C) ndarray
        The images to augment.

    heatmaps : None or list of imgaug.augmentables.heatmaps.HeatmapsOnImage
        The heatmaps to augment.

    segmentation_maps : None or list of imgaug.augmentables.segmaps.SegmentationMapsOnImage
        The segmentation maps to augment.

    keypoints : None or list of imgaug.augmentables.kps.KeypointOnImage
        The keypoints to augment.

    bounding_boxes : None or list of imgaug.augmentables.bbs.BoundingBoxesOnImage
        The bounding boxes to augment.

    polygons : None or list of imgaug.augmentables.polys.PolygonsOnImage
        The polygons to augment.

    line_strings : None or list of imgaug.augmentables.lines.LineStringsOnImage
        The line strings to augment.

    """

    # Added in 0.4.0.
    def __init__(self, images=None, heatmaps=None, segmentation_maps=None,
                 keypoints=None, bounding_boxes=None, polygons=None,
                 line_strings=None, data=None):
        """Create a new :class:`_BatchInAugmentation` instance."""
        self.images = images
        self.heatmaps = heatmaps
        self.segmentation_maps = segmentation_maps
        self.keypoints = keypoints
        self.bounding_boxes = bounding_boxes
        self.polygons = polygons
        self.line_strings = line_strings
        self.data = data

    @property
    def empty(self):
        """Estimate whether this batch is empty, i.e. contains no data.

        Added in 0.4.0.

        Returns
        -------
        bool
            ``True`` if the batch contains no data to augment.
            ``False`` otherwise.

        """
        return self.nb_rows == 0

    @property
    def nb_rows(self):
        """Get the number of rows (i.e. examples) in this batch.

        Note that this method assumes that all columns have the same number
        of rows.

        Added in 0.4.0.

        Returns
        -------
        int
            Number of rows or ``0`` if there is no data in the batch.

        """
        for augm_name in _AUGMENTABLE_NAMES:
            value = getattr(self, augm_name)
            if value is not None:
                return len(value)
        return 0

    @property
    def columns(self):
        """Get the columns of data to augment.

        Each column represents one datatype and its corresponding data,
        e.g. images or polygons.

        Added in 0.4.0.

        Returns
        -------
        list of _AugmentableColumn
            The columns to augment within this batch.

        """
        return _get_columns(self, "")

    def get_column_names(self):
        """Get the names of types of augmentables that contain data.

        This method is intended for situations where one wants to know which
        data is contained in the batch that has to be augmented, visualized
        or something similar.

        Added in 0.4.0.

        Returns
        -------
        list of str
            Names of types of augmentables. E.g. ``["images", "polygons"]``.

        """
        return _get_column_names(self, "")

    def get_rowwise_shapes(self):
        """Get the shape of each row within this batch.

        Each row denotes the data of different types (e.g. image array,
        polygons) corresponding to a single example in the batch.

        This method assumes that all ``.shape`` attributes contain the same
        shape and that it is identical to the image's shape.
        It also assumes that there are no columns containing only ``None`` s.

        Added in 0.4.0.

        Returns
        -------
        list of tuple of int
            The shapes of each row.

        """
        nb_rows = self.nb_rows
        columns = self.columns
        shapes = [None] * nb_rows
        found = np.zeros((nb_rows,), dtype=bool)
        for column in columns:
            if column.name == "images" and ia.is_np_array(column.value):
                shapes = [column.value.shape[1:]] * nb_rows
            else:
                for i, item in enumerate(column.value):
                    if item is not None:
                        shapes[i] = item.shape
                        found[i] = True
            if np.all(found):
                return shapes
        return shapes

    def subselect_rows_by_indices(self, indices):
        """Reduce this batch to a subset of rows based on their row indices.

        Added in 0.4.0.

        Parameters
        ----------
        indices : iterable of int
            Row indices to select.

        Returns
        -------
        _BatchInAugmentation
            Batch containing only a subselection of rows.

        """
        kwargs = {"data": self.data}
        for augm_name in _AUGMENTABLE_NAMES:
            rows = getattr(self, augm_name)
            if rows is not None:
                if augm_name == "images" and ia.is_np_array(rows):
                    rows = rows[indices]  # pylint: disable=unsubscriptable-object
                else:
                    rows = [rows[index] for index in indices]

                if len(rows) == 0:
                    rows = None
            kwargs[augm_name] = rows

        return _BatchInAugmentation(**kwargs)

    def invert_subselect_rows_by_indices_(self, indices, batch_subselected):
        """Reverse the subselection of rows in-place.

        This is the inverse of
        :func:`_BatchInAugmentation.subselect_rows_by_indices`.

        This method has to be executed on the batch *before* subselection.

        Added in 0.4.0.

        Parameters
        ----------
        indices : iterable of int
            Row indices that were selected. (This is the input to

        batch_subselected : _BatchInAugmentation
            The batch after
            :func:`_BatchInAugmentation.subselect_rows_by_indices` was called.

        Returns
        -------
        _BatchInAugmentation
            The updated batch. (Modified in-place.)

        Examples
        --------
        >>> import numpy as np
        >>> from imgaug.augmentables.batches import _BatchInAugmentation
        >>> images = np.zeros((2, 10, 20, 3), dtype=np.uint8)
        >>> batch = _BatchInAugmentation(images=images)
        >>> batch_sub = batch.subselect_rows_by_indices([0])
        >>> batch_sub.images += 1
        >>> batch = batch.invert_subselect_rows_by_indices_([0], batch_sub)

        """
        for augm_name in _AUGMENTABLE_NAMES:
            column = getattr(self, augm_name)
            if column is not None:
                column_sub = getattr(batch_subselected, augm_name)
                if column_sub is None:
                    # list of indices was empty, resulting in the columns
                    # in the subselected batch being empty and replaced
                    # by Nones. We can just re-use the columns before
                    # subselection.
                    pass
                elif augm_name == "images" and ia.is_np_array(column):
                    # An array does not have to stay an array after
                    # augmentation. The shapes and/or dtypes of rows may
                    # change, turning the array into a list.
                    if ia.is_np_array(column_sub):
                        shapes = {column.shape[1:], column_sub.shape[1:]}
                        dtypes = {column.dtype.name, column_sub.dtype.name}
                    else:
                        shapes = set(
                            [column.shape[1:]]
                            + [image.shape for image in column_sub])
                        dtypes = set(
                            [column.dtype.name]
                            + [image.dtype.name for image in column_sub])

                    if len(shapes) == 1 and len(dtypes) == 1:
                        column[indices] = column_sub  # pylint: disable=unsupported-assignment-operation
                    else:
                        self.images = list(column)
                        for ith_index, index in enumerate(indices):
                            self.images[index] = column_sub[ith_index]
                else:
                    for ith_index, index in enumerate(indices):
                        column[index] = column_sub[ith_index]  # pylint: disable=unsupported-assignment-operation

        return self

    def propagation_hooks_ctx(self, augmenter, hooks, parents):
        """Start a context in which propagation hooks are applied.

        Added in 0.4.0.

        Parameters
        ----------
        augmenter : imgaug.augmenters.meta.Augmenter
            Augmenter to provide to the propagation hook function.

        hooks : imgaug.imgaug.HooksImages or imgaug.imgaug.HooksKeypoints
            The hooks that might contain a propagation hook function.

        parents : list of imgaug.augmenters.meta.Augmenter
            The list of parents to provide to the propagation hook function.

        Returns
        -------
        _BatchInAugmentationPropagationContext
            The progagation hook context.

        """
        return _BatchInAugmentationPropagationContext(
            self, augmenter=augmenter, hooks=hooks, parents=parents)

    def apply_propagation_hooks_(self, augmenter, hooks, parents):
        """Set columns in this batch to ``None`` based on a propagation hook.

        This method works in-place.

        Added in 0.4.0.

        Parameters
        ----------
        augmenter : imgaug.augmenters.meta.Augmenter
            Augmenter to provide to the propagation hook function.

        hooks : imgaug.imgaug.HooksImages or imgaug.imgaug.HooksKeypoints
            The hooks that might contain a propagation hook function.

        parents : list of imgaug.augmenters.meta.Augmenter
            The list of parents to provide to the propagation hook function.

        Returns
        -------
        list of tuple of str
            Information about which columns were set to ``None``.
            Each tuple contains
            ``(column attribute name, column value before setting it to None)``.
            This information is required when calling
            :func:`_BatchInAugmentation.invert_apply_propagation_hooks_`.

        """
        if hooks is None:
            return None

        noned_info = []
        for column in self.columns:
            is_prop = hooks.is_propagating(
                column.value, augmenter=augmenter, parents=parents,
                default=True)
            if not is_prop:
                setattr(self, column.attr_name, None)
                noned_info.append((column.attr_name, column.value))
        return noned_info

    def invert_apply_propagation_hooks_(self, noned_info):
        """Set columns from ``None`` back to their original values.

        This is the inverse of
        :func:`_BatchInAugmentation.apply_propagation_hooks_`.

        This method works in-place.

        Added in 0.4.0.

        Parameters
        ----------
        noned_info : list of tuple of str
            Information about which columns were set to ``None`` and their
            original values. This is the output of
            :func:`_BatchInAugmentation.apply_propagation_hooks_`.

        Returns
        -------
        _BatchInAugmentation
            The updated batch. (Modified in-place.)

        """
        for attr_name, value in noned_info:
            setattr(self, attr_name, value)
        return self

    def to_batch_in_augmentation(self):
        """Convert this batch to a :class:`_BatchInAugmentation` instance.

        This method simply returns the batch itself. It exists for consistency
        with the other batch classes.

        Added in 0.4.0.

        Returns
        -------
        imgaug.augmentables.batches._BatchInAugmentation
            The batch itself. (Not copied.)

        """
        return self

    def fill_from_batch_in_augmentation_(self, batch_in_augmentation):
        """Set the columns in this batch to the column values of another batch.

        This method works in-place.

        Added in 0.4.0.

        Parameters
        ----------
        batch_in_augmentation : _BatchInAugmentation
            Batch of which to use the column values.
            The values are *not* copied. Only their references are used.

        Returns
        -------
        _BatchInAugmentation
            The updated batch. (Modified in-place.)

        """
        if batch_in_augmentation is self:
            return self

        self.images = batch_in_augmentation.images
        self.heatmaps = batch_in_augmentation.heatmaps
        self.segmentation_maps = batch_in_augmentation.segmentation_maps
        self.keypoints = batch_in_augmentation.keypoints
        self.bounding_boxes = batch_in_augmentation.bounding_boxes
        self.polygons = batch_in_augmentation.polygons
        self.line_strings = batch_in_augmentation.line_strings

        return self

    def to_batch(self, batch_before_aug):
        """Convert this batch into a :class:`Batch` instance.

        Added in 0.4.0.

        Parameters
        ----------
        batch_before_aug : imgaug.augmentables.batches.Batch
            The batch before augmentation. It is required to set the input
            data of the :class:`Batch` instance, e.g. ``images_unaug``
            or ``data``.

        Returns
        -------
        imgaug.augmentables.batches.Batch
            Batch, with original unaugmented inputs from `batch_before_aug`
            and augmented outputs from this :class:`_BatchInAugmentation`
            instance.

        """
        batch = Batch(
            images=batch_before_aug.images_unaug,
            heatmaps=batch_before_aug.heatmaps_unaug,
            segmentation_maps=batch_before_aug.segmentation_maps_unaug,
            keypoints=batch_before_aug.keypoints_unaug,
            bounding_boxes=batch_before_aug.bounding_boxes_unaug,
            polygons=batch_before_aug.polygons_unaug,
            line_strings=batch_before_aug.line_strings_unaug,
            data=batch_before_aug.data
        )
        batch.images_aug = self.images
        batch.heatmaps_aug = self.heatmaps
        batch.segmentation_maps_aug = self.segmentation_maps
        batch.keypoints_aug = self.keypoints
        batch.bounding_boxes_aug = self.bounding_boxes
        batch.polygons_aug = self.polygons
        batch.line_strings_aug = self.line_strings
        return batch

    def deepcopy(self):
        """Copy this batch and all of its column values.

        Added in 0.4.0.

        Returns
        -------
        _BatchInAugmentation
            Deep copy of this batch.

        """
        batch = _BatchInAugmentation(data=utils.deepcopy_fast(self.data))

        for augm_name in _AUGMENTABLE_NAMES:
            value = getattr(self, augm_name)
            if value is not None:
                setattr(batch, augm_name, utils.copy_augmentables(value))

        return batch

max_stars_repo_name: chroneus/imgaug

max_issues_repo_name: chroneus/imgaug

max_forks_repo_name: chroneus/imgaug

reason: import numpy


text:
#!/usr/bin/env python
import rospy
import roslib
import tf
import sys
import os
import random
import math
import nav_msgs
import geometry_msgs
from nav_msgs.msg import OccupancyGrid
import numpy as np
from mavros_msgs.msg import PositionTarget
from mavros_msgs.msg import Trajectory
from px4_rrt_avoidance.msg import Path_coll
from px4_rrt_avoidance.msg import Maximizer
from px4_rrt_avoidance.msg import RRT_Star_Call
import mavros_msgs
import time
from debugging_functions import show_checked_path
from debugging_functions import pub_node_list

class Drone():
    def __init__(self, drone_pose, desired_altitude):
        self.drone_pose = drone_pose
        self.desired_altitude = desired_altitude

    def set_drone_position(self, msg):
        position = msg.pose.position
        orientation =  msg.pose.orientation
        x = position.x
        y = position.y
        z = position.z
        pitch = orientation.x
        yaw = orientation.y
        roll = orientation.z
        self.drone_pose = [x, y, z, roll, pitch, yaw]
        rrt_star.start = rrt_star.Node(x , y)

        x_ind , y_ind = get_voxel_inds(x, y)
        try:
            local_map_value = local_map.map_matrix[y_ind][x_ind]
            if local_map_value >= 30:
                new_x, new_y = local_map.search_closest_save_point([x_ind, y_ind])
                rrt_star.start = rrt_star.Node(new_x, new_y)
        except IndexError:
            rrt_star.start = rrt_star.Node(x , y)

class Local_Map():
    def __init__(self, map_origin, resolution, map_matrix, target_position, maximizer, start_planning = False ):
        self.resolution = resolution            # width, height, pixel_size
        self.map_origin = map_origin            # [x,y,z, x,y,z,w] x, y, z and orientation: x, y, z, w
        self.map_matrix = map_matrix            # [[-1, -1, 0, 100], [-1, -1, -1, 0], [....]...] -1: unknown; 0: not occupied; 100: occupied
        self.target_position = target_position          # [x, y, z]
        self.maximizer = maximizer
        self.start_planning = start_planning    # starts planning when requested by Masternode

    def local_map_splitter(self, map):

        info = map.info
        width = info.width
        height = info.height
        position = info.origin.position
        origin = [position.x, position.y]
        pixel_size = round(map.info.resolution,3)
        map_vector = map.data

        map_matrix = np.zeros((height, width))

        i = 0 #height
        k = 0 # position in Map_matrix
        while (i <= height -1):
            j = 0
            while (j <= width -1):
                map_matrix[i][j] = map_vector[k]
                k = k +1
                j = j +1
            i = i +1
        self.map_matrix = map_matrix
        self.resolution = {"width": width, "height": height, "pixel_size": pixel_size}
        self.map_origin = origin

        rrt_star.min_rand_height = (self.map_origin[1] / local_map.resolution.get('pixel_size')) +1                         # reduce search space to nod fly at edge
        rrt_star.max_rand_height = ((self.map_origin[1] + height*pixel_size) / local_map.resolution.get('pixel_size')) -1
        rrt_star.min_rand_width = ((self.map_origin[0]) / local_map.resolution.get('pixel_size')) +1
        rrt_star.max_rand_width = ((self.map_origin[0] + width*pixel_size)  / local_map.resolution.get('pixel_size') ) -1

    def set_target(self, target_position):

        rrt_star.start_planning = target_position.start_planning
        if target_position.start_planning == False:
            local_map.maximizer = 0
            pub_map_maximizer(local_map.maximizer)

        x_target = target_position.target_position[0]
        y_target = target_position.target_position[1]
        z_target = target_position.target_position[2]

        print("target_pos: x: " + str(x_target) + " y: " + str(y_target))

        x_ind, y_ind = get_voxel_inds(x_target, y_target)
        target_voxel = [x_ind, y_ind] #x_ind and y_nd at begining = nan
        #CAUTION
        print("Index: " + str(target_voxel))
        try:
            local_map_value = local_map.map_matrix[target_voxel[1]][target_voxel[0]]
        except IndexError:
            print("Index was out of bounds!!!")
            local_map_value = 0     #we assume that unknwon space is free

        if ((local_map_value >= 30) and drone_1.drone_pose[2] > 1):
            #print("Waypoint too close to obstacle")
            print("waypoint inside occupied cell")
            x_target, y_target = self.search_closest_save_point(target_voxel)

        if math.isnan(z_target) == True:
            z_target = 0
        drone_1.desired_altitude = z_target

        if( abs(x_target -drone_1.drone_pose[0]) >= 400  or abs(y_target - drone_1.drone_pose[1]) >= 400):
            #print("Next waypoint is too far away")
            x_target = round(drone_1.drone_pose[0],1)
            y_target = round(drone_1.drone_pose[1], 1)

        local_map.target_position = [x_target, y_target, z_target]
        rrt_star.end = rrt_star.Node(x_target , y_target)
        rrt_star.goal_node.x = x_target
        rrt_star.goal_node.y = y_target
       # test_print()

    def search_closest_save_point(self, t_voxel):
        search_radius = 7 #PARAMETER number of voxels to check in radius/square
                          #should not be higher than minimum of local_map radius
        width_start = t_voxel[0] - search_radius
        height_start = t_voxel[1] - search_radius
        available_voxels = []

        i = height_start
        while (i <= height_start + (2*search_radius) -1 and i <= local_map.resolution.get('height')-1 ):
            j = width_start
            while (j <= width_start + (2*search_radius) -1 and j <= local_map.resolution.get('width')-1 ):

                if local_map.map_matrix[i][j] == 0 or local_map.map_matrix[i][j] == -1 : #there must be free space when an obstacle was detected
                    available_voxels.append([j,i]) #[width, height]
                j = j +1
            j = width_start
            i = i + 1

        #search for closest voxel:
        distance_list = []
        z = len(available_voxels) - 1
        while z >=0:
            voxel = available_voxels[z]
            d_x = distance(t_voxel[0], voxel[0])
            d_y = distance(t_voxel[1], voxel[1])
            true_distance = math.sqrt(d_x**2 + d_y**2)
            distance_list.append(true_distance)
            z = z-1

        min_dis = min(distance_list)
        min_index = distance_list.index(min_dis)

        save_voxel = available_voxels[min_index]

        save_geo = self.get_geo_pos(save_voxel)
        return save_geo[0], save_geo[1]

    @staticmethod
    def get_geo_pos(voxel):
        x_geo = local_map.map_origin[0] + voxel[0]*local_map.resolution.get('pixel_size')
        y_geo = local_map.map_origin[1] + voxel[1]*local_map.resolution.get('pixel_size')

        save_geo = [x_geo, y_geo]
        return save_geo

############################ RRT STAR ORIGINE NOTICE #################################
#   Copyright (c) 2016 - 2021 Atsushi Sakai

#   Permission is hereby granted, free of charge, to any person obtaining a copy
#   of this software and associated documentation files (the "Software"), to deal
#   in the Software without restriction, including without limitation the rights
#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#   copies of the Software, and to permit persons to whom the Software is
#   furnished to do so, subject to the following conditions:

#   The above copyright notice and this permission notice shall be included in all
#   copies or substantial portions of the Software.

#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#   SOFTWARE.
#######################################################################################
# For more details vist: https://github.com/AtsushiSakai/PythonRobotics               #
#######################################################################################

class RRT_Star(object):

    class Node():
        def __init__(self, x, y):
            super(RRT_Star).__init__(RRT_Star)  #ERROR WEIL PYTHON2.7 not Python3 -> switch to ROS2 (not recommended for now)
            self.cost = 0.0
            self.x = x
            self.y = y
            self.path_x = []
            self.path_y = []
            self.parent = None

    def __init__(self, start, goal, rand_area_width,rand_area_height,
                expand_dis = 6,
                path_resolution = 1.5,
                goal_sample_rate = 20,          #[%] will check goal node
                max_iter = 500,                 #CHANGE VALUE IN FUTURE
                connect_circle_dist = 7.0):

        super(RRT_Star, self).__init__( )

        self.connect_circle_dist = connect_circle_dist
        self.goal_node = self.Node(goal[0], goal[1])

        #from RRT planner:
        self.start = self.Node(start[0], start[1])
        self.end = self.Node(goal[0], goal[1])
        self.min_rand_width = rand_area_width[0] / local_map.resolution.get('pixel_size')       #later multiplied by pixel size
        self.max_rand_width = rand_area_width[1] / local_map.resolution.get('pixel_size')       #later multiplied by pixel size
        self.min_rand_height = rand_area_height[0] / local_map.resolution.get('pixel_size')     #later multiplied by pixel size
        self.max_rand_height = rand_area_height[1] / local_map.resolution.get('pixel_size')     #later multiplied by pixel size
        self.expand_dis = expand_dis
        self.path_resolution = path_resolution
        self.goal_sample_rate = goal_sample_rate        #percent: checks if goal can be reached
        self.max_iter = max_iter
        self.node_list = []
        self.start_planning = False    # starts planning when requested by Masternode

    def calc_dist_to_goal(self, x, y):
        dx = x - self.end.x
        dy = y - self.end.y
        return math.hypot(dx, dy)

    def planning(self, search_until_max_iter = False):

        self.node_list = [self.start]
        for i in range(self.max_iter):

            rnd = self.get_random_node()
            nearest_ind = self.get_nearest_node_index(self.node_list, rnd)
            new_node = self.steer(self.node_list[nearest_ind], rnd, self.expand_dis)
            near_node = self.node_list[nearest_ind]
            new_node.cost = near_node.cost + \
                math.hypot(new_node.x-near_node.x, new_node.y-near_node.y)

            if self.check_collision(new_node):
                near_inds = self.find_near_nodes(new_node)

                node_with_updated_parent = self.choose_parent(new_node, near_inds)
                if node_with_updated_parent:
                    self.rewire(node_with_updated_parent, near_inds)
                    self.node_list.append(node_with_updated_parent)
                else:
                    self.node_list.append(new_node)

                if( not search_until_max_iter and new_node):

                    last_index = self.search_best_goal_node()
                    if last_index is not None:

                        return self.generate_final_course(last_index)   #line 885 in RRT Star

        last_index = self.search_best_goal_node()

        if last_index is not None:

            return self.generate_final_course(last_index)

        return None

    def generate_final_course(self, goal_ind):

        path = [[self.end.x, self.end.y]]
        node = self.node_list[goal_ind]
        while node.parent is not None:
            path.append([node.x, node.y])
            node = node.parent

        path.append([node.x, node.y])

        return path

    def search_best_goal_node(self):
        dist_to_goal_list = [self.calc_dist_to_goal(n.x, n.y) for n in self.node_list]
        goal_inds = [dist_to_goal_list.index(i) for i in dist_to_goal_list if i <= self.expand_dis]

        safe_goal_inds=[]

        for goal_ind in goal_inds:

            t_node = self.steer(self.node_list[goal_ind], self.goal_node)

            if self.check_collision(t_node):
                safe_goal_inds.append(goal_ind)

            if not safe_goal_inds:
                return None

            min_cost = min([self.node_list[i].cost for i in safe_goal_inds])
            for i in safe_goal_inds:
                if self.node_list[i].cost == min_cost:
                    return i

            return None

    def get_random_node(self):
        if random.randint(0, 100) > self.goal_sample_rate:
            rnd_1 = random.randint(self.min_rand_width, self.max_rand_width) * round(local_map.resolution.get('pixel_size'),2)
            rnd_2 = random.randint(self.min_rand_height, self.max_rand_height) * round(local_map.resolution.get('pixel_size'),2)
            rnd = self.Node(rnd_1, rnd_2)

        else: #goal point sampling
            rnd = self.Node(self.end.x, self.end.y)

        return rnd

    def rewire(self, new_node, near_inds):
        for i  in near_inds:
            near_node = self.node_list[i]
            edge_node = self.steer(new_node, near_node)
            if not edge_node:
                continue

            edge_node.cost = self.calc_new_cost(new_node, near_node)

            no_collision = self.check_collision(edge_node)
            improve_cost = near_node.cost > edge_node.cost

            if no_collision and improve_cost:
                near_node.x = edge_node.x
                near_node.y = edge_node.y
                near_node.cost = edge_node.cost
                near_node.path_x = edge_node.path_x
                near_node.path_y = edge_node.path_y
                near_node.parent = edge_node.parent
                self.propagate_cost_to_leaves(new_node)

    def propagate_cost_to_leaves(self, parent_node):
        for node in self.node_list:
            if node.parent == parent_node:
                node.cost = self.calc_new_cost(parent_node, node)
                self.propagate_cost_to_leaves(node)

    def find_near_nodes(self, new_node):
        nnode = len(self.node_list) + 1
        r = self.connect_circle_dist * math.sqrt((math.log(nnode) / nnode ))

        if hasattr(self, 'expand_dis'):        #hasattr(): true if object has the beforementioned attribute
            r = min(r, self.expand_dis)

        dist_list = [(node.x -new_node.x)**2 + (node.y - new_node.y)**2 for node in self.node_list]
        near_inds = [dist_list.index(i) for i in dist_list if i <= r**2 ]     #r^2 = ~866

        return near_inds

    def choose_parent( self, new_node, near_inds):
        if not near_inds:
            return None

        costs = []
        for i in near_inds:
            near_node = self.node_list[i]
            t_node = self.steer(near_node, new_node)
            if t_node and self.check_collision(t_node):
                costs.append(self.calc_new_cost(near_node, new_node))

            else:
                costs.append(float("inf"))       #the cost of collision node

        min_cost = min(costs)

        if min_cost == float("inf"):
            return None

        min_ind = near_inds[costs.index(min_cost)]
        new_node = self.steer(self.node_list[min_ind], new_node)
        new_node.cost = min_cost

        return new_node

    def calc_new_cost(self, from_node, to_node):
        d, _ = self.calc_distance_and_angle(from_node, to_node)
        return from_node.cost + d

    def steer(self, from_node, to_node, extend_length=float('inf')):

        new_node = self.Node(from_node.x, from_node.y)
        d, theta = self.calc_distance_and_angle(new_node, to_node)

        new_node.path_x = [new_node.x]
        new_node.path_y = [new_node.y]

        if extend_length > d:
            extend_length = d

        n_expand = int(math.floor(extend_length     / self.path_resolution))

        for _ in range(n_expand):
            new_node.x += self.path_resolution * math.cos(theta)
            new_node.y += self.path_resolution * math.sin(theta)

            new_node.path_x.append(new_node.x)
            new_node.path_y.append(new_node.y)

        d, _ = self.calc_distance_and_angle(new_node, to_node)

        if d <= self.path_resolution:
            new_node.path_x.append(to_node.x)
            new_node.path_y.append(to_node.y)
            new_node.x = to_node.x
            new_node.y = to_node.y

        new_node.parent = from_node

        return new_node

    @staticmethod
    def calc_distance_and_angle(from_node, to_node):

        #print(from_node)
        dx = (to_node.x) - (from_node.x)
        dy = to_node.y - from_node.y

        d = math.hypot(dx, dy)
        theta = math.atan2(dy, dx)
        return d, theta

    @staticmethod
    def get_nearest_node_index(node_list, rnd_node):
        dlist = [(node.x -rnd_node.x)**2 + (node.y - rnd_node.y)**2 for node in node_list]
        minind = dlist.index(min(dlist))

        return minind

    @staticmethod
    def remove_node_element(node, i):
        del node.path_x[i]
        del node.path_y[i]

    #@staticmethod
    def check_collision(self, node):
        if node is None:
            return False # collision

        i = 0

        while (i <= len(node.path_x) -1):

            x_node = node.path_x[i]
            y_node = node.path_y[i]
            in_map = check_node(x_node, y_node)

            if in_map == False:

                self.remove_node_element(node, i)
                break
            x_int, y_int = get_voxel_inds(x_node, y_node)
            if (math.isnan(x_int) == True):
                i = i+1
                continue
            if ( x_int >= local_map.resolution.get('width') or y_int >= local_map.resolution.get('height')):
                break

            try:
                if (local_map.map_matrix[y_int][x_int] >= 30 ):
                    return False
            except IndexError:
                pass
            i = i +1

        return True

    @staticmethod
    def check_2node_collision( start_node, end_node ):
        x_1 , y_1 = start_node.x, start_node.y
        x_2 , y_2 = end_node.x, end_node.y

        simple_angle, angle = calc_angle([x_1, y_1], [x_2, y_2])

        d_x = distance(x_1, x_2)
        d_y = distance(y_1, y_2)

        if(d_x is None or d_y is None): #Needed for Take-off
            return True
        if (abs(d_x) <= local_map.resolution.get('pixel_size') and abs(d_y) <= local_map.resolution.get('pixel_size')):

            return True

        start_position = [x_1, y_1]
        next_node = [x_2, y_2]
        list_to_check = get_list_of_voxels_to_check(simple_angle, d_x, d_y, start_position, next_node )
        waited_once = False #if neccesarry waiting for new local_map
        for voxel in list_to_check:

            if voxel[0] >= local_map.resolution.get('width'):
                if waited_once == False:
                    rospy.sleep(0.05)
                    waited_once = True
                voxel[0] = local_map.resolution.get('width') -1

            if voxel[1] >= local_map.resolution.get('height'):
                if waited_once == False:
                    rospy.sleep(0.05)
                    waited_once = True
                #print("voxel height outside of map")
                voxel[1] = local_map.resolution.get('height') -1

            if (local_map.map_matrix[voxel[1]][voxel[0]] >= 30):

                return False
            #that is a crash

        return True

def distance(x_1, x_2):
    if x_1 > x_2:   #returns negative per definition
        return abs(x_1 - x_2)*(-1)

    if x_2 > x_1: #returns positiv number
        return abs(x_2 - x_1)

    if x_2 == x_1:
        return 0

def get_list_of_voxels_to_check(angle, d_x, d_y, start, next_node): #d_x & d_y in: +- ; angle = simple_angle ; start = [x, y] in meter

    #simple_angle is the angle defined by d_x and d_y
    simple_angle = math.radians(angle)
    resolution = local_map.resolution.get('pixel_size')

    voxel_to_check= []
    multiplier_x = 1    #in positive direction
    multiplier_y = 1    #in pos direction
    if d_x < 0:
        multiplier_x = -1
    if d_y < 0:
        multiplier_y = -1

    i = 0
    while ( ( (abs(d_x) - abs((i*resolution*math.sin(simple_angle)))) >= 0) or ( (abs(d_y) - abs((i*resolution*math.cos(simple_angle)))) >=0) ):
        if d_x >= 3000 or d_y >= 3000:
            width = local_map.resolution.get('width')
            height = local_map.resolution.get('height')
            res = local_map.resolution.get('pixel_size')
            #before take off mavros/trajectory/desired is < -10000
            return [[width/res , height/res ]]

        x = start[0] + (d_x - multiplier_x*i*math.sin(simple_angle)*resolution)
        y = start[1] + (d_y - multiplier_y*i*math.cos(simple_angle)*resolution)

        x_ind, y_ind = get_voxel_inds(x, y)
        if (math.isnan(x_ind) == True or math.isnan(y_ind) == True ):
            i = i+1
            continue

        check_voxel = [x_ind, y_ind] # width, height = x , y
        voxel_to_check.append(check_voxel)
        i = i +1

    return voxel_to_check

def check_node(x_node, y_node):
    pixel = local_map.resolution.get('pixel_size')
    width = local_map.resolution.get('width')
    height = local_map.resolution.get('height')
    if (local_map.map_origin[0] >= x_node or local_map.map_origin[1] >= y_node):
        return False

    else:
        if (x_node > local_map.map_origin[0]+ pixel*width):
            return False

        if (y_node > local_map.map_origin[1] + pixel*height):
            return False
        return True

def calc_map_pixel_steps(angle, d_x, d_y):

    delta_x = 0     # moves cells each line by delta_x to a side
    x_step = 1      #numbers of cells in line to check
    pixel_size = local_map.resolution.get('pixel_size')
    i = 1
    while i <= d_x:
        #print("x_step: " +str(i))
        if (i*pixel_size*math.tan(abs(angle)) >= pixel_size):
            x_step = i

            delta_x = i -1

            if ( delta_x == 0):

                x_step = x_step +1
                delta_height = pixel_size*math.tan(abs(angle))
                numerator = int(delta_height / pixel_size) +1

                delta_x = float(1.0/numerator)

            break
        i = i +1

    return delta_x, x_step

def calc_angle(start_vec, end_vec):
    if math.isnan(start_vec[0]) or math.isnan(start_vec[1]):
        return 0 ,0
    if math.isnan(end_vec[0]) or math.isnan(end_vec[1]):
        return 0 ,0

    d_x = distance(start_vec[0], end_vec[0])
    d_y = distance( start_vec[1], end_vec[1])

    if d_y != 0:
        simple_angle = abs(math.degrees(math.atan(d_x / d_y)))
    else:
        simple_angle = 0

    if (d_y >= 0):
        if (d_x > 0):
            angle =  math.degrees(math.atan(d_y/d_x))
            return simple_angle, angle
        else:
            if (d_x < 0):
                angle = (180 + math.degrees(math.atan(d_y / d_x)))
                return simple_angle, angle
            if (d_x == 0):
                return simple_angle, 90
    else:
        #print("angle is negative")
        if d_x > 0:
            angle = -180 -  math.degrees(math.atan(d_y / d_x))
            return simple_angle, angle
        if d_x < 0:
            angle = -1*  math.degrees((math.atan(d_y / d_x)))
            return simple_angle, angle
        if d_y == 0 and d_x >0:
            return simple_angle, 0
        if d_y == 0 and d_x < 0:
            return simple_angle, -180

        else: #d_x == 0
            return simple_angle, -90

#returns the indexs of local_map.map_matrix[x_int][y_int] (column and height)
def get_voxel_inds(x , y):
    x_origin = local_map.map_origin[0]
    y_origin = local_map.map_origin[1]
    resolution = local_map.resolution.get('pixel_size')

    if (x >= x_origin and y >= y_origin):
        x_int = 0
        while (x_int +1 <= local_map.resolution.get('width') ):
            if (x_origin + resolution + (x_int * resolution) >= x):
                break
            x_int = x_int +1
        y_int = 0
        while (y_int <= local_map.resolution.get('height') ):

            if (y_origin + resolution + (y_int*resolution) >= y) :

                break
            y_int = y_int +1

    else:
        x_int = float('NaN')   #CAN CAUSE ERROR???
        y_int = float('NaN')   #CAN Cause ERROR????

    #print(x_int, y_int)
    return x_int , y_int

###########################END CLASS#########################################

def subscribe_to_topics():
    rospy.Subscriber("/local_map", nav_msgs.msg.OccupancyGrid, local_map.local_map_splitter)
    rospy.Subscriber("/mavros/local_position/pose", geometry_msgs.msg.PoseStamped, drone_1.set_drone_position)
    rospy.Subscriber("/start_planning", RRT_Star_Call, local_map.set_target)
    #print("I did subscribe")

def publish_path(path):#, pub): #path = [[x_1, y_1], [x_2,..], ..] #first= target_position
    safe_path = Path_coll()  #new ROS message
    safe_path.stamp = rospy.get_rostime()
    pub = rospy.Publisher('safe_path', Path_coll, queue_size=10)
    #path is in opposite order: 1.node target_position
    j = len(path) - 1
    while j >= 0:
        for node in path[j]:
            safe_path.safe_path.append(node)
        j = j -1

    #GONE for good
    #safe_path.cost = calculate_cost(path)
    ##Always accept directr pass: cost should be lower than the true way
    #if len(path)==2:
    #    safe_path.cost = safe_path.cost -3
    #safe_path.desired_altitude = drone_1.desired_altitude

    #path is delivered in reverse order, hence the zero
    safe_path.target_position = path[0]
    pub.publish(safe_path)

def calculate_cost(path):

    i = 0
    cost = 0
    while i <= len(path) -2:

        x_1 ,y_1 = path[i][0], path[i][1]
        x_2, y_2 = path[i+1][0], path[i +1][1]
        d_x = (x_2 - x_1) **2
        d_y = (y_2 - y_1) **2
        cost = cost + math.sqrt(d_x + d_y)
        i = i +1
    return cost

def pub_map_maximizer(i):

    pub = rospy.Publisher('/maximize_map', Maximizer, queue_size= 10)
    msg = Maximizer()

    if i <= 5:
        msg.expand_map = i
        msg.stamp = rospy.get_rostime()
        pub.publish(msg)

if __name__ == "__main__":
    rospy.init_node('local_planner', anonymous= True)
    drone_position = [0, 0, 0, 0, 0, 0]
    drone_1 = Drone(drone_position, desired_altitude=0)
    local_map = Local_Map([0, 0], {"width": 0, "height": 0, "pixel_size": 1}, [[0][0]], [0, 0, 0], maximizer=0 )
    subscribe_to_topics()
    rospy.sleep(0.5)  #NEEDED TO GET DATA FROM Subscriber

    start_0 = [drone_1.drone_pose[0], drone_1.drone_pose[1]]
    goal_0 =  [local_map.target_position[0], local_map.target_position[1]]

    local_map_origin = local_map.map_origin
    pixel_size = local_map.resolution.get('pixel_size')
    width = local_map.resolution.get('width')
    height = local_map.resolution.get('height')

    rrt_star = RRT_Star( start_0 , goal_0, rand_area_width=[local_map.map_origin[0], local_map_origin[0]+ width*pixel_size ], rand_area_height= [local_map.map_origin[1], local_map_origin[1] + height*pixel_size ] )
    start_time = time.time()
    path = rrt_star.planning()
    pub = rospy.Publisher('safe_path', Path_coll, queue_size=10)

    if path is None:
        del rrt_star.node_list

    else:

        publish_path(path)#, pub )
        del rrt_star.node_list

    last_time_published = time.time()
    while not rospy.is_shutdown():
        if rrt_star.start_planning == True:
            try:
                print("#############")
                print("Start")
                for value in local_map.target_position:
                    if math.isnan(value) == True:
                        path = [ [ ] ]
                        #publish_path(path) #, pub)
                        pass
                start_time = time.time()
                check_also_direct_path = False
                if check_also_direct_path == True:
                    if (rrt_star.check_2node_collision(rrt_star.start , rrt_star.goal_node) == True):
                        print("end: " +str(rrt_star.goal_node.x)+ ", " + str(rrt_star.goal_node.y))
                        path = [ [ rrt_star.goal_node.x, rrt_star.goal_node.y ], [rrt_star.start.x, rrt_star.start.y] ]
                        publish_path(path) #, pub)
                        if local_map.maximizer >= 0:
                            local_map.maximizer = 0
                            pub_map_maximizer(local_map.maximizer)
                        #rospy.sleep(0.025)  #publishes very fast(~200 per sec)

                path = rrt_star.planning()
                if path is not None :
                    #PATH IS SOMETIMES JUST 2 Points without target position is path_list
                    print("safe::: " + str(path))
                    publish_path(path)#, pub)
                    last_time_published = time.time()
                    del rrt_star.node_list
                    #rospy.sleep(0.025)
                else:
                    if (time.time() - last_time_published >= 2):
                        last_time_published = time.time()
                        if local_map.maximizer <= 3:
                            local_map.maximizer = local_map.maximizer + 1
                            pub_map_maximizer(local_map.maximizer)
                    del rrt_star.node_list

                print(rrt_star.start_planning)
                print("End")
                print("#############")
            except KeyboardInterrupt:
                sys.exit()
        else:
            print("I am not needed right now!")
            rospy.sleep(0.5)

        #rospy.spin()