text:
# Purpose: Generate Elastic Net Predictions
# Inputs: outputs of MI Feature Selection
# Outputs: Prediction file
# Machine: Laptop, ~10/15 mins

# Packages Used
import argparse
import sys
import numpy as np
import utils
import models
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

# Random Seed
np.random.seed(0)

parser = argparse.ArgumentParser(
  description='Elastic net script.')
parser.add_argument('input_features_filename',
                    help='input file containing the all the features per child.')
parser.add_argument('input_labels_filename',
                    help='input file containing the all the known labels per child. '
                    'It must be joinable with input_features_filename by the challenge '
                    'id column.')
parser.add_argument('input_predictions_filename',
                    help='input file containing the features we need to make prediction '
                    'for their unknown labels.')
parser.add_argument('output_prediction_filename',
                    help='The output file will contain our predictions for the rows in '
                    'input_prediction_filename.')
parser.add_argument('-al', '--algorithm', dest='algorithm',
                    default='elastic-net', choices=['elastic-net', 'lasso', 'lars'],
                    help='The algorithm to use for training/prediction.')
parser.add_argument('-t', '--transformation', dest='transformation',
                    default='standard', choices=['standard', 'normal', 'minmax'],
                    help='The method to transform features: normalize, standard scaling '
                    'or minmax scaling.')
parser.add_argument('-trs', '--train_size', dest='train_size',
                    type=int, default=-1,
                    help='The size of train set. It must be a whole integer, greater '
                    'than 1 which represents the absolute number of data points to use '
                    'for training. Default is -1 which means all available data after '
                    'splitting into test and train is used. There will be some '
                    'undersampling of the majority class, if train data needs to be '
                    'balanced.')
parser.add_argument('-tes', '--test_size', dest='test_size',
                    type=int, default=800,
                    help='The size of test set. It must be a whole integer, greater '
                    'than 1 which represents the absolute number of data points to use '
                    'for testing. You need to make sure test and train size do not '
                    'exceed the total size and there are enough samples of each class '
                    'in case train data is balanced.')
parser.add_argument('-lc', '--label_columns', dest='label_columns',
                    nargs='+', default=['gpa', 'grit', 'materialHardship'],
                    help='The column names in input_labels_filename to predict. '
                    'We essentially build a model for each requested column.')
parser.add_argument('-nj', '--num_jobs', dest='num_jobs', type=int, default=-1,
                    help='Number of jobs to instantiate for grid search. Default is '
                    '-1, which corresponds to the number of cores in the machine')

# cross validation flags: they will be ignored if cross validation is not skipped.
parser.add_argument('-sc', '--skip_cross_validation', dest='skip_cross_validation',
                    default=False, action='store_true',
                    help='If specified, skips cross validation and uses the parameters '
                    'given below. ')
parser.add_argument('-na', '--num_alphas', dest='num_alphas', type=int, default=200,
                    help='The number of alphas to try for each l1 ratio, in case '
                    'cross validation is not skipped. Defaults to 100.')
parser.add_argument('-l', '--l1_ratio', dest='l1_ratio', default=0.15, type=float,
                    help='The Elastic Net mixing parameter, must be a float between 0 '
                    'and 1, with 0 <= l1_ratio <= 1. It determines how aggressive is the '
                    'feature selection in elastic net. Must be specified only if penatly '
                    'is elasticnet, as default (0.15) might not be good.  Relevent only '
                    'if cross validation is skipped.')
parser.add_argument('-a', '--alpha', dest='alpha', default=1, type=float,
                    help='alpha multiplies the penalty terms and thus determines '
                    'the level of regularization. Defaults to 1.0. alpha = 0 is '
                    'equivalent to an ordinary least square.')
args = parser.parse_args()

def main():
  add_transformed_vars = True

  all_predictions = None
  for label_column in args.label_columns:
    print('Training for Label {}'.format(label_column))
    (all_features_predictions_df,
     train_features, transformed_train_labels, train_labels,
     test_features, transformed_test_labels, test_labels,
     all_feature_names, all_label_names) = utils.prepare_data(args.input_features_filename,
                                                              args.input_labels_filename,
                                                              args.input_predictions_filename,
                                                              args.transformation,
                                                              label_column,
                                                              args.train_size,
                                                              args.test_size,
                                                              add_transformed_vars,
                                                              False,
                                                              False)
    all_features = np.concatenate((train_features, test_features), axis=0)
    all_labels = np.concatenate((train_labels, test_labels), axis=0)
    all_transformed_labels = np.concatenate((transformed_train_labels, transformed_test_labels), axis=0)
    #print 'Max label values: {}'.format(np.max(all_labels))
    #print 'Min label values: {}'.format(np.min(all_labels))
    if all_predictions is None:
      all_predictions = all_features_predictions_df[all_label_names]

    #print('\n*****************************\n')
    (model, best_params) = models.train_model(all_features,
                                              all_transformed_labels,
                                              args.algorithm,
                                              args.num_alphas,
                                              args.skip_cross_validation,
                                              args.alpha, args.l1_ratio,
                                              args.num_jobs)
    train_y_true = all_labels
    train_y_pred = utils.predict(all_features, model, label_column)
    #print 'R-Squared on train is {}'.format(r2_score(train_y_true, train_y_pred))

    #print('\n')
    #print 'MAE on train: {}'.format(mean_absolute_error(train_y_true, train_y_pred))
    #print 'MSE on train: {}'.format(mean_squared_error(train_y_true, train_y_pred))

    coefs = model.coef_
    print('Optimal parameters in final model is {}'.format(best_params))
    # info on most predictive coefs
    print('Model has {} non-zero coefficients.'.format(len(coefs[coefs != 0])))
    highest_indices = np.abs(coefs).argsort()[::-1]
    #print 'Top predictors of {}'.format(label_column)
    for i in range(min(10, len(coefs[coefs != 0]))):
      idx = highest_indices[i]
    #  print '{}: {}'.format(all_feature_names[idx], coefs[idx])

    # predict on all and write to file
    prediction_features = all_features_predictions_df[all_feature_names]
    label_predictions = utils.predict(prediction_features, model, label_column)
    # correct if below/above min/max
    label_predictions[label_predictions < np.min(all_labels)] = np.min(all_labels)
    label_predictions[label_predictions > np.max(all_labels)] = np.max(all_labels)
    all_predictions.loc[:,label_column] = label_predictions

    print('\n*****************************')
    print('*****************************\n')
  all_predictions.to_csv(args.output_prediction_filename)

if __name__ == '__main__':
  main()


text:
import sys
import numpy as np
import gym
import minerl
import torch
import argparse
from collections import Counter, namedtuple
import torch.nn as nn
import torch.nn.functional as F
from model import Model
import utils
from utils import *
from losses import *
from random import randint

NUM_ACTIONS = 11

# Test J_L2 Loss
behavior_network = Model()
print(J_L2(behavior_network))

# Test J_E loss
pred = torch.arange(9).reshape((3,3))
demo = torch.tensor([[0,0,0],[0,1,0],[1,0,0]])
margin = 1
pred + margin*demo
x,_ = torch.max(pred + margin*demo,dim=1)
y = torch.sum(pred * demo,dim=1)
print(J_E(pred, demo))

target_network = Model()
def generate_dummy_states(batch_size):
    return torch.rand(batch_size,3,64,64), torch.rand(batch_size, 2)
def generate_dummy_is_demo(batch_size):
    is_demo = torch.zeros((batch_size, NUM_ACTIONS))
    for i in range(batch_size):
        is_demo[i,randint(0,1)] = 1
    return is_demo
samples = (generate_dummy_states(3),
            None,
            None,
            generate_dummy_is_demo(3),
            None)
loss = J_Q(target_network,
           behavior_network,
           samples)
print(loss)

ext: py

max_stars_repo_name: jarbus/minerl

max_issues_repo_name: jarbus/minerl

max_forks_repo_name: jarbus/minerl

reason: import numpy


text:
# pylint: disable=missing-function-docstring
"""Tests for '_discrete.py' file"""
import numpy as np
from numpy.testing import assert_array_equal, assert_array_almost_equal
import scipy.stats.distributions as distrs
import pytest

from randomvars._discrete import Disc
from randomvars._utils import _tolerance
from randomvars.tests.commontests import (
    DECIMAL,
    h,
    _test_equal_rand,
    _test_input_coercion,
    _test_from_rv_rand,
    _test_from_sample_rand,
    _test_log_fun,
    _test_one_value_input,
    _test_rvs_method,
)
from randomvars.options import config

DISTRIBUTIONS_FINITE = {
    "fin_bernoulli": distrs.bernoulli(p=0.9),
    "fin_binom": distrs.binom(n=10, p=0.7),
    "fin_binom_wide": distrs.binom(n=101, p=0.7),
    "fin_randint": distrs.randint(low=-5, high=4),
}

DISTRIBUTIONS_INFINITE = {
    "inf_geom": distrs.geom(p=0.1),
    "inf_poisson": distrs.poisson(mu=5),
    "inf_poisson_big": distrs.poisson(mu=1001),
}

DISTRIBUTIONS_SHIFTED = {
    "shift_binom": distrs.binom(n=10, p=0.7, loc=-np.pi),
    "shift_randint": distrs.randint(low=-5, high=4, loc=np.pi),
}

DISTRIBUTIONS = {
    **DISTRIBUTIONS_FINITE,
    **DISTRIBUTIONS_INFINITE,
    **DISTRIBUTIONS_SHIFTED,
}

class TestDisc:
    """Regression tests for `Disc` class"""

    def test_init_errors(self):
        def check_one_input(def_args, var):
            with pytest.raises(TypeError, match=f"`{var}`.*numpy array"):
                def_args[var] = {"a": None}
                Disc(**def_args)
            with pytest.raises(TypeError, match=f"`{var}`.*float"):
                def_args[var] = ["a", "a"]
                Disc(**def_args)
            with pytest.raises(TypeError, match=f"`{var}`.*finite values"):
                def_args[var] = [0, np.nan]
                Disc(**def_args)
            with pytest.raises(TypeError, match=f"`{var}`.*finite values"):
                def_args[var] = [0, np.inf]
                Disc(**def_args)
            with pytest.raises(ValueError, match=f"`{var}`.*1d array"):
                def_args[var] = [[0, 1]]
                Disc(**def_args)

        check_one_input({"p": [0.2, 0.8]}, "x")
        check_one_input({"x": [0, 1]}, "p")

        with pytest.raises(ValueError, match="[Ll]engths.*match"):
            Disc([0, 1], [1, 1, 1])

        with pytest.warns(UserWarning, match="`x`.*not sorted.*`x` and `p`"):
            rv = Disc([1, 0], [0.2, 0.8])
            rv_ref = Disc([0, 1], [0.8, 0.2])
            _test_equal_rand(rv, rv_ref)

        with pytest.raises(ValueError, match="`p`.*negative"):
            Disc([0, 1], [0.8, -1])

        with pytest.raises(ValueError, match="`p`.*no positive"):
            Disc([0, 1], [0, 0])

    def test_init(self):
        x_ref = np.array([0.1, 1, 2])
        p_ref = np.array([0.1, 0.2, 0.7])
        rv_ref = Disc(x_ref, p_ref)

        # Simple case with non-numpy input
        rv_1 = Disc(x=x_ref.tolist(), p=p_ref.tolist())
        _test_equal_rand(rv_1, rv_ref)

        # Check if `p` is normalized
        rv_2 = Disc(x=x_ref, p=10 * p_ref)
        _test_equal_rand(rv_2, rv_ref)

        # Check that zero probability is allowed
        rv_3 = Disc(x=[0, 1, 3], p=[0, 0.5, 0.5])
        assert_array_equal(rv_3.x, np.array([0, 1, 3]))
        assert_array_equal(rv_3.p, np.array([0, 0.5, 0.5]))

        # Check if `x` and `p` are rearranged if not sorted
        with pytest.warns(UserWarning, match="`x`.*not sorted"):
            rv_4 = Disc(x=x_ref[[1, 0, 2]], p=p_ref[[1, 0, 2]])
            _test_equal_rand(rv_4, rv_ref)

        # Check if duplicated values are removed from `x`
        with pytest.warns(UserWarning, match="duplicated"):
            # First pair of xy-grid is taken among duplicates
            rv_5 = Disc(x=x_ref[[0, 1, 1, 2]], p=p_ref[[0, 1, 2, 2]])
            _test_equal_rand(rv_5, rv_ref)

    def test_str(self):
        rv = Disc([0, 2, 4], [0.125, 0, 0.875])
        assert str(rv) == "Discrete RV with 3 values (support: [0.0, 4.0])"

        # Uses singular noun with one value
        rv = Disc([1], [1])
        assert str(rv) == "Discrete RV with 1 value (support: [1.0, 1.0])"

    def test_properties(self):
        x = np.arange(10)
        p = np.repeat(0.1, 10)
        rv = Disc(x, p)

        assert list(rv.params.keys()) == ["x", "p"]
        assert_array_equal(rv.params["x"], x)
        assert_array_equal(rv.params["p"], p)

        assert_array_equal(rv.x, x)
        assert_array_equal(rv.p, p)
        assert rv.a == 0
        assert rv.b == 9

    def test_support(self):
        rv = Disc([0.5, 1.5, 4.5], [0.25, 0.375, 0.375])
        assert rv.support() == (0.5, 4.5)

    def test_compress(self):
        # Keep only x-values with positive probabilities
        _test_equal_rand(
            Disc([0, 1, 2, 3], [0.0, 0.5, 0.0, 0.5]).compress(),
            Disc([1, 3], [0.5, 0.5]),
        )

        # If nothing to compress, self should be returned
        rv = Disc([0, 1], [0.5, 0.5])
        assert rv.compress() is rv

    def test_from_rv_basic(self):
        x = [0, 1, 5]
        p = [0.1, 0.4, 0.5]
        rv = distrs.rv_discrete(values=(x, p))
        rv_out = Disc.from_rv(rv)
        rv_ref = Disc(x=x, p=p)
        _test_equal_rand(rv_out, rv_ref)

        # Objects of `Rand` class should be `convert()`ed
        _test_from_rv_rand(cls=Disc, to_class="Disc")

        # Works with single-valued rv
        rv_single = distrs.rv_discrete(values=(2, 1))
        rv_out = Disc.from_rv(rv_single)
        assert_array_equal(rv_out.x, [2])

    def test_from_rv_errors(self):
        # Absence of either `cdf` or `ppf` method should result intro error
        class Tmp:
            pass

        tmp1 = Tmp()
        tmp1.ppf = lambda x: np.where((0 <= x) & (x <= 1), 1, 0)
        with pytest.raises(ValueError, match="cdf"):
            Disc.from_rv(tmp1)

        tmp2 = Tmp()
        tmp2.cdf = lambda x: np.where((0 <= x) & (x <= 1), 1, 0)
        with pytest.raises(ValueError, match="ppf"):
            Disc.from_rv(tmp2)

        # Having bad `cdf()` or `ppf()` that result into infinite loop should
        # result into error
        tmp3 = Tmp()
        tmp3.cdf = lambda x: x
        tmp3.ppf = lambda x: 0.5
        with pytest.raises(ValueError, match="Couldn't get increase"):
            Disc.from_rv(tmp3)

    def test_from_rv_options(self):
        # Usage of `small_prob` option
        x = [1, 2, 3]
        p = [0.5, 0.125, 0.375]
        rv = distrs.rv_discrete(values=(x, p))

        with config.context({"small_prob": 0.125 + 1e-5}):
            rv_out = Disc.from_rv(rv)
            rv_ref = Disc([1, 3], [0.5, 0.5])
            _test_equal_rand(rv_out, rv_ref)

    def test_from_sample_basic(self):
        x = np.array([0.1, -100, 1, np.pi, np.pi, 1, 3, 0.1])

        rv = Disc.from_sample(x)
        rv_ref = Disc(x=[-100, 0.1, 1, 3, np.pi], p=[0.125, 0.25, 0.25, 0.125, 0.25])
        assert isinstance(rv, Disc)
        _test_equal_rand(rv, rv_ref)

    def test_from_sample_errors(self):
        with pytest.raises(TypeError, match="numpy array with float"):
            Disc.from_sample(["a"])

        with pytest.raises(ValueError, match="1d"):
            Disc.from_sample([[1], [2]])

        # Warning is given with default discrete estimator if not all values
        # are finite
        with pytest.warns(UserWarning, match="has non-finite"):
            rv = Disc.from_sample([1, 2, np.nan])
            rv_ref = Disc(x=[1, 2], p=[0.5, 0.5])
            _test_equal_rand(rv, rv_ref)

        # Error is given with default discrete estimator if there is no finite
        # values
        with pytest.raises(ValueError, match="doesn't have finite values"):
            Disc.from_sample([-np.inf, np.nan, np.inf])

    def test_from_sample_options(self):
        binom = distrs.binom(n=10, p=0.5)

        rng = np.random.default_rng(101)
        x = binom.rvs(100, random_state=rng)

        # "estimator_disc"
        def single_value_estimator(x):
            return np.array([1.0]), np.array([1.0])

        with config.context({"estimator_disc": single_value_estimator}):
            rv = Disc.from_sample(x)
            _test_equal_rand(rv, Disc(x=[1.0], p=[1.0]))

        # "estimator_disc" which returns allowed classes
        ## `Rand` class should be forwarded to `from_rv()` method
        _test_from_sample_rand(
            cls=Disc,
            sample=x,
            estimator_option="estimator_disc",
        )

        ## "Scipy" distribution should be forwarded to `Disc.from_rv()`
        rv_binom = distrs.binom(n=10, p=0.5)
        with config.context({"estimator_disc": lambda x: rv_binom}):
            rv = Disc.from_sample(np.asarray([0, 1, 2]))
            rv_ref = Disc.from_rv(rv_binom)
            _test_equal_rand(rv, rv_ref)

    def test_pdf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        with pytest.raises(AttributeError, match=r"Use `pmf\(\)`"):
            rv.pdf(0.5)

    def test_logpdf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        with pytest.raises(AttributeError, match=r"Use `logpmf\(\)`"):
            rv.logpdf(0.5)

    def test_pmf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])

        # Regular checks
        x = np.array([0, 0.5, 1, 3, 3 + 0.5 * _tolerance(3), 3 + 1.5 * _tolerance(3)])
        assert_array_equal(rv.pmf(x), np.array([0.0, 0.1, 0.2, 0.7, 0.7, 0.0]))

        # Coercion of not ndarray input
        _test_input_coercion(rv.pmf, x)

        # Bad input
        x = np.array([-np.inf, np.nan, np.inf])
        assert_array_equal(rv.pmf(x), np.array([0, np.nan, 0]))

        # Using different tolerance
        with config.context({"base_tolerance": 0.1}):
            assert_array_equal(rv.pmf([1 + 0.11, 1 + 0.09]), [0.0, 0.2])

        # Broadcasting
        x = np.array([[-1, 0.5], [2, 4]])
        assert_array_equal(rv.pmf(x), np.array([[0.0, 0.1], [0.0, 0.0]]))

        # One value input
        _test_one_value_input(rv.pmf, 0.5)
        _test_one_value_input(rv.pmf, -1)
        _test_one_value_input(rv.pmf, np.nan)

    def test_logpmf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        _test_log_fun(rv.logpmf, rv.pmf, x_ref=[-1, 0.5, 3, np.inf, np.nan])

    def test_cdf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])

        # Regular checks
        x = np.array([-10, 0.5 - h, 0.5, 0.5 + h, 1 - h, 1, 1 + h, 3 - h, 3, 3 + h, 10])
        assert_array_almost_equal(
            rv.cdf(x),
            np.array([0, 0, 0.1, 0.1, 0.1, 0.3, 0.3, 0.3, 1, 1, 1]),
            decimal=DECIMAL,
        )

        # Coercion of not ndarray input
        _test_input_coercion(rv.cdf, x)

        # Bad input
        x = np.array([-np.inf, np.nan, np.inf])
        assert_array_equal(rv.cdf(x), np.array([0, np.nan, 1]))

        # Broadcasting
        x = np.array([[-1, 0.5], [2, 4]])
        assert_array_almost_equal(
            rv.cdf(x), np.array([[0.0, 0.1], [0.3, 1.0]]), decimal=DECIMAL
        )

        # One value input
        _test_one_value_input(rv.cdf, 0.5)
        _test_one_value_input(rv.cdf, -1)
        _test_one_value_input(rv.cdf, np.nan)

    def test_logcdf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        _test_log_fun(rv.logcdf, rv.cdf, x_ref=[-1, 0.5, 3, np.inf, np.nan])

    def test_sf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        x_ref = [-1, 0.5, 3, np.inf, np.nan]
        assert_array_equal(rv.sf(x_ref), 1 - rv.cdf(x_ref))

    def test_logsf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        _test_log_fun(rv.logsf, rv.sf, x_ref=[-1, 0.5, 3, np.inf, np.nan])

    def test_ppf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])

        # Regular checks
        ## Outputs for q=0 and q=1 should be equal to minimum and maximum elements
        q = np.array([0, 0.1 - h, 0.1, 0.1 + h, 0.3 - h, 0.3, 0.3 + h, 1 - h, 1])
        assert_array_equal(rv.ppf(q), np.array([0.5, 0.5, 0.5, 1, 1, 1, 3, 3, 3]))

        # Coercion of not ndarray input
        _test_input_coercion(rv.ppf, q)

        # Bad input
        q = np.array([-np.inf, -h, np.nan, 1 + h, np.inf])
        assert_array_equal(
            rv.ppf(q), np.array([np.nan, np.nan, np.nan, np.nan, np.nan])
        )

        # Broadcasting
        q = np.array([[0, 0.5], [0.0, 1.0]])
        assert_array_equal(rv.ppf(q), np.array([[0.5, 3], [0.5, 3]]))

        # One value input
        _test_one_value_input(rv.ppf, 0.25)
        _test_one_value_input(rv.ppf, -1)
        _test_one_value_input(rv.ppf, np.nan)

    def test_isf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        q_ref = np.array([0, 0.7 - h, 0.7, 0.7 + h, 0.9 - h, 0.9, 0.9 + h, 1 - h, 1])

        # Output is lowest value `t` (within supp.) for which `P(X > t) <= q`
        assert_array_equal(rv.isf(q_ref), rv.ppf(1 - q_ref))

    def test_rvs(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])

        _test_rvs_method(rv)

    def test__cdf_spline(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        x = np.array([-10, 0.5 - h, 0.5, 0.5 + h, 1 - h, 1, 1 + h, 3 - h, 3, 3 + h, 10])
        assert_array_equal(rv._cdf_spline(x), rv.cdf(x))

    def test_integrate_cdf(self):
        rv = Disc([0.5, 1, 3], [0.1, 0.2, 0.7])
        assert np.allclose(rv.integrate_cdf(-10, 10), 0.1 * 0.5 + 0.3 * 2 + 1 * 7)

    def test_convert(self):
        import randomvars._boolean as bool
        import randomvars._continuous as cont
        import randomvars._mixture as mixt

        rv = Disc([0.0, 0.5, 1, 3], [0.05, 0.05, 0.2, 0.7])

        # By default and supplying `None` should return self
        assert rv.convert() is rv
        assert rv.convert(None) is rv

        # Converting to Bool should result into boolean with probability of
        # `False` being probability of 0
        out_bool = rv.convert("Bool")
        assert isinstance(out_bool, bool.Bool)
        assert_array_almost_equal(out_bool.prob_false, rv.pmf(0.0), decimal=DECIMAL)

        # Converting to Cont should result into continuous RV with the same `x`
        # values as in input's xp-grid
        out_cont = rv.convert("Cont")
        assert isinstance(out_cont, cont.Cont)
        assert_array_equal(out_cont.x, rv.x)

        ## Converting to Cont if there is one element should raise an error
        with pytest.raises(ValueError, match="one"):
            Disc([1], [1]).convert("Cont")

        # Converting to own class should return self
        out_disc = rv.convert("Disc")
        assert out_disc is rv

        # Converting to Mixt should result into degenerate mixture with only
        # discrete component
        out_mixt = rv.convert("Mixt")
        assert isinstance(out_mixt, mixt.Mixt)
        assert out_mixt.disc is rv
        assert out_mixt.weight_disc == 1.0

        # Any other target class should result into error
        with pytest.raises(ValueError, match="one of"):
            rv.convert("aaa")

class TestFromRVAccuracy:
    """Accuracy of `Disc.from_rv()`"""

    def test_tails(self):
        test_passed = {
            name: TestFromRVAccuracy.is_from_rv_small_tails(distr)
            for name, distr in DISTRIBUTIONS.items()
        }

        assert all(test_passed.values())

    def test_small_prob_detection(self):
        # Currently small probabilities can be not detected due to "stepping"
        # procedure of x-values detection.

        # If small probability element is not detected, its probability is
        # "squashed" to the next (bigger) detected element
        rv_1 = distrs.rv_discrete(values=([1, 2, 3], [0.5, 0.125, 0.375]))
        with config.context({"small_prob": 0.125 + 1e-5}):
            rv_1_out = Disc.from_rv(rv_1)
            rv_1_ref = Disc([1, 3], [0.5, 0.5])
            _test_equal_rand(rv_1_out, rv_1_ref)

        # Currently not all elements with small probabilities are not detected,
        # but only those, which cumulative probability after previous detected
        # element is less than threshold.
        rv_2 = distrs.rv_discrete(values=([1, 2, 3, 4], [0.5, 0.0625, 0.0625, 0.375]))
        with config.context({"small_prob": 0.1}):
            rv_2_out = Disc.from_rv(rv_2)
            rv_2_ref = Disc([1, 3, 4], [0.5, 0.125, 0.375])
            _test_equal_rand(rv_2_out, rv_2_ref)

    def test_last_value(self):
        # If last (biggest) value has probability equal to `small_prob`, it
        # should nevertheless be included
        rv = distrs.rv_discrete(values=([0, 1], [0.875, 0.125]))
        with config.context({"small_prob": 0.125}):
            rv_out = Disc.from_rv(rv)
            rv_ref = Disc([0, 1], [0.875, 0.125])
            _test_equal_rand(rv_out, rv_ref)

        # If last value cuts some right tail, all its probability should be
        # transferred into last x-value
        rv = distrs.rv_discrete(values=([0, 1, 2, 3], [0.625, 0.125, 0.125, 0.125]))
        with config.context({"small_prob": 0.2}):
            rv_out = Disc.from_rv(rv)
            rv_ref = Disc([0, 2], [0.625, 0.375])
            _test_equal_rand(rv_out, rv_ref)

    @staticmethod
    def is_from_rv_small_tails(rv):
        small_prob = config.small_prob

        rv_disc = Disc.from_rv(rv)
        x = rv_disc.x

        # Left tail
        left_tail_prob = rv.cdf(x[0] - 1e-13)

        # Right tail
        right_tail_prob = 1 - rv.cdf(x[-1])

        return (left_tail_prob <= small_prob) and (right_tail_prob <= small_prob)

max_stars_repo_name: echasnovski/randomvars

max_issues_repo_name: echasnovski/randomvars

max_forks_repo_name: echasnovski/randomvars

reason: import numpy,from numpy,import scipy


text:
# https://stable-baselines.readthedocs.io/en/master/guide/custom_env.html
import math
import copy
import json

import gym
from gym import spaces
import pandas as pd
import numpy as np
import shapely.geometry as geom
from analysis.obspy_copy import degrees2kilometers

from gym.envs.classic_control import rendering
from pyglet.window import key, mouse

import pyglet
from pyglet import shapes, window
from pyglet import gl

WINDOW_W = 800
WINDOW_H = 800

CAR_WIDTH = 1.680
CAR_LENGTH = 2.817
WHEEL_SPACING = 1.480
WHEEL_LENGTH = 0.3
WHEEL_WIDTH = 0.195
WHEEL_BASE = 2.345

CAR_COLOR = (1, 0, 0)  # RED
WINDOW_COLOR = (.5, .5, .8)  # BLUE
WHEEL_COLOR = (0, 0, 0)  # BLACK

ROAD_WIDTH = 7

ZOOM = 50

env = None
close = False

# environment with continuous observation space
class CarEnv(gym.Env):
    """
    Environment class containing all the methods necessary for reinforcement
    learning.

    Attributes
    ----------
        df : pd.DataFrame
            DataFrame containing info about the ref. trace.
        road : geom.LineString
            Geometric multiline representing the ref. trace.
        init_x : float
            Initial x position for reset.
        init_y : float
            Initial y position for reset.
        init_theta : float
            Initial vehicle heading angle for reset (radians).
        init_v : float
            Initial velocity of the vehicle for reset (m/s).
        init_psi : float
            Initial steering command of the vehicle given in radians.
        init_observation : np.array
            Initial observation used to reset the environment. It's described
            by array [x, y, theta, v, psi].
        state : np.array or None
            Current state of the car given by array [x, y, theta, v, psi].
        viewer : rendering.Viewer()
            Scene with objects being rendered
        track : list
            Coordinates x,y of all points of the track
        car_trans : rendering.Transformation()
            Attribute for making transformations with car object
        track_trans : rendering.Transformation()
            Attribute for making transformations with track
        left_wheel_trans, right_wheel_trans : rendering.Transformation()
            Attribute for making transformations with front wheels

    """

    metadata = {'render.modes': ['human']}

    def __init__(self, reference_track, type='continuous',
                 action_dim=2, verbose=1):

        super(CarEnv, self).__init__()
        self.df = self._read_df(reference_track)

        self._create_road()
        self.map_center = self._get_track_start()

        self._create_road_borders()

        self.init_x = 0
        self.init_y = 0
        self.init_theta = 0
        self.init_v = 2.0
        self.init_psi = 0.0
        self.time = 0
        self.const_t = False
        self.const_v = False

        self.track = None
        self.outer_border = None
        self.inner_border = None

        self.car_trans = None
        self.left_wheel_trans = None
        self.right_wheel_trans = None
        self.left_rearwheel_trans = None
        self.right_rearwheel_trans = None
        self.track_trans = None

        self.viewer = None

        self.zoom = 0
        self.is_zoomed = False
        self.car_rendered = False

    @staticmethod
    def _read_df(df) -> pd.DataFrame:
        # covert degrees to meters and degrees to radians
        df.loc[:,'LAT'] = df.loc[:,'LAT'].apply(lambda deg: degrees2kilometers(deg) * 1000)
        df.loc[:,'LON'] = df.loc[:,'LON'].apply(lambda deg: degrees2kilometers(deg) * 1000)

        return df

    def _create_road(self):
        """
        Creates road from given points
        """

        self.track_points = list()
        for index, row in self.df.iterrows():
            self.track_points.append((row['LON'], row['LAT']))

    def _create_road_borders(self):
        """
        Creates point forming borders (inner anf outer)
        of track with given road width
        """
        self.outer = list()
        self.inner = list()

        for i, point in enumerate(self.track_points):

            x_center, y_center = self.map_center
            dist = ROAD_WIDTH/2  # half_width of road

            AL = []  # transformed point A to the left
            AR = []  # transformed point A to the right
            BL = []  # transformed point B to the left
            BR = []  # transformed point B to the right

            A = [point[0], point[1]]  # start point X, Y
            if i == len(self.track_points) - 1:
                B = [self.track_points[0][0], self.track_points[0][1]]  # end point X, Y
            else:
                B = [self.track_points[i + 1][0], self.track_points[i + 1][1]]  # end point X, Y

            V = [B[0] - A[0], B[1] - A[1]]  # vector of points A and B

            # angle between vector V and X axis in coordinate system
            cos_alfa = math.acos(V[0] / math.sqrt(V[0]**2 + V[1]**2))

            # if angle is greater than 90 degrees, subtract 90 degerees
            if cos_alfa > (math.pi/2):
                beta = cos_alfa - (math.pi/2)
            else:
                beta = (math.pi/2) - cos_alfa

            delta_x = dist * math.cos(beta)  # difference on X axis
            delta_y = dist * math.sin(beta)  # difference on Y axis

            # all possible directions of vector V
            if V[0] > 0:  # direction to right
                if V[1] == 0:
                    AL = [A[0], A[1] + dist]
                    AR = [A[0], A[1] - dist]
                    BL = [B[0], B[1] + dist]
                    BR = [B[0], B[1] - dist]
                if V[1] > 0:  # direction right, up
                    AL = [A[0] - delta_x, A[1] + delta_y]
                    AR = [A[0] + delta_x, A[1] - delta_y]
                    BL = [B[0] - delta_x, B[1] + delta_y]
                    BR = [B[0] + delta_x, B[1] - delta_y]
                if V[1] < 0:  # direction right, down
                    AL = [A[0] + delta_x, A[1] + delta_y]
                    AR = [A[0] - delta_x, A[1] - delta_y]
                    BL = [B[0] + delta_x, B[1] + delta_y]
                    BR = [B[0] - delta_x, B[1] - delta_y]

            if V[0] < 0:  # direction to left
                if V[1] == 0:
                    AL = [A[0], A[1] - dist]
                    AR = [A[0], A[1] + dist]
                    BL = [B[0], B[1] - dist]
                    BR = [B[0], B[1] + dist]
                if V[1] > 0:  # direction left, up
                    AL = [A[0] - delta_x, A[1] - delta_y]
                    AR = [A[0] + delta_x, A[1] + delta_y]
                    BL = [B[0] - delta_x, B[1] - delta_y]
                    BR = [B[0] + delta_x, B[1] + delta_y]
                if V[1] < 0:  # direction left, down
                    AL = [A[0] + delta_x, A[1] - delta_y]
                    AR = [A[0] - delta_x, A[1] + delta_y]
                    BL = [B[0] + delta_x, B[1] - delta_y]
                    BR = [B[0] - delta_x, B[1] + delta_y]

            if V[1] > 0 and V[0] == 0:  # direction up
                AL = [A[0] - dist, A[1]]
                AR = [A[0] + dist, A[1]]
                BL = [B[0] - dist, B[1]]
                BR = [B[0] + dist, B[1]]
            if V[1] < 0 and V[0] == 0:  # direction down
                AL = [A[0] + dist, A[1]]
                AR = [A[0] - dist, A[1]]
                BL = [B[0] + dist, B[1]]
                BR = [B[0] - dist, B[1]]

            self.outer.append(((AL[0] - x_center), (AL[1] - y_center)))
            self.outer.append(((BL[0] - x_center), (BL[1] - y_center)))
            self.inner.append(((AR[0] - x_center), (AR[1] - y_center)))
            self.inner.append(((BR[0] - x_center), (BR[1] - y_center)))

    def _get_track_start(self):
        """
        Gets track centre based on starting point of the track
        """
        x = self.track_points[0][0]
        y = self.track_points[0][1]
        return ((x, y))

    def _create_track(self, points, mode):
        """
        Transforms world coordinates into x,y coordinates by finding center
        and translating points and creates track represented by x,y points

        :returns - list of x,y points representing track in
        """

        if mode == 'track':
            center_x, center_y = self.map_center
        else:
            center_x = 0
            center_y = 0

        # append all x,y points to the track list
        coordinates = []
        for point in points:
            x = point[0] - center_x
            y = point[1] - center_y
            coordinates.append((x, y))

        coordinates.append(((points[0][0] - center_x), (points[0][1] - center_y)))

        return coordinates

    def _create_track_object(self):
        """
        Creats track polygon from precalculated border coordinates
        with given color
        """

        """
        Connecting all points as individual polygons is too slow,
        so connecions are between every 5th point
        """
        for index in range(0,len(self.outer),5):
            if index < (len(self.outer) - 5):
                ll_x, ll_y = self.inner[index]
                ul_x, ul_y = self.inner[index + 5]
                lr_x, lr_y = self.outer[index]
                ur_x, ur_y = self.outer[index + 5]
            else:
                ll_x, ll_y = self.inner[index]
                ul_x, ul_y = self.inner[0]
                lr_x, lr_y = self.outer[index]
                ur_x, ur_y = self.outer[0]

            track = rendering.FilledPolygon([
                    (ll_x + WINDOW_W / 2, ll_y + WINDOW_H / 2),
                    (ul_x + WINDOW_W / 2, ul_y + WINDOW_H / 2),
                    (ur_x + WINDOW_W / 2, ur_y + WINDOW_H / 2),
                    (lr_x + WINDOW_W / 2, lr_y + WINDOW_H / 2)
            ])
            # creates track represented by line
            self.track_trans = rendering.Transform()
            track.add_attr(self.track_trans)
            track.add_attr(self.transform)
            track.set_color(0.25, 0.25, 0.25)
            self.viewer.add_geom(track)

    def _create_car_object(self, car_width, car_length, car_color):
        """
        Creates car object (polygon) which includes body and rear wheels
        """
        car_body = list()

        R, G, B = car_color

        left, right, top, bottom = (
            -car_width / 2,
            car_width / 2,
            car_length / 2,
            -car_length / 2
        )
        # CAR BODY
        car = rendering.FilledPolygon([
            (left, bottom),
            (left + car_width/10, top),
            (right - car_width/10, top),
            (right, bottom)
        ])
        self.car_trans = rendering.Transform()
        car.add_attr(self.car_trans)
        car.add_attr(self.transform)
        car.set_color(R, G, B)
        car_body.append(car)

        wheel_oversize = WHEEL_WIDTH - (CAR_WIDTH - WHEEL_SPACING)/2

        # REAR WHEELS
        left_rear_wheel = rendering.FilledPolygon([
            (left - wheel_oversize, bottom),
            (left - wheel_oversize, bottom + WHEEL_LENGTH),
            (left + wheel_oversize, bottom + WHEEL_LENGTH),
            (left + wheel_oversize, bottom)
        ])
        left_rear_wheel.add_attr(self.car_trans)
        left_rear_wheel.add_attr(self.transform)
        left_rear_wheel.set_color(.1, .1, .1)
        car_body.append(left_rear_wheel)

        right_rear_wheel = rendering.FilledPolygon([
            (right - wheel_oversize, bottom),
            (right - wheel_oversize, bottom + WHEEL_LENGTH),
            (right + wheel_oversize, bottom + WHEEL_LENGTH),
            (right + wheel_oversize, bottom)
        ])
        right_rear_wheel.add_attr(self.car_trans)
        right_rear_wheel.add_attr(self.transform)
        right_rear_wheel.set_color(.1, .1, .1)
        car_body.append(right_rear_wheel)

        return car_body

    def _define_label(self, text, offset):
        label = pyglet.text.Label(
            text,
            font_size=12,
            x=5,
            y=WINDOW_H - offset,
            color=(255, 255, 255, 255),
        )
        return label

    def _create_track_line(self, line, color):
        """
        Function to create middle line and borders of the track
        """
        R,G,B = color
        track_line = rendering.PolyLine(line, 1)
        track_line.add_attr(self.track_trans)
        track_line.add_attr(self.transform)
        track_line.set_color(R, G, B)
        return track_line

    def step(self, action):
        """
        Updates parameters with each step
        """

        new_x = action['x']
        new_y = action['y']

        # scp does not contain velocity and time, so we want to set it to constant
        if 'velocity' in action:
            if action['velocity'] == 'const':
                new_v = 1
                self.const_v = True
            else:
                new_v = action['velocity']
        else:
            new_v = 1

        if 'time' in action:
            new_t = action['time']
        else:
            self.const_t = True
            new_t = self.time

        new_theta = action['heading_angle']
        new_psi = action['steering_angle']

        self.state = [new_x, new_y, new_v, new_theta, new_psi, new_t]

        return self.state

    def reset(self):
        """
        Resets the state of the car to the initial state and creates track
        """
        self.state = np.array([
            self.init_x,
            self.init_y,
            self.init_v,
            self.init_theta,
            self.init_psi,
            self.time
        ])
        self.track = self._create_track(self.track_points, 'track')
        self.outer_border = self._create_track(self.outer, 'border')
        self.inner_border = self._create_track(self.inner, 'border')

    def render(self, mode='human'):

        x, y, v, theta, psi, time = self.state
        center_x, center_y = self.map_center

        self.time = self.time + 0.1 if self.const_t else time

        car_w = CAR_WIDTH * ZOOM
        car_l = CAR_LENGTH * ZOOM
        wheel_w = WHEEL_WIDTH * ZOOM
        wheel_l = WHEEL_LENGTH * ZOOM

        if self.viewer is None:
            x = 0
            y = 0

            self.viewer = rendering.Viewer(WINDOW_W, WINDOW_H)
            self.transform = rendering.Transform()

            # -----DEFINE LABELS-----
            self.time_label = self._define_label('time', 20)
            self.velocity_label = self._define_label('time', 40)
            self.steering_angle_label = self._define_label('time', 60)
            self.upper_left_indicator = shapes.Rectangle(0, WINDOW_H - 80, 200, 80, color=(0,0,0))

            # -----BACKGROUND GRASS-----
            background = rendering.FilledPolygon([
                (-WINDOW_W, -WINDOW_H),
                (-WINDOW_W, WINDOW_H * 2),
                (WINDOW_W * 2, WINDOW_H * 2),
                (WINDOW_W * 2, -WINDOW_H)
            ])
            self.background_trans = rendering.Transform()
            background.add_attr(self.background_trans)
            background.add_attr(self.transform)
            background.set_color(0.051, 0.455, 0.024)
            self.viewer.add_geom(background)

            # -----TRACK AND TRACK LINES RENDERING-----
            self._create_track_object()

            # centers points of the track to the center of the screen
            for i in range(len(self.track)):
                x, y = self.track[i]
                self.track[i] = x + WINDOW_W / 2, y + WINDOW_H / 2

            for i in range(len(self.inner_border)):
                x, y = self.inner_border[i]
                self.inner_border[i] = x + WINDOW_W / 2, y + WINDOW_H / 2

            for i in range(len(self.outer_border)):
                x, y = self.outer_border[i]
                self.outer_border[i] = x + WINDOW_W / 2, y + WINDOW_H / 2

            # creates track represented by line
            self.track_trans = rendering.Transform()
            track_line = self._create_track_line(self.track, color=(1, 1, 1))
            self.viewer.add_geom(track_line)
            track_inner = self._create_track_line(self.inner_border, color=(1, 0, 0))
            self.viewer.add_geom(track_inner)
            track_outer = self._create_track_line(self.outer_border, color=(1, 0, 0))
            self.viewer.add_geom(track_outer)

            # -----CAR OBJECT RENDERING-----
            if self.car_rendered == False:
                car_body = self._create_car_object(CAR_WIDTH, CAR_LENGTH, CAR_COLOR)

                for part in car_body:
                    self.viewer.add_geom(part)

                # FRONT WHEELS
                wheel = rendering.FilledPolygon([
                    (-WHEEL_WIDTH / 2, -WHEEL_LENGTH / 2),
                    (WHEEL_WIDTH / 2, -WHEEL_LENGTH / 2),
                    (WHEEL_WIDTH / 2, WHEEL_LENGTH / 2),
                    (-WHEEL_WIDTH / 2, WHEEL_LENGTH / 2)
                ])

                left_front_wheel = copy.deepcopy(wheel)
                self.left_wheel_trans = rendering.Transform()
                left_front_wheel.add_attr(self.left_wheel_trans)
                left_front_wheel.add_attr(self.car_trans)
                left_front_wheel.add_attr(self.transform)
                left_front_wheel.set_color(.1, .1, .1)
                self.viewer.add_geom(left_front_wheel)

                right_front_wheel = copy.deepcopy(wheel)
                self.right_wheel_trans = rendering.Transform()
                right_front_wheel.add_attr(self.right_wheel_trans)
                right_front_wheel.add_attr(self.car_trans)
                right_front_wheel.add_attr(self.transform)
                right_front_wheel.set_color(.1, .1, .1)
                self.viewer.add_geom(right_front_wheel)

                self.car_rendered = True

                # TODO - Make funcion from this
                # -----MINIATURE OF CAR IN LOWER LEFT INDICATOR-----

                lower_left_indicator = rendering.FilledPolygon([
                    (0, 0),
                    (0, 250),
                    (250, 250),
                    (250, 0)
                ])
                lower_left_indicator.set_color(0.784, 0.784, 0.784)
                self.viewer.add_geom(lower_left_indicator)

                self.compass_trans = rendering.Transform()
                self.compass_trans.set_translation(125, 125)
                compass = rendering.Image('animation/compass.png', 250, 250)
                compass.add_attr(self.compass_trans)
                self.viewer.add_geom(compass)

                left, right, top, bottom = (
                    -CAR_WIDTH * ZOOM / 2,
                    CAR_WIDTH * ZOOM / 2,
                    CAR_LENGTH * ZOOM / 2,
                    -CAR_LENGTH * ZOOM / 2
                )
                # CAR BODY

                car = rendering.FilledPolygon([
                    (left, bottom),
                    (left + CAR_WIDTH*ZOOM/10, top),
                    (right - CAR_WIDTH*ZOOM/10, top),
                    (right, bottom)
                ])
                self.mini_car_trans = rendering.Transform()
                self.mini_car_trans.set_translation(125,125)
                car.add_attr(self.mini_car_trans)
                car.set_color(1, 0, 0)
                self.viewer.add_geom(car)

                # WHEELS REAR
                wheel_oversize = WHEEL_WIDTH * ZOOM - (CAR_WIDTH * ZOOM - WHEEL_SPACING * ZOOM) / 2

                left_rear_wheel = rendering.FilledPolygon([
                    (left - wheel_oversize, bottom),
                    (left - wheel_oversize, bottom + WHEEL_LENGTH * ZOOM),
                    (left + wheel_oversize, bottom + WHEEL_LENGTH * ZOOM),
                    (left + wheel_oversize, bottom)
                ])

                left_rear_wheel.add_attr(self.mini_car_trans)
                left_rear_wheel.set_color(.1, .1, .1)
                self.viewer.add_geom(left_rear_wheel)

                right_rear_wheel = rendering.FilledPolygon([
                    (right - wheel_oversize, bottom),
                    (right - wheel_oversize, bottom + WHEEL_LENGTH * ZOOM),
                    (right + wheel_oversize, bottom + WHEEL_LENGTH * ZOOM),
                    (right + wheel_oversize, bottom)
                ])
                right_rear_wheel.add_attr(self.mini_car_trans)
                right_rear_wheel.set_color(.1, .1, .1)
                self.viewer.add_geom(right_rear_wheel)

                # WHEELS FRONT MINI
                wheel = rendering.FilledPolygon([
                    (-WHEEL_WIDTH*ZOOM / 2, -WHEEL_LENGTH*ZOOM / 2),
                    (WHEEL_WIDTH*ZOOM / 2, -WHEEL_LENGTH*ZOOM / 2),
                    (WHEEL_WIDTH*ZOOM / 2, WHEEL_LENGTH*ZOOM / 2),
                    (-WHEEL_WIDTH*ZOOM / 2, WHEEL_LENGTH*ZOOM / 2)
                ])

                left_front_wheel_mini = copy.deepcopy(wheel)
                self.mini_left_wheel_trans = rendering.Transform()
                left_front_wheel_mini.add_attr(self.mini_left_wheel_trans)
                left_front_wheel_mini.add_attr(self.mini_car_trans)
                left_front_wheel_mini.set_color(.1, .1, .1)
                self.viewer.add_geom(left_front_wheel_mini)

                right_front_wheel_mini = copy.deepcopy(wheel)
                self.mini_right_wheel_trans = rendering.Transform()
                right_front_wheel_mini.add_attr(self.mini_right_wheel_trans)
                right_front_wheel_mini.add_attr(self.mini_car_trans)
                right_front_wheel_mini.set_color(.1, .1, .1)
                self.viewer.add_geom(right_front_wheel_mini)

        # -----ZOOMING INTO CAR VIEW-----
        if self.zoom == ZOOM:
            self.is_zoomed = True
            # Zooms into car view until reached specific zoom
        if self.is_zoomed is not True:
            if self.time == 0:
                self.time = 0.05
            self.zoom = 0.01 * max(1 - self.time, 0) + (ZOOM * min(self.time, 1))

        # -----FOLLOWING CAR VIEW-----
        scroll_x = x
        scroll_y = y
        angle = -theta
        self.transform.set_scale(self.zoom, self.zoom)
        self.transform.set_translation(
            -WINDOW_W * (self.zoom / 2 - 0.5) - (scroll_x * self.zoom),
            -WINDOW_H * (self.zoom / 2 - 0.5) - (scroll_y * self.zoom)
        )

        # -----TRANSFORMATION AND ROTATION OF CAR ON TRACK-----

        # car rotation and translation
        self.car_trans.set_translation(x + WINDOW_W / 2, y + WINDOW_H / 2)
        self.car_trans.set_rotation(theta - (math.pi) / 2)

        # wheels rotation and translation
        wheel_oversize = WHEEL_WIDTH - (CAR_WIDTH - WHEEL_SPACING) / 2
        self.left_wheel_trans.set_translation(-WHEEL_SPACING / 2 - wheel_oversize, WHEEL_BASE / 2)
        self.right_wheel_trans.set_translation(WHEEL_SPACING / 2 + wheel_oversize, WHEEL_BASE / 2)
        self.left_wheel_trans.set_rotation(psi)
        self.right_wheel_trans.set_rotation(psi)

        # -----TRANSFORMATION AND ROTATION OF CAR IN LOWER LEFT INDICATOR-----
        wheel_oversize = WHEEL_WIDTH * ZOOM - (CAR_WIDTH * ZOOM - WHEEL_SPACING * ZOOM) / 2
        self.mini_car_trans.set_rotation(theta - (math.pi) / 2)
        self.mini_left_wheel_trans.set_translation(-WHEEL_SPACING * ZOOM / 2 - wheel_oversize, WHEEL_BASE * ZOOM / 2)
        self.mini_right_wheel_trans.set_translation(WHEEL_SPACING * ZOOM / 2 + wheel_oversize, WHEEL_BASE * ZOOM / 2)
        self.mini_left_wheel_trans.set_rotation(psi)
        self.mini_right_wheel_trans.set_rotation(psi)

        # -----RENDERING-----
        win = self.viewer.window
        win.switch_to()
        win.dispatch_events()
        win.clear()

        for geom in self.viewer.geoms:
            geom.render()

        # ----------LABELS-----------
        self.upper_left_indicator.draw()
        self.time_label.text = f'time: {round(time, 1)}'
        self.time_label.draw()

        self.velocity_label.text = 'velocity: constant' if self.const_v else f'velocity: {v:.2f} km/h'
        self.velocity_label.draw()

        self.steering_angle_label.text = f'steering_angle: {math.degrees(psi):.2f} deg'
        self.steering_angle_label.draw()

        win.flip()
        self.viewer.isopen

    def zoom_in(self):
        self.zoom /= 0.8
        return

    def zoom_out(self):
        self.zoom *= 0.8
        return

    def close(self):
        """
        Closes the viewer window.
        """
        if self.viewer is not None:
            self.viewer.close()
            self.viewer = None

def key_press(k, mod):
    global close
    if k == key.ESCAPE or k == key.X:
        close = True
        print('Animation closed by user.')

def mouse_scroll(x, y, scroll_x, scroll_y):
    print(env.zoom)
    if scroll_y > 0 and env.zoom > 3:
        env.zoom_out()
    if scroll_y < 0 and env.zoom < 80:
        env.zoom_in()

def on_close_window():
    global close
    close = True

def run_animation(data, car_dimensions, coords):
    """
    Runs complete animation
    Parameters
    ---------
        env : CarEnv
            A car environment
        close : Boolean
            True if animation is finished or canceled
    """
    global env, close, CAR_WIDTH, CAR_LENGTH, WHEEL_WIDTH, WHEEL_LENGTH, WHEEL_BASE, WHEEL_SPACING

    if car_dimensions:
        CAR_LENGTH, CAR_WIDTH, WHEEL_SPACING, WHEEL_BASE, WHEEL_LENGTH = car_dimensions

    env = CarEnv(reference_track=coords)
    env.reset()
    env.render()

    env.viewer.window.on_mouse_scroll = mouse_scroll
    env.viewer.window.on_key_press = key_press
    env.viewer.window.on_close = on_close_window

    for index, row in data.iterrows():
        if close is True:
            close = False
            break
        observation = env.step(row)
        env.render()
    env.viewer.window.set_visible(False)
    env.close()
    print('Animation ended.')
    return

if __name__ == '__main__':

    df = pd.read_csv('animation/track_ref.csv')
    data = pd.read_csv('animation/data.csv')

    print(df)

    env = CarEnv(reference_track=df)
    env.reset()
    env.render()

    env.viewer.window.on_key_press = key_press
    env.viewer.window.on_mouse_scroll = mouse_scroll

    for index, row in data.iterrows():
        if close is True:
            close = False
            break
        observation = env.step(row)
        env.render()
    env.viewer.window.set_visible(False)
    env.close()


text:
# In settings.json first activate computer vision mode:
# https://github.com/Microsoft/AirSim/blob/master/docs/image_apis.md#computer-vision-mode

# import setup_path
import airsim

import pprint
import tempfile
import os
import sys
import time

import numpy as np
import cv2

from noise_model import *
DELAY = 0.01 # must be > 0

# RegEx for the object IDs of the objects of interest (that is, the objects
# around which you want the bounding boxes).
REGEX_OBJECTS_OF_INTEREST = [
    "Landscape[\w]*",
    "prp_trafficLight_Blueprint[\w]*",
    "prp_manhole[\w]*",
    "prp_light[\w]*",
    "prp_garbageCan[\w]*",
    "prp_bench[\w]*",
    "prp_fireHydrant[\w]*",
    "prp_chair[\w]*",
    "prp_table[\w]*",
    "CameraActor[\w]*",
    "prp_sewerGrate[\w]*",
    "prp_tableUmbrella[\w]*",
    "prp_electricalBox[\w]*",
    "Awning_mdl[\w]*",
    "PowerLine[\w]*",
    "prp_crosswalkSign[\w]*",
    "Building[\w]*",
    "prp_metalPillar[\w]*",
    "flg_tree[\w]*",

    "prp_chainlinkFence[\w]*",
    "prp_tarp[\w]*",
    "prp_vent[\w]*",
    "door_A[\w]*",
    "prp_ac[\w]*",
    "StoreSign[\w]*",
    "prp_pylong_Sml[\w]*",
    "flg_grass[\w]*",
    "prp_parkingMeter[\w]*",
    "Flg_hedge_Short[\w]*",
    "prp_metalFence[\w]*",
    "prp_barrier_Lrg[\w]*",
    "TrashBag[\w]*",
    "prp_potSquare[\w]*",
    "prp_cementBarrier[\w]*",
    "flg_fern[\w]*",
    "prp_flag[\w]*"
]

# Classes corresponding to the objects in the REGEX_OBJECTS_OF_INTEREST list above.
OBJECT_OF_INTEREST_CLASSES = [
    "ground",
    "trafficLight",
    "manhole",
    "fixLight",
    "garbageCan",
    "bench",
    "fireHydrant",
    "chair",
    "table",
    "CameraActor",
    "sewerGrate",
    "tableUmbrella",
    "electricalBox",
    "Awning",
    "PowerLine",
    "crosswalkSign",
    "Building",
    "metalPillar",
    "tree",

    "chainlinkFence",
    "tarp",
    "vent",
    "door",
    "ac",
    "StoreSign",
    "pylong_Sml",
    "grass",
    "parkingMeter",
    "hedge_Short",
    "metalFence",
    "barrier_Lrg",
    "TrashBag",
    "potSquare",
    "cementBarrier",
    "fern",
    "flag"
]

# An exhaustive list of all classes
# Note that this is separate to the above list, as we may want classes
# which are not currently exhibited in OBJECT_OF_INTEREST_CLASSES.
# CLASSES = ['mountainbike',
#            'car',
#            'truck',
#            'dog',
#            'horse',
#            'sheep',
#            'giraffe']

# Confusion matrix for object misclassification emulation.
# See https://en.wikipedia.org/wiki/Confusion_matrix for more information
CONFUSION_MATRIX = np.array(
        [[0.85714286,0.,0.14285714,0.,0.,0.,0.],
        [0.06666667,0.8,0.06666667,0.,0.,0.06666667,0.],
        [0.,0.,1.,0.,0.,0.,0.],
        [0.03636364,0.,0.01818182,0.94545455,0.,0.,0.],
        [0.,0.,0.,0.,1.,0.,0.],
        [0.,0.,0.,0.,0.,1.,0.],
        [0.,0.,0.,0.01408451,0.,0.,0.98591549]]).T

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

MESH_COLS = np.array([
          [57 ,181, 55],
          [6  ,108,153],
          [191,105,112],
          [72 ,121, 89],
          [64 ,225,190],
          [59 ,190,206],
          [36 ,13 , 81],
          [195,176,115],
          [27 ,171,161],
          [180,169,135],
          [199,26 ,29],
          [239,16 ,102],
          [146,107,242],
          [23 ,198,156],
          [160,89 ,49],
          [116,218,68],
        ])

# deal with swapped RB values
# cp = MESH_COLS.copy()
# MESH_COLS[:,0] = cp[:,2]
# MESH_COLS[:,2] = cp[:,0]

client = None

def setup():
    # Clear background
    msg = 'Setting everything to ID 255 (to clear unwanted objects)...'
    print(msg, end='')
    found = client.simSetSegmentationObjectID("[\w]*", 255, True);
    print(' ' * (65 - len(msg)) + ('[SUCCESS]' if found else '[FAILED!]'))

    # Set objects of interest
    for key, val in zip(REGEX_OBJECTS_OF_INTEREST, range(len(REGEX_OBJECTS_OF_INTEREST))):
        msg = 'Setting %s to ID %d...' % (key, val)
        print(msg, end='')
        found = client.simSetSegmentationObjectID(key, val, True);
        print(' ' * (40 - len(msg)) + ('[SUCCESS]' if found else '[FAILED!]'))

if __name__ == '__main__':
    file_path = "/Users/aqiu/Documents/AirSim/2022-03-07-02-02-06/airsim_rec.txt" # sys.argv[1]
    cur_dir = file_path[:file_path.rfind(os.path.sep)] + os.path.sep
    print(cur_dir)

    pp = pprint.PrettyPrinter(indent=4)
    client = airsim.VehicleClient()
    client.confirmConnection()

    setup()

    # print_pose(client)

    # airsim.wait_key('Press any key to get camera parameters')
    for camera_id in range(2):
        camera_info = client.simGetCameraInfo(str(camera_id))
        print("CameraInfo %d: %s" % (camera_id, pp.pprint(camera_info)))
    tmp_dir = os.path.join(cur_dir, "airsim_drone")
    print ("Saving images to %s" % tmp_dir)
    try:
        for n in range(4):
            os.makedirs(os.path.join(tmp_dir, str(n)))
    except OSError:
        if not os.path.isdir(tmp_dir):
            raise

    fin = open(file_path, "r")
    line = fin.readline().strip()
    line = fin.readline().strip()
    while (line):
        parts = line.split("\t")
        timestamp = parts[1] # ms
        # timestamp = float(parts[1]) / 1000.0 # s

        pos_x = float(parts[2])
        pos_y =float( parts[3])
        pos_z = float(parts[4])

        quat_w = float(parts[5])
        quat_x = float(parts[6])
        quat_y = float(parts[7])
        quat_z = float(parts[8])

        client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(pos_x, pos_y, pos_z), airsim.Quaternionr(quat_x,quat_y,quat_z,quat_w)), True)
        time.sleep(0.1)

        responses = client.simGetImages([
        airsim.ImageRequest("0", airsim.ImageType.Scene, False, False),
        airsim.ImageRequest("0", airsim.ImageType.DepthPerspective, True, False),
        airsim.ImageRequest("0", airsim.ImageType.Infrared,False, False),
        airsim.ImageRequest("0", airsim.ImageType.Segmentation,False, False)]
        )

        for i, response in enumerate(responses):
            if response.pixels_as_float:
                print("Type %d, size %d, pos %s" % (response.image_type, len(response.image_data_float), pprint.pformat(response.camera_position)))
                airsim.write_pfm(os.path.normpath(os.path.join(tmp_dir,  str(i), str(timestamp) + "_" + str(i) + '.pfm')), airsim.get_pfm_array(response))
            else:
                print("Type %d, size %d, pos %s" % (response.image_type, len(response.image_data_uint8), pprint.pformat(response.camera_position)))
                airsim.write_file(os.path.normpath(os.path.join(tmp_dir, str(i), str(timestamp) + "_" + str(i) + '.png')), response.image_data_uint8)

                if  i == 3:
                    im = np.fromstring(response.image_data_uint8, dtype=np.uint8) # get numpy array
                    im = im.reshape(response.height, response.width, 3)           # reshape array to 4 channel image
                    # im = np.flipud(im)                                            # original image is flipped vertically

                    # find unique colors and draw bounding boxes
                    colours = np.unique(im.reshape((-1,3)), axis=0)
                    print("",i, " ", len(colours))

                    # store the BBs and their classes
                    bbs = []
                    classes = []
                    for col in colours:
                        colours_of_interest = np.sum(np.all(MESH_COLS == col, axis=-1))

                        # ignore if this colour does not correspond to an object of interest.
                        if colours_of_interest == 0:
                            continue
                        elif colours_of_interest > 1:
                            print("[WARNING] Multiple objects have the same color in segmented view! Using lowest index...")

                        index = np.where(np.all(MESH_COLS == col, axis=-1))[0][0]
                        objClass = OBJECT_OF_INTEREST_CLASSES[index]

                        mask = np.all(im == col, axis=-1)
                        locs = np.array(np.where(mask))

                        # find the BB
                        min_x = np.min(locs[0,:])
                        max_x = np.max(locs[0,:])
                        min_y = np.min(locs[1,:])
                        max_y = np.max(locs[1,:])
                        bbs.append((min_x, max_x, min_y, max_y))
                        classes.append(objClass)

                    bbs_clean = np.array(bbs).copy()

                    #################################################
                    ##### Add noise to the BBs
                    #################################################
                    # # first do some mis-classification
                    # classes = misclassify(classes, CLASSES, CONFUSION_MATRIX)
                    # # now add some error to the BBs
                    # bbs = add_jitter(bbs, shape=im.shape, length_scale_fraction=0.05, center_error_fraction=0.05)
                    # bbs, classes = introduce_false_negatives(bbs, classes, p=0.01, min_size=4)
                    # bbs, classes = introduce_false_positives(bbs, CLASSES, classes, shape=im.shape, p=0.01)
                    # bbs, classes = merge_close_bbs(bbs, classes, area_similarity_factor=1.7, overlap_factor=0.5)
                    #################################################

                    quiet = False
                    if not quiet:
                        # Draw the images
                        # boxedTrue = draw_bbs_on_image(trueIm, bbs)
                        boxedSeg =  draw_bbs_on_image(im, bbs)

                        cv2.imshow('infrared:2, seg:3, i:'+str(i), boxedSeg)
                        cv2.waitKey(int(DELAY * 1000))
                        cv2.destroyAllWindows()

                        # display the images
                        # cv2.imshow('Scene + BBs', swapRB(np.flipud(boxedTrue)))
                        # cv2.imshow('Segmented + BBs', swapRB(np.flipud(boxedSeg)))
                        # cv2.waitKey(int(DELAY * 1000))

        pose = client.simGetVehiclePose()
        pp.pprint(pose)

        # print_pose(client)

        time.sleep(3)

        line = fin.readline().strip()

    # pp = pprint.PrettyPrinter(indent=4)

    # client = airsim.VehicleClient()

    # airsim.wait_key('Press any key to get camera parameters')
    # for camera_id in range(2):
    # for camera_id in range(2):
    #     camera_info = client.simGetCameraInfo(str(camera_id))
    #     print("CameraInfo %d: %s" % (camera_id, pp.pprint(camera_info)))

    # airsim.wait_key('Press any key to get images')
    # tmp_dir = os.path.join(tempfile.gettempdir(), "airsim_drone")
    # print ("Saving images to %s" % tmp_dir)
    # try:
    #     for n in range(3):
    #         os.makedirs(os.path.join(tmp_dir, str(n)))
    # except OSError:
    #     if not os.path.isdir(tmp_dir):
    #         raise

    # for x in range(50): # do few times
    #     #xn = 1 + x*5  # some random number
    #     client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(x, 0, -2), airsim.to_quaternion(0, 0, 0)), True)
    #     time.sleep(0.1)

    #     responses = client.simGetImages([
    #         airsim.ImageRequest("0", airsim.ImageType.Scene),
    #         airsim.ImageRequest("1", airsim.ImageType.Scene),
    #         airsim.ImageRequest("2", airsim.ImageType.Scene)])

    #     for i, response in enumerate(responses):
    #         if response.pixels_as_float:
    #             print("Type %d, size %d, pos %s" % (response.image_type, len(response.image_data_float), pprint.pformat(response.camera_position)))
    #             airsim.write_pfm(os.path.normpath(os.path.join(tmp_dir, str(x) + "_" + str(i) + '.pfm')), airsim.get_pfm_array(response))
    #         else:
    #             print("Type %d, size %d, pos %s" % (response.image_type, len(response.image_data_uint8), pprint.pformat(response.camera_position)))
    #             airsim.write_file(os.path.normpath(os.path.join(tmp_dir, str(i), str(x) + "_" + str(i) + '.png')), response.image_data_uint8)

    #     pose = client.simGetVehiclePose()
    #     pp.pprint(pose)

    #     time.sleep(3)

    # currently reset() doesn't work in CV mode. Below is the workaround
    client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(0, 0, 0), airsim.to_quaternion(0, 0, 0)), True)

    # import pprint
    # from decimal import *
    # def print_pose(client):
    #     # 1 vehicle pose
    #     pose = client.simGetVehiclePose()
    #     print("vehicle pose: x={}, y={}, z={}".format(pose.position.x_val, pose.position.y_val, pose.position.z_val))
    #     angles = airsim.to_eularian_angles(pose.orientation)
    #     print("pitch={}, roll={}, yaw={}".format(angles[0], angles[1], angles[2]))
        # 2 multirotor pose
        # state = client.getMultirotorState()
        # # s = pprint.pformat(state)
        # # print("state: %s" % s)
        # print("multirotor pose: x={}, y={}, z={}".format(state.kinematics_estimated.position.x_val, state.kinematics_estimated.position.y_val, state.kinematics_estimated.position.z_val))
        # angles = airsim.to_eularian_angles(state.kinematics_estimated.orientation)
        # print("pitch={}, roll={}, yaw={}".format(angles[0], angles[1], angles[2]))
        # kinematics = client.simGetGroundTruthKinematics()
        # environment = client.simGetGroundTruthEnvironment()
        # print("Kinematics: %s\nEnvironemt %s" % (
        #     pprint.pformat(kinematics), pprint.pformat(environment)))
        # 3 imu pose
        # 4 camera pose
        # img_position = client.simGetImages([airsim.ImageRequest(0, airsim.ImageType.Scene)])[0].camera_position
        # img_orientation = client.simGetImages([airsim.ImageRequest(0, airsim.ImageType.Scene)])[0].camera_orientation
        # print("camera pose: x={}, y={}, z={}".format(img_position.x_val, img_position.y_val, img_position.z_val))
        # angles = airsim.to_eularian_angles(img_orientation)
        # print("pitch={}, roll={}, yaw={}".format(angles[0], angles[1], angles[2]))

max_stars_repo_name: aqiugroup/AirSim

max_issues_repo_name: aqiugroup/AirSim

max_forks_repo_name: aqiugroup/AirSim


text:
import numpy as np
import cv2
from keras.layers import Input
from keras.models import Model
from keras.models import load_model

decoder = load_model('roses_decoder.h5')
perceptron = load_model('decoder-perceptron.h5')

path = 'dataset/rose'
id=25 # sample code
param0 = np.loadtxt(path+'{:04d}.txt'.format(id))
id=26 # sample code
param1 = np.loadtxt(path+'{:04d}.txt'.format(id))
id=2 # sample code
param2 = np.loadtxt(path+'{:04d}.txt'.format(id))
id=235 # sample code
param3 = np.loadtxt(path+'{:04d}.txt'.format(id))
param = np.copy(param0)

last_value = -1

cv2.namedWindow("generator")
cv2.imshow('generator',np.zeros((112*5,112*5),np.uint8))

value = int(param[7])

def update_value( *args ):
    global value
    print(args[0])
    value = float(args[0])

cv2.createTrackbar("value", "generator", value, 90, update_value)

id = 0
while True:

    if last_value != value:
        last_value = value
        param[7] = value
        coded = perceptron.predict(param.reshape(1,-1))
        decoded = decoder.predict(coded)
        decoded = np.asarray(decoded[0]*255,np.uint8)
        decoded = cv2.resize(decoded,(112*5,112*5))
        cv2.imshow('generator',decoded)

    key = cv2.waitKey(10)
    if key == 27:
        break
    elif key == ord('r') or key == ord('0'):
        param = np.copy(param0)
        update_value(int(param[7]))
    elif key == ord('1'):
        param = np.copy(param1)
        update_value(int(param[7]))
    elif key == ord('2'):
        param = np.copy(param2)
        update_value(int(param[7]))
    elif key == ord('3'):
        param = np.copy(param3)
        update_value(int(param[7]))
    elif key == ord('s'):
        cv2.imwrite('generator/final'+str(id)+'.png',decoded)
        np.savetxt('generator/final'+str(id)+'.txt',param)
        id += 1

cv2.destroyAllWindows()


text:
import os, sys
from os.path import join
ROOT_DIR = os.path.abspath(os.curdir)
if ROOT_DIR not in sys.path:
    sys.path.append(join(ROOT_DIR,'opmatch'))
from util import utils
import numpy as np
from tests import create_test_data
d1 = {'a': [1,3]}
d2 = {'a':2, 'b':2}
d = utils.combine_dicts(d1,d2)
print(d)
print('first func tested')
df = create_test_data.get_test_data(True, 30, 3, .3)

exp_nexp_dic = {0:[1,2], 3:[4,5], 6:[7,8] }
avg_dist = utils.compute_avg_dist(df, exp_nexp_dic)
print(avg_dist)

max_stars_repo_name: kirilklein/opmatch

max_issues_repo_name: kirilklein/opmatch

max_forks_repo_name: kirilklein/opmatch


text:
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import pytest

from mmocr.utils import (bezier_to_polygon, is_on_same_line, sort_points,
                         stitch_boxes_into_lines)

def test_box_on_line():
    # regular boxes
    box1 = [0, 0, 1, 0, 1, 1, 0, 1]
    box2 = [2, 0.5, 3, 0.5, 3, 1.5, 2, 1.5]
    box3 = [4, 0.8, 5, 0.8, 5, 1.8, 4, 1.8]
    assert is_on_same_line(box1, box2, 0.5)
    assert not is_on_same_line(box1, box3, 0.5)

    # irregular box4
    box4 = [0, 0, 1, 1, 1, 2, 0, 1]
    box5 = [2, 1.5, 3, 1.5, 3, 2.5, 2, 2.5]
    box6 = [2, 1.6, 3, 1.6, 3, 2.6, 2, 2.6]
    assert is_on_same_line(box4, box5, 0.5)
    assert not is_on_same_line(box4, box6, 0.5)

def test_stitch_boxes_into_lines():
    boxes = [  # regular boxes
        [0, 0, 1, 0, 1, 1, 0, 1],
        [2, 0.5, 3, 0.5, 3, 1.5, 2, 1.5],
        [3, 1.2, 4, 1.2, 4, 2.2, 3, 2.2],
        [5, 0.5, 6, 0.5, 6, 1.5, 5, 1.5],
        # irregular box
        [6, 1.5, 7, 1.25, 7, 1.75, 6, 1.75]
    ]
    raw_input = [{'box': boxes[i], 'text': str(i)} for i in range(len(boxes))]
    result = stitch_boxes_into_lines(raw_input, 1, 0.5)
    # Final lines: [0, 1], [2], [3, 4]
    # box 0, 1, 3, 4 are on the same line but box 3 is 2 pixels away from box 1
    # box 3 and 4 are on the same line since the length of overlapping part >=
    # 0.5 * the y-axis length of box 5
    expected_result = [{
        'box': [0, 0, 3, 0, 3, 1.5, 0, 1.5],
        'text': '0 1'
    }, {
        'box': [3, 1.2, 4, 1.2, 4, 2.2, 3, 2.2],
        'text': '2'
    }, {
        'box': [5, 0.5, 7, 0.5, 7, 1.75, 5, 1.75],
        'text': '3 4'
    }]
    result.sort(key=lambda x: x['box'][0])
    expected_result.sort(key=lambda x: x['box'][0])
    assert result == expected_result

def test_bezier_to_polygon():
    bezier_points = [
        37.0, 249.0, 72.5, 229.55, 95.34, 220.65, 134.0, 216.0, 132.0, 233.0,
        82.11, 240.2, 72.46, 247.16, 38.0, 263.0
    ]
    pts = bezier_to_polygon(bezier_points)
    target = np.array([[37.0, 249.0], [42.50420761043885, 246.01570199737577],
                       [47.82291296107305, 243.2012392477038],
                       [52.98102930456334, 240.5511007435486],
                       [58.00346989357049, 238.05977547747486],
                       [62.91514798075522, 235.721752442047],
                       [67.74097681877824, 233.53152062982943],
                       [72.50586966030032, 231.48356903338674],
                       [77.23473975798221, 229.57238664528356],
                       [81.95250036448464, 227.79246245808432],
                       [86.68406473246829, 226.13828546435346],
                       [91.45434611459396, 224.60434465665548],
                       [96.28825776352238, 223.18512902755504],
                       [101.21071293191426, 221.87512756961655],
                       [106.24662487243039, 220.6688292754046],
                       [111.42090683773145, 219.5607231374836],
                       [116.75847208047819, 218.5452981484181],
                       [122.28423385333137, 217.6170433007727],
                       [128.02310540895172, 216.77044758711182],
                       [134.0, 216.0], [132.0, 233.0],
                       [124.4475521213005, 234.13617728531858],
                       [117.50700976818779, 235.2763434903047],
                       [111.12146960198277, 236.42847645429362],
                       [105.2340282840064, 237.6005540166205],
                       [99.78778247557953, 238.80055401662054],
                       [94.72582883802303, 240.0364542936288],
                       [89.99126403265781, 241.31623268698053],
                       [85.52718472080478, 242.64786703601104],
                       [81.27668756378483, 244.03933518005545],
                       [77.1828692229188, 245.49861495844874],
                       [73.18882635952762, 247.0336842105263],
                       [69.23765563493221, 248.65252077562326],
                       [65.27245371045342, 250.3631024930748],
                       [61.23631724741216, 252.17340720221605],
                       [57.07234290712931, 254.09141274238226],
                       [52.723627350925796, 256.12509695290856],
                       [48.13326724012247, 258.2824376731302],
                       [43.24435923604024, 260.5714127423822], [38.0, 263.0]])
    assert np.allclose(pts, target)

    bezier_points = [0, 0, 0, 1, 0, 2, 0, 3, 1, 0, 1, 1, 1, 2, 1, 3]
    pts = bezier_to_polygon(bezier_points, num_sample=3)
    target = np.array([[0, 0], [0, 1.5], [0, 3], [1, 0], [1, 1.5], [1, 3]])
    assert np.allclose(pts, target)

    with pytest.raises(AssertionError):
        bezier_to_polygon(bezier_points, num_sample=-1)

    bezier_points = [0, 1]
    with pytest.raises(AssertionError):
        bezier_to_polygon(bezier_points)

def test_sort_points():
    points = np.array([[1, 1], [0, 0], [1, -1], [2, -2], [0, 2], [1, 1],
                       [0, 1], [-1, 1], [-1, -1]])
    target = np.array([[-1, -1], [0, 0], [-1, 1], [0, 1], [0, 2], [1, 1],
                       [1, 1], [2, -2], [1, -1]])
    assert np.allclose(target, sort_points(points))

    points = np.array([[1, 1], [1, -1], [-1, 1], [-1, -1]])
    target = np.array([[-1, -1], [-1, 1], [1, 1], [1, -1]])
    assert np.allclose(target, sort_points(points))

    points = [[1, 1], [1, -1], [-1, 1], [-1, -1]]
    assert np.allclose(target, sort_points(points))

    with pytest.raises(AssertionError):
        sort_points([1, 2])

max_stars_repo_name: yangrisheng/mmocr

max_issues_repo_name: yangrisheng/mmocr

max_forks_repo_name: yangrisheng/mmocr


text:
import numpy as np
import matplotlib.pyplot as plt
import math
import cv2
import time

infile = open('points_inliers_r.pcd','rt')
data = infile.readlines()
data = data[11:]
data = [k.split(" ") for k in data]
for i in range(0,len(data)):
    data[i] = [float(k) for k in data[i]]

realpoints = np.zeros((len(data),3))

for i in range(0,len(data)):
    realpoints[i,0] = data[i][0]
    realpoints[i,1] = data[i][1]
    realpoints[i,2] = data[i][2]

fig2 = plt.figure(figsize=(10, 5))
bx = fig2.add_subplot(111)
bx.scatter(realpoints[:,0],realpoints[:,1])
bx.scatter(0,0,color = 'red')
bx.scatter(1,1,color = 'green')
plt.savefig("maptest2.png")

f = open("wab.txt","w+")
f.write("The Co-ordinates of features x, y, z , feature\n")
f.write("o:cup obstacle, b: ball, j:junk noise\n")

flag = 0
img = cv2.imread('maptest2.png')
for i in range(0,img.shape[0]):
	for j in range(0,img.shape[1]):
		if(img[i][j][2] == 255 and img[i][j][0] == 0 and img[i][j][1] == 0):
			origin = (j,i)
			flag = 1
			break
	if(flag==1):
		break

flag = 0
for i in range(0,img.shape[0]):
	for j in range(0,img.shape[1]):
		if(img[i][j][2] == 0 and img[i][j][0] == 0 and img[i][j][1] > 0):
			oneone = (j,i)
			flag = 1
			break
	if(flag==1):
		break

dydx = (oneone[0] - origin[0],oneone[1] - origin[1])
# print(img.shape)

img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
i1 = cv2.medianBlur(img,5)
i2 = cv2.GaussianBlur(img,(3,3),0)

'''
# global thresholding
ret1,th1 = cv2.threshold(img,127,255,cv2.THRESH_BINARY)
ret2,th2 = cv2.threshold(i1,127,255,cv2.THRESH_BINARY)
ret3,th3 = cv2.threshold(i2,127,255,cv2.THRESH_BINARY)
mega = np.c_[img,th1,th2,th3]
'''

'''
# adaptive mean thresholding
th1 = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,11,2)
th2 = cv2.adaptiveThreshold(i1,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,11,2)
th3 = cv2.adaptiveThreshold(i2,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,11,2)
mega = np.c_[img,th1,th2,th3]
'''

# adaptive gaussian thresholding
th1 = cv2.adaptiveThreshold(img,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,11,2)
th2 = cv2.adaptiveThreshold(i1,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,11,2)
th3 = cv2.adaptiveThreshold(i2,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,11,2)
mega = np.c_[img,th1,th2,th3]
'''
#Otsu
ret1,th1 = cv2.threshold(img,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
ret2,th2 = cv2.threshold(i1,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
ret3,th3 = cv2.threshold(i2,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
mega = np.c_[img,th1,th2,th3]
'''

#Contours
image, contours, hierarchy = cv2.findContours(th1,1,2)

centrad = np.zeros((len(contours),3))
actpoints = np.zeros((len(contours),3))

th2c = th1
j = 0
for i in range(0,len(contours)):
	(x,y),radius = cv2.minEnclosingCircle(contours[i])
	if(int(radius)>=8 and int(radius)<= 40):

		centrad[j,0] = int(x)
		centrad[j,1] = int(y)
		centrad[j,2] = int(radius)

		if(j>=1):
			flag = 0
			for k in range(0,j-1):
				if(abs(centrad[j,0] - centrad[k,0]) <= 2 and abs(centrad[j,1]- centrad[k,1]) <= 3):
					flag = 1
					break

			# print(flag)
			if(flag==0):

				actpoints[j,0] = float("{0:.2f}".format((x-origin[0])/dydx[0]))
				actpoints[j,1] = float("{0:.2f}".format((y-origin[1])/dydx[1]))
				th2c = cv2.circle(th2c,(int(x),int(y)),int(radius),(0,255,0),1)
				j = j+1

		if(j==0):
			actpoints[j,0] = float("{0:.2f}".format((x-origin[0])/dydx[0]))
			actpoints[j,1] = float("{0:.2f}".format((y-origin[1])/dydx[1]))
			th2c = cv2.circle(th2c,(int(x),int(y)),int(radius),(0,255,0),1)
			j = j+1

centrad = centrad[0:j]
actpoints = actpoints[0:j]

for i in range(0,len(actpoints)):
	l = []
	rad = centrad[i,2]
	for j in range(0,len(realpoints)):
		resx = float("{0:.2f}".format(abs((rad-origin[0])/dydx[0])))
		resy = float("{0:.2f}".format(abs((rad-origin[1])/dydx[0])))

		if (abs(actpoints[i,0] - realpoints[j,0]) <= 0.05 and abs(actpoints[i,1] - realpoints[j,1]) <= 0.05):
			l.append(realpoints[j,2])

	actpoints[i,2] = max(l)

	if(abs(actpoints[i,2])>=0.12):
		c = 'c'
	else:
		if(abs(actpoints[i,2])>=0.025 and abs(actpoints[i,2])<=0.06):
			c = 'b'
		else:
			c = 'j'

	f.write("%f %f %f %c\n" %(actpoints[i,0],actpoints[i,1],actpoints[i,2],c))

	'''if( resx<= 0.02  and  resy<= 0.02):
		actpoints[i,2] = realpoints[j,2]
		break
	'''

	'''
	if( resx<= resx  and  resy<= resy):
		actpoints[i,2] = realpoints[j,2]
		print(actpoints[i,:],i)
		break
	'''

f.close()

print("Action Points sampled:")
print(actpoints)
print("Begining Path Planning")
time.sleep(1)
cv2.imwrite("thmp3.png",th2c)
# cv2.waitKey(0)
cv2.destroyAllWindows()


text:
from policyopt.wrapper import *

import cv2
import gym
import numpy as np
from ppaquette_gym_doom.wrappers.action_space import ToDiscrete

from collections import deque

class ProcessFrame84x84Doom(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessFrame84x84Doom, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 3))

    def _observation(self, obs):
        return ProcessFrame84x84Doom.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 3:
            img = np.reshape(frame, [120, 160, 3]).astype(np.float32)
        else:
            assert False, "Unknown resolution."
        #img = img[:, :, 0] * 0.299 + img[:, :, 1] * 0.587 + img[:, :, 2] * 0.114
        #resized_screen = cv2.resize(img, (84, 84), interpolation=cv2.INTER_AREA)
        x_t = cv2.resize(img, (84, 84), interpolation=cv2.INTER_LINEAR)
        #x_t = resized_screen[18:102, :]
        #x_t = np.reshape(x_t, [84, 84, 1])
        assert x_t.shape == (84, 84, 3)
        #x_t = np.transpose(x_t, axes=[2, 0, 1])
        #assert x_t.shape == (3, 84, 84)
        return x_t.astype(np.uint8)

def wrap_doom_84x84(env, frame_skip=4, history_len=4):
  env = ToDiscrete("minimal")(env)
  env = SkipFrameEnv(env, skip=frame_skip)
  env = ProcessFrame84x84Doom(env)
  env = Transpose3DEnv(env)
  if history_len > 1:
    env = StackFrameEnv(env, history_len)
  return env

class PreprocDoomRGBDMask(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(PreprocDoomRGBDMask, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(4, 120, 160))

    def _observation(self, obs):
        if obs.size == 120 * 160 * 4:
            img = np.reshape(obs, (4, 120, 160))
        else:
            assert False, "Unknown resolution: {}".format(obs.size)
        img[:,30:90,50:110] = 0
        return img

class PreprocDoomRGBDMaskBig(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(PreprocDoomRGBDMaskBig, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(4, 120, 160))

    def _observation(self, obs):
        if obs.size == 120 * 160 * 4:
            img = np.reshape(obs, (4, 120, 160))
        else:
            assert False, "Unknown resolution: {}".format(obs.size)
        img[:,10:110,30:130] = 0
        return img

class PreprocDoomRGBDMaskOut(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(PreprocDoomRGBDMaskOut, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(4, 120, 160))

    def _observation(self, obs):
        if obs.size == 120 * 160 * 4:
            img = np.reshape(obs, (4, 120, 160))
        else:
            assert False, "Unknown resolution: {}".format(obs.size)
        img[:,:40,:] = 0
        img[:,80:,:] = 0
        img[:,:,:60] = 0
        img[:,:,100:] = 0
        return img

class ProcessDoomRGBD1x1(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessDoomRGBD1x1, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(3, 1, 1))

    def _observation(self, obs):
        return ProcessDoomRGBD1x1.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 4:
            img = np.reshape(frame, (4, 120, 160)).astype(np.float32)
            img_color = img[:3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(frame.size)
        img_color = np.transpose(img_color, (1, 2, 0))
        x_t = cv2.resize(img_color, (1, 1), interpolation=cv2.INTER_AREA)
        assert x_t.shape == (1, 1, 3)
        x_t = np.reshape(x_t, (3, 1, 1))
        return x_t.astype(np.uint8)

def wrap_doom_1x1_v2(env, frame_skip=4, history_len=4, zclip=False):
  env = ToDiscrete("minimal")(env)
  if frame_skip > 1:
    env = SkipFrameEnv(env, skip=frame_skip)
  if zclip:
    raise NotImplementedError
  else:
    env = ProcessDoomRGBD1x1(env)
  if history_len > 1:
    env = StackFrameEnv(env, history_len)
  return env

class ProcessDoomRGBD20(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessDoomRGBD20, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(20, 20, 3))

    def _observation(self, obs):
        return ProcessDoomRGBD20.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 4:
            img = np.reshape(frame, (4, 120, 160)).astype(np.float32)
            img_color = img[:3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(frame.size)
        img_color = np.transpose(img_color, (1, 2, 0))
        x_t = cv2.resize(img_color, (20, 20), interpolation=cv2.INTER_AREA)
        assert x_t.shape == (20, 20, 3)
        return x_t.astype(np.uint8)

def wrap_doom_20x20_v2(env, frame_skip=4, action_delay=None, zclip=False, depth=False, frame_flicker=False, history_len=4):
  env = ToDiscrete("minimal")(env)
  if frame_skip > 1:
    env = SkipFrameEnv(env, skip=frame_skip)
  if action_delay is not None and action_delay >= 1:
    raise NotImplementedError
  if zclip:
    raise NotImplementedError
  elif depth:
    raise NotImplementedError
  elif frame_flicker:
    raise NotImplementedError
  else:
    env = ProcessDoomRGBD20(env)
  env = Transpose3DEnv(env)
  if history_len > 1:
    env = StackFrameEnv(env, history_len)
  return env

class ProcessDoomRGBD44(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessDoomRGBD44, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(44, 44, 3))

    def _observation(self, obs):
        return ProcessDoomRGBD44.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 4:
            img = np.reshape(frame, (4, 120, 160)).astype(np.float32)
            img_color = img[:3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(frame.size)
        img_color = np.transpose(img_color, (1, 2, 0))
        x_t = cv2.resize(img_color, (44, 44), interpolation=cv2.INTER_AREA)
        assert x_t.shape == (44, 44, 3)
        return x_t.astype(np.uint8)

def wrap_doom_44x44_v2(env, frame_skip=4, history_len=4, zclip=False):
  env = ToDiscrete("minimal")(env)
  if frame_skip > 1:
    env = SkipFrameEnv(env, skip=frame_skip)
  if zclip:
    raise NotImplementedError
  else:
    env = ProcessDoomRGBD44(env)
  env = Transpose3DEnv(env)
  if history_len > 1:
    env = StackFrameEnv(env, history_len)
  return env

class ProcessDoomRGBD(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessDoomRGBD, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 3))

    def _observation(self, obs):
        return ProcessDoomRGBD.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 4:
            img = np.reshape(frame, (4, 120, 160)).astype(np.float32)
            img_color = img[:3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(frame.size)
        img_color = np.transpose(img_color, (1, 2, 0))
        x_t = cv2.resize(img_color, (84, 84), interpolation=cv2.INTER_LINEAR)
        assert x_t.shape == (84, 84, 3)
        return x_t.astype(np.uint8)

class ProcessDoomRGBDDepth(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessDoomRGBDDepth, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 1))

    def _observation(self, obs):
        return ProcessDoomRGBDDepth.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 4:
            img = np.reshape(frame, (4, 120, 160)).astype(np.float32)
            img_depth = img[3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(frame.size)
        #img_depth = np.reshape(img_depth, (120, 160, 1))
        x_t = cv2.resize(img_depth, (84, 84), interpolation=cv2.INTER_LINEAR)
        assert x_t.shape == (84, 84)
        x_t = np.reshape(x_t, (84, 84, 1))
        return x_t.astype(np.uint8)

class ProcessDoomRGBDHardZClip(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(ProcessDoomRGBDHardZClip, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 3))

    def _observation(self, obs):
        return ProcessDoomRGBDHardZClip.process(obs)

    @staticmethod
    def process(frame):
        if frame.size == 120 * 160 * 4:
            img = np.reshape(frame, (4, 120, 160)).astype(np.float32)
            img_color = img[:3,:,:]
            img_depth = img[3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(frame.size)
        img_color = img_color * (img_depth <= 14).astype(np.float32)
        img_color = np.transpose(img_color, (1, 2, 0))
        x_t = cv2.resize(img_color, (84, 84), interpolation=cv2.INTER_LINEAR)
        assert x_t.shape == (84, 84, 3)
        return x_t.astype(np.uint8)

class ProcessDoomRGBDFlicker(gym.ObservationWrapper):
    def __init__(self, env=None, flicker_prob=None):
        super(ProcessDoomRGBDFlicker, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 3))

    def _observation(self, obs):
        r = np.random.choice(2)
        if r == 0:
            if obs.size == 120 * 160 * 4:
                img = np.reshape(obs, (4, 120, 160)).astype(np.float32)
                img_color = img[:3,:,:]
            else:
                assert False, "Unknown resolution: {}".format(obs.size)
            img_color = np.transpose(img_color, (1, 2, 0))
            x_t = cv2.resize(img_color, (84, 84), interpolation=cv2.INTER_LINEAR)
            assert x_t.shape == (84, 84, 3)
            return x_t.astype(np.uint8)
        elif r == 1:
            return np.zeros((84, 84, 3), dtype=np.uint8)
        else:
            assert False

class ProcessDoomRGBDRepeat(gym.Wrapper):
    def __init__(self, env=None, frame_repeat=None):
        super(ProcessDoomRGBDRepeat, self).__init__(env)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 3))
        self._frame_repeat = frame_repeat
        #self._obs_queue = deque([], maxlen=frame_repeat)
        self._frame_ctr = 0
        self._curr_obs = None

    def _reset(self):
        obs = self.reset()
        self._frame_ctr = 0
        self._curr_obs = obs
        return self._process_obs(self._curr_obs)

    def _step(self, action):
        obs, reward, done, info = self.env.step(action)
        self._frame_ctr += 1
        if self._frame_ctr == self._frame_repeat:
            self._frame_ctr = 0
            self._curr_obs = obs
        return self._process_obs(self._curr_obs), reward, done, info

    def _process_obs(self, obs):
        if obs.size == 120 * 160 * 4:
            img = np.reshape(obs, (4, 120, 160)).astype(np.float32)
            img_color = img[:3,:,:]
        else:
            assert False, "Unknown resolution: {}".format(obs.size)
        img_color = np.transpose(img_color, (1, 2, 0))
        x_t = cv2.resize(img_color, (84, 84), interpolation=cv2.INTER_LINEAR)
        assert x_t.shape == (84, 84, 3)
        return x_t.astype(np.uint8)

def wrap_doom_84x84_v2(env, frame_skip=4, action_delay=None, preproc_mask=False, preproc_mask_big=False, preproc_mask_out=False, zclip=False, depth=False, frame_flicker=False, frame_repeat=None, history_len=4):
  env = ToDiscrete("minimal")(env)
  if frame_skip > 1:
    env = SkipFrameEnv(env, skip=frame_skip)
  if action_delay is not None and action_delay > 0:
    env = DelayedActionEnv(env, action_delay)
  if preproc_mask:
    env = PreprocDoomRGBDMask(env)
  elif preproc_mask_big:
    env = PreprocDoomRGBDMaskBig(env)
  elif preproc_mask_out:
    env = PreprocDoomRGBDMaskOut(env)
  if zclip:
    env = ProcessDoomRGBDHardZClip(env)
  elif depth:
    env = ProcessDoomRGBDDepth(env)
  elif frame_flicker:
    env = ProcessDoomRGBDFlicker(env, 0.5)
  elif frame_repeat is not None and frame_repeat > 1:
    env = ProcessDoomRGBDRepeat(env, frame_repeat)
  else:
    env = ProcessDoomRGBD(env)
  env = Transpose3DEnv(env)
  if history_len > 1:
    env = StackFrameEnv(env, history_len)
  return env

max_stars_repo_name: peterhj/arm-pytorch

max_issues_repo_name: peterhj/arm-pytorch

max_forks_repo_name: peterhj/arm-pytorch


text:
import torch
import numpy as np
from colour_demosaicing.bayer import masks_CFA_Bayer
import torch.nn as nn
import torch.nn.functional as F

def demosaicing_CFA_Bayer_bilinear(CFA, pattern='RGGB'):
    """
    Returns the demosaiced *RGB* colourspace array from given *Bayer* CFA using
    bilinear interpolation.

    Parameters
    ----------
    CFA : array_like
        *Bayer* CFA.
    pattern : unicode, optional
        **{'RGGB', 'BGGR', 'GRBG', 'GBRG'}**,
        Arrangement of the colour filters on the pixel array.

    Returns
    -------
    ndarray
        *RGB* colourspace array.

    Notes
    -----
    -   The definition output is not clipped in range [0, 1] : this allows for
        direct HDRI / radiance image generation on *Bayer* CFA data and post
        demosaicing of the high dynamic range data as showcased in this
        `Jupyter Notebook <https://github.com/colour-science/colour-hdri/\
blob/develop/colour_hdri/examples/\
examples_merge_from_raw_files_with_post_demosaicing.ipynb>`__.

    References
    ----------
    :cite:`Losson2010c`

    Examples
    --------
    >>> import numpy as np
    >>> CFA = np.array(
    ...     [[0.30980393, 0.36078432, 0.30588236, 0.3764706],
    ...      [0.35686275, 0.39607844, 0.36078432, 0.40000001]])
    >>> demosaicing_CFA_Bayer_bilinear(CFA)
    array([[[ 0.69705884,  0.17941177,  0.09901961],
            [ 0.46176472,  0.4509804 ,  0.19803922],
            [ 0.45882354,  0.27450981,  0.19901961],
            [ 0.22941177,  0.5647059 ,  0.30000001]],
    <BLANKLINE>
           [[ 0.23235295,  0.53529412,  0.29705883],
            [ 0.15392157,  0.26960785,  0.59411766],
            [ 0.15294118,  0.4509804 ,  0.59705884],
            [ 0.07647059,  0.18431373,  0.90000002]]])
    >>> CFA = np.array(
    ...     [[0.3764706, 0.360784320, 0.40784314, 0.3764706],
    ...      [0.35686275, 0.30980393, 0.36078432, 0.29803923]])
    >>> demosaicing_CFA_Bayer_bilinear(CFA, 'BGGR')
    array([[[ 0.07745098,  0.17941177,  0.84705885],
            [ 0.15490197,  0.4509804 ,  0.5882353 ],
            [ 0.15196079,  0.27450981,  0.61176471],
            [ 0.22352942,  0.5647059 ,  0.30588235]],
    <BLANKLINE>
           [[ 0.23235295,  0.53529412,  0.28235295],
            [ 0.4647059 ,  0.26960785,  0.19607843],
            [ 0.45588237,  0.4509804 ,  0.20392157],
            [ 0.67058827,  0.18431373,  0.10196078]]])
    """

    ## Above is the original version on mosaicing_demosaicing package processing image based on numpy arrays, we adapt it to a torch tensor version as follows:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    batch, h, w= CFA.size()

    R_m, G_m, B_m = masks_CFA_Bayer([h, w], pattern)

    R_m = R_m[np.newaxis, np.newaxis, :]
    R_m = np.repeat(R_m, batch, axis = 0)
    G_m = G_m[np.newaxis, np.newaxis, :]
    G_m = np.repeat(G_m, batch, axis=0)
    B_m = B_m[np.newaxis, np.newaxis, :]
    B_m = np.repeat(B_m, batch, axis=0)

    R_m = torch.from_numpy(R_m).to(device)
    G_m = torch.from_numpy(G_m).to(device)
    B_m = torch.from_numpy(B_m).to(device)

    H_G = np.array(
        [[0, 1, 0],
         [1, 4, 1],
         [0, 1, 0]]) / 4  # yapf: disable

    H_G = H_G[np.newaxis, np.newaxis, :]
    H_G = torch.from_numpy(H_G).to(device)

    H_RB = np.array(
        [[1, 2, 1],
         [2, 4, 2],
         [1, 2, 1]]) / 4  # yapf: disable

    H_RB = H_RB[np.newaxis, np.newaxis, :]
    H_RB = torch.from_numpy(H_RB).to(device)
    CFA = CFA.unsqueeze(1)

    R = F.conv2d(CFA * R_m, H_RB, stride=1, padding=1)
    G = F.conv2d(CFA * G_m, H_G, stride=1, padding=1)
    B = F.conv2d(CFA * B_m, H_RB, stride=1, padding=1)

    R = R.squeeze(1)
    G = G.squeeze(1)
    B = B.squeeze(1)

    del R_m, G_m, B_m, H_RB, H_G
    torch.cuda.empty_cache()

    return torch.stack((R, G, B), dim = 3)

def mosaicing_CFA_Bayer(RGB, pattern = 'RGGB'):
    """
    Returns the *Bayer* CFA mosaic for a given *RGB* colourspace array.

    Parameters
    ----------
    RGB : array_like
        *RGB* colourspace array.
    pattern : unicode, optional
        **{'RGGB', 'BGGR', 'GRBG', 'GBRG'}**,
        Arrangement of the colour filters on the pixel array.

    Returns
    -------
    ndarray
        *Bayer* CFA mosaic.

    Examples
    --------
    >>> import numpy as np
    >>> RGB = np.array([[[0, 1, 2],
    ...                  [0, 1, 2]],
    ...                 [[0, 1, 2],
    ...                  [0, 1, 2]]])
    >>> mosaicing_CFA_Bayer(RGB)
    array([[ 0.,  1.],
           [ 1.,  2.]])
    >>> mosaicing_CFA_Bayer(RGB, pattern='BGGR')
    array([[ 2.,  1.],
           [ 1.,  0.]])
    """

    ## Above is the original version on mosaicing_demosaicing package processing image based on numpy arrays, we adapt it to a torch tensor version as follows:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    R = RGB[:, :, :, 0]
    G = RGB[:, :, :, 1]
    B = RGB[:, :, :, 2]

    batch, _, _, _ = RGB.shape
    R_m, G_m, B_m = masks_CFA_Bayer(RGB.shape[1:3], pattern)

    G_m = G_m[np.newaxis, :]
    G_m = np.repeat(G_m, batch, axis = 0)
    B_m = B_m[np.newaxis, :]
    B_m = np.repeat(B_m, batch, axis = 0)
    R_m = R_m[np.newaxis, :]
    R_m = np.repeat(R_m, batch, axis = 0)

    R_m = torch.from_numpy(R_m).to(device)
    G_m = torch.from_numpy(G_m).to(device)
    B_m = torch.from_numpy(B_m).to(device)

    CFA = R * R_m + G * G_m + B * B_m
    del R_m, G_m, B_m
    torch.cuda.empty_cache()

    return CFA