text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

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

from contextlib import contextmanager
import tensorflow as tf
import numpy as np
from scipy import stats
from scipy.special import logsumexp, factorial, gammaln

from tests.distributions import utils
from zhusuan.distributions.multivariate import *

class TestMultivariateNormalCholesky(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                MultivariateNormalCholesky(tf.zeros([]), tf.zeros([]))
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                MultivariateNormalCholesky(tf.zeros([1]), tf.zeros([1]))
            with self.assertRaisesRegexp(ValueError, 'compatible'):
                MultivariateNormalCholesky(
                    tf.zeros([1, 2]), tf.placeholder(tf.float32, [1, 2, 3]))
            u = tf.placeholder(tf.float32, [None])
            len_u = tf.shape(u)[0]
            dst = MultivariateNormalCholesky(
                tf.zeros([2]), tf.zeros([len_u, len_u]))
            with self.assertRaisesRegexp(
                    tf.errors.InvalidArgumentError, 'compatible'):
                dst.sample().eval(feed_dict={u: np.ones((3,))})

    def test_shape_inference(self):
        with self.test_session(use_gpu=True):
            # Static
            mean = 10 * np.random.normal(size=(10, 11, 2)).astype('d')
            cov = np.zeros((10, 11, 2, 2))
            dst = MultivariateNormalCholesky(
                tf.constant(mean), tf.constant(cov))
            self.assertEqual(dst.get_batch_shape().as_list(), [10, 11])
            self.assertEqual(dst.get_value_shape().as_list(), [2])
            # Dynamic
            unk_mean = tf.placeholder(tf.float32, None)
            unk_cov = tf.placeholder(tf.float32, None)
            dst = MultivariateNormalCholesky(unk_mean, unk_cov)
            self.assertEqual(dst.get_value_shape().as_list(), [None])
            feed_dict = {unk_mean: np.ones(2), unk_cov: np.eye(2)}
            self.assertEqual(list(dst.batch_shape.eval(feed_dict)), [])
            self.assertEqual(list(dst.value_shape.eval(feed_dict)), [2])

    def _gen_test_params(self, seed):
        np.random.seed(seed)
        mean = 10 * np.random.normal(size=(10, 11, 3)).astype('d')
        cov = np.zeros((10, 11, 3, 3))
        cov_chol = np.zeros_like(cov)
        for i in range(10):
            for j in range(11):
                cov[i, j] = stats.invwishart.rvs(3, np.eye(3))
                cov[i, j] /= np.max(np.diag(cov[i, j]))
                cov_chol[i, j, :, :] = np.linalg.cholesky(cov[i, j])
        return mean, cov, cov_chol

    @contextmanager
    def fixed_randomness_session(self, seed):
        with tf.Graph().as_default() as g:
            with self.test_session(use_gpu=True, graph=g):
                tf.set_random_seed(seed)
                yield

    def test_sample(self):
        with self.fixed_randomness_session(233):
            def test_sample_with(seed):
                mean, cov, cov_chol = self._gen_test_params(seed)
                dst = MultivariateNormalCholesky(
                    tf.constant(mean), tf.constant(cov_chol))
                n_exp = 20000
                samples = dst.sample(n_exp)
                sample_shape = (n_exp, 10, 11, 3)
                self.assertEqual(samples.shape.as_list(), list(sample_shape))
                samples = dst.sample(n_exp).eval()
                self.assertEqual(samples.shape, sample_shape)
                self.assertAllClose(
                    np.mean(samples, axis=0), mean, rtol=5e-2, atol=5e-2)
                for i in range(10):
                    for j in range(11):
                        self.assertAllClose(
                            np.cov(samples[:, i, j, :].T), cov[i, j],
                            rtol=1e-1, atol=1e-1)

            for seed in [23, 233, 2333]:
                test_sample_with(seed)

    def test_prob(self):
        with self.fixed_randomness_session(233):
            def test_prob_with(seed):
                mean, cov, cov_chol = self._gen_test_params(seed)
                dst = MultivariateNormalCholesky(
                    tf.constant(mean), tf.constant(cov_chol),
                    check_numerics=True)
                n_exp = 200
                samples = dst.sample(n_exp).eval()
                log_pdf = dst.log_prob(tf.constant(samples))
                pdf_shape = (n_exp, 10, 11)
                self.assertEqual(log_pdf.shape.as_list(), list(pdf_shape))
                log_pdf = log_pdf.eval()
                self.assertEqual(log_pdf.shape, pdf_shape)
                for i in range(10):
                    for j in range(11):
                        log_pdf_exact = stats.multivariate_normal.logpdf(
                                samples[:, i, j, :], mean[i, j], cov[i, j])
                        self.assertAllClose(
                            log_pdf_exact, log_pdf[:, i, j])
                self.assertAllClose(
                    np.exp(log_pdf), dst.prob(tf.constant(samples)).eval())

            for seed in [23, 233, 2333]:
                test_prob_with(seed)

    def test_sample_reparameterized(self):
        mean, cov, cov_chol = self._gen_test_params(23)
        mean = tf.constant(mean)
        cov_chol = tf.constant(cov_chol)
        mvn_rep = MultivariateNormalCholesky(mean, cov_chol)
        samples = mvn_rep.sample(tf.placeholder(tf.int32, shape=[]))
        mean_grads, cov_grads = tf.gradients(samples, [mean, cov_chol])
        self.assertTrue(mean_grads is not None)
        self.assertTrue(cov_grads is not None)
        mvn_no_rep = MultivariateNormalCholesky(
                mean, cov_chol, is_reparameterized=False)
        samples = mvn_no_rep.sample(tf.placeholder(tf.int32, shape=[]))
        mean_grads, cov_grads = tf.gradients(samples, [mean, cov_chol])
        self.assertEqual(mean_grads, None)
        self.assertEqual(cov_grads, None)

class TestMultinomial(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                Multinomial(tf.zeros([]), n_experiments=10)

    def test_init_n(self):
        dist = Multinomial(tf.ones([2]), n_experiments=10)
        self.assertTrue(isinstance(dist.n_categories, int))
        self.assertEqual(dist.n_categories, 2)
        self.assertTrue(isinstance(dist.n_experiments, int))
        self.assertEqual(dist.n_experiments, 10)

        with self.test_session(use_gpu=True) as sess:
            logits = tf.placeholder(tf.float32, None)
            n_experiments = tf.placeholder(tf.int32, None)
            dist2 = Multinomial(logits, n_experiments=n_experiments)
            self.assertEqual(
                sess.run([dist2.n_categories, dist2.n_experiments],
                         feed_dict={logits: np.ones([2]), n_experiments: 10}),
                [2, 10])
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should have rank"):
                dist2.n_categories.eval(feed_dict={logits: 1.,
                                                   n_experiments: 10})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should be a scalar"):
                dist2.n_experiments.eval(feed_dict={logits: [1.],
                                                    n_experiments: [10]})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "must be positive"):
                dist2.n_experiments.eval(feed_dict={logits: [1.],
                                                    n_experiments: 0})

    def test_value_shape(self):
        # static
        dist = Multinomial(tf.placeholder(tf.float32, [None, 2]),
                           n_experiments=10)
        self.assertEqual(dist.get_value_shape().as_list(), [2])

        # dynamic
        logits = tf.placeholder(tf.float32, None)
        dist2 = Multinomial(logits, n_experiments=10)
        self.assertTrue(dist2._value_shape().dtype is tf.int32)
        with self.test_session(use_gpu=True):
            self.assertEqual(dist2._value_shape().eval(
                feed_dict={logits: np.ones([2])}).tolist(), [2])

        self.assertEqual(dist._value_shape().dtype, tf.int32)

    def test_batch_shape(self):
        def _distribution(param):
            return Multinomial(param, n_experiments=10)
        utils.test_batch_shape_1parameter(
            self, _distribution, np.zeros, is_univariate=False)

    def test_sample(self):
        def _distribution(param):
            return Multinomial(param, n_experiments=10)
        utils.test_1parameter_sample_shape_one_rank_less(
            self, _distribution, np.zeros)
        dist = Multinomial(np.ones([2, 2]), n_experiments=None)
        with self.assertRaisesRegexp(ValueError,
                                     "Cannot sample when `n_experiments`"):
            dist.sample()

    def test_log_prob_shape(self):
        def _distribution(param):
            return Multinomial(param, n_experiments=10)

        def _make_samples(shape):
            samples = np.zeros(shape)
            samples = samples.reshape((-1, shape[-1]))
            samples[:, 0] = 1
            return samples.reshape(shape)

        utils.test_1parameter_log_prob_shape_one_rank_less(
            self, _distribution, _make_samples, _make_samples)

    def test_value(self):
        with self.test_session(use_gpu=True):
            def _test_value(logits, n_experiments, given, normalize_logits):
                logits = np.array(logits, np.float32)
                given = np.array(given)
                dist = Multinomial(logits, n_experiments=None,
                                   normalize_logits=normalize_logits)
                log_p = dist.log_prob(given)
                if n_experiments is not None:
                    dist2 = Multinomial(logits, n_experiments=n_experiments,
                                        normalize_logits=normalize_logits)
                    log_p_2 = dist2.log_prob(given)
                    self.assertAllClose(log_p.eval(), log_p_2.eval())

                maybe_normalized_logits = logits
                if normalize_logits:
                    maybe_normalized_logits -= logsumexp(
                        logits, axis=-1, keepdims=True)
                n_experiments = np.sum(given, axis=-1)
                target_log_p = np.log(factorial(n_experiments)) - \
                    np.sum(np.log(factorial(given)), -1) + \
                    np.sum(given * maybe_normalized_logits, -1)
                self.assertAllClose(log_p.eval(), target_log_p)
                p = dist.prob(given)
                target_p = np.exp(target_log_p)
                self.assertAllClose(np.log(p.eval()), np.log(target_p))

            for normalize_logits in [True, False]:
                _test_value([-50., -20., 0.], 4, [1, 0, 3], normalize_logits)
                _test_value([1., 10., 1000.], 1, [1, 0, 0], normalize_logits)
                _test_value([[2., 3., 1.], [5., 7., 4.]], 7,
                            np.array([3, 1, 3], dtype=np.int32),
                            normalize_logits)
                _test_value([-10., 10., 20., 50.], 100,
                            [[0, 1, 49, 50], [50, 49, 1, 0]],
                            normalize_logits)

    def test_dtype(self):
        def _distribution(param, **kwargs):
            return Multinomial(param, n_experiments=10, **kwargs)
        utils.test_dtype_1parameter_discrete(self, _distribution)

        with self.assertRaisesRegexp(TypeError, "n_experiments must be"):
            Multinomial([1., 1.], n_experiments=tf.placeholder(tf.float32, []))
        with self.assertRaisesRegexp(TypeError,
                                     "n_experiments must be integer"):
            Multinomial([1., 1.], n_experiments=2.0)

class TestUnnormalizedMultinomial(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                UnnormalizedMultinomial(tf.zeros([]))

    def test_init_n(self):
        dist = UnnormalizedMultinomial(tf.ones([2]))
        self.assertTrue(isinstance(dist.n_categories, int))
        self.assertEqual(dist.n_categories, 2)

        with self.test_session(use_gpu=True) as sess:
            logits = tf.placeholder(tf.float32, None)
            dist2 = UnnormalizedMultinomial(logits)
            self.assertEqual(
                sess.run(dist2.n_categories, feed_dict={logits: np.ones([2])}),
                2)
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should have rank"):
                dist2.n_categories.eval(feed_dict={logits: 1.})

    def test_value_shape(self):
        # static
        dist = UnnormalizedMultinomial(tf.placeholder(tf.float32, [None, 2]))
        self.assertEqual(dist.get_value_shape().as_list(), [2])

        # dynamic
        logits = tf.placeholder(tf.float32, None)
        dist2 = UnnormalizedMultinomial(logits)
        self.assertTrue(dist2._value_shape().dtype is tf.int32)
        with self.test_session(use_gpu=True):
            self.assertEqual(dist2._value_shape().eval(
                feed_dict={logits: np.ones([2])}).tolist(), [2])

        self.assertEqual(dist._value_shape().dtype, tf.int32)

    def test_batch_shape(self):
        utils.test_batch_shape_1parameter(
            self, UnnormalizedMultinomial, np.zeros, is_univariate=False)

    def test_sample(self):
        dist = UnnormalizedMultinomial(np.ones([2, 2]))
        with self.assertRaisesRegexp(NotImplementedError,
                                     "Unnormalized multinomial distribution"
                                     " does not support sampling"):
            dist.sample()

    def test_log_prob_shape(self):
        def _distribution(param):
            return UnnormalizedMultinomial(param)

        def _make_samples(shape):
            samples = np.zeros(shape)
            samples = samples.reshape((-1, shape[-1]))
            samples[:, 0] = 1
            return samples.reshape(shape)

        utils.test_1parameter_log_prob_shape_one_rank_less(
            self, _distribution, _make_samples, _make_samples)

    def test_value(self):
        with self.test_session(use_gpu=True):
            def _test_value(logits, given, normalize_logits):
                logits = np.array(logits, np.float32)
                given = np.array(given)
                dist = UnnormalizedMultinomial(
                    logits, normalize_logits=normalize_logits)
                log_p = dist.log_prob(given)

                maybe_normalized_logits = logits
                if normalize_logits:
                    maybe_normalized_logits -= logsumexp(
                        logits, axis=-1, keepdims=True)
                target_log_p = np.sum(given * maybe_normalized_logits, -1)
                self.assertAllClose(log_p.eval(), target_log_p)
                p = dist.prob(given)
                target_p = np.exp(target_log_p)
                self.assertAllClose(p.eval(), target_p)

            for normalize_logits in [True, False]:
                _test_value([-50., -20., 0.], [1, 0, 3], normalize_logits)
                _test_value([1., 10., 1000.], [1, 0, 0], normalize_logits)
                _test_value([[2., 3., 1.], [5., 7., 4.]],
                            np.ones([3, 1, 3], dtype=np.int32),
                            normalize_logits)
                _test_value([-10., 10., 20., 50.],
                            [[0, 1, 99, 100], [100, 99, 1, 0]],
                            normalize_logits)

    def test_dtype(self):
        utils.test_dtype_1parameter_discrete(self, UnnormalizedMultinomial,
                                             prob_only=True)

class TestOnehotCategorical(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                OnehotCategorical(logits=tf.zeros([]))

    def test_init_n_categories(self):
        cat = OnehotCategorical(tf.ones([10]))
        self.assertTrue(isinstance(cat.n_categories, int))
        self.assertEqual(cat.n_categories, 10)

        with self.test_session(use_gpu=True):
            logits = tf.placeholder(tf.float32, None)
            cat2 = OnehotCategorical(logits)
            self.assertEqual(
                cat2.n_categories.eval(feed_dict={logits: np.ones([10])}), 10)
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should have rank"):
                cat2.n_categories.eval(feed_dict={logits: 1.})

    def test_value_shape(self):
        # static
        cat = OnehotCategorical(tf.placeholder(tf.float32, [None, 10]))
        self.assertEqual(cat.get_value_shape().as_list(), [10])

        # dynamic
        logits = tf.placeholder(tf.float32, None)
        cat2 = OnehotCategorical(logits)
        self.assertTrue(cat2._value_shape().dtype is tf.int32)
        with self.test_session(use_gpu=True):
            self.assertEqual(cat2._value_shape().eval(
                feed_dict={logits: np.ones([2, 1, 3])}).tolist(), [3])

        self.assertEqual(cat._value_shape().dtype, tf.int32)

    def test_batch_shape(self):
        utils.test_batch_shape_1parameter(
            self, OnehotCategorical, np.zeros, is_univariate=False)

    def test_sample_shape(self):
        utils.test_1parameter_sample_shape_one_rank_less(
            self, OnehotCategorical, np.zeros)

    def test_log_prob_shape(self):
        def _make_samples(shape):
            samples = np.zeros(shape)
            samples = samples.reshape((-1, shape[-1]))
            samples[:, 0] = 1
            return samples.reshape(shape)

        utils.test_1parameter_log_prob_shape_one_rank_less(
            self, OnehotCategorical, _make_samples, _make_samples)

    def test_value(self):
        with self.test_session(use_gpu=True):
            def _test_value(logits, given):
                logits = np.array(logits, np.float32)
                normalized_logits = logits - logsumexp(
                    logits, axis=-1, keepdims=True)
                given = np.array(given, np.int32)
                cat = OnehotCategorical(logits)
                log_p = cat.log_prob(tf.one_hot(given, logits.shape[-1],
                                                dtype=tf.int32))

                def _one_hot(x, depth):
                    n_elements = x.size
                    ret = np.zeros((n_elements, depth))
                    ret[np.arange(n_elements), x.flat] = 1
                    return ret.reshape(list(x.shape) + [depth])

                target_log_p = np.sum(_one_hot(
                    given, logits.shape[-1]) * normalized_logits, -1)
                self.assertAllClose(log_p.eval(), target_log_p)
                p = cat.prob(tf.one_hot(given, logits.shape[-1],
                                        dtype=tf.int32))
                target_p = np.sum(_one_hot(
                    given, logits.shape[-1]) * np.exp(normalized_logits), -1)
                self.assertAllClose(p.eval(), target_p)

            _test_value([0.], [0, 0, 0])
            _test_value([-50., -10., -50.], [0, 1, 2, 1])
            _test_value([0., 4.], [[0, 1], [0, 1]])
            _test_value([[2., 3., 1.], [5., 7., 4.]],
                        np.ones([3, 1, 1], dtype=np.int32))

    def test_dtype(self):
        utils.test_dtype_1parameter_discrete(self, OnehotCategorical)

class TestDirichlet(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                Dirichlet(alpha=tf.zeros([]))

    def test_init_n_categories(self):
        dist = Dirichlet(tf.ones([10]))
        self.assertTrue(isinstance(dist.n_categories, int))
        self.assertEqual(dist.n_categories, 10)
        with self.assertRaisesRegexp(ValueError,
                                     "n_categories.*should be at least 2"):
            Dirichlet(tf.ones([3, 1]))
        dist2 = Dirichlet(tf.placeholder(tf.float32, [3, None]))
        self.assertTrue(dist2.n_categories is not None)

        with self.test_session(use_gpu=True):
            alpha = tf.placeholder(tf.float32, None)
            dist3 = Dirichlet(alpha)
            self.assertEqual(
                dist3.n_categories.eval(feed_dict={alpha: np.ones([10])}), 10)
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should have rank"):
                dist3.n_categories.eval(feed_dict={alpha: 1.})

    def test_value_shape(self):
        # static
        dist = Dirichlet(tf.placeholder(tf.float32, [None, 10]))
        self.assertEqual(dist.get_value_shape().as_list(), [10])

        # dynamic
        alpha = tf.placeholder(tf.float32, None)
        dist2 = Dirichlet(alpha)
        self.assertEqual(dist2.get_value_shape().as_list(), [None])
        self.assertTrue(dist2._value_shape().dtype is tf.int32)
        with self.test_session(use_gpu=True):
            self.assertEqual(dist2._value_shape().eval(
                feed_dict={alpha: np.ones([2, 1, 3])}).tolist(), [3])
        self.assertEqual(dist._value_shape().dtype, tf.int32)

    def test_batch_shape(self):
        utils.test_batch_shape_1parameter(
            self, Dirichlet, np.zeros, is_univariate=False)

    def test_sample_shape(self):
        utils.test_1parameter_sample_shape_one_rank_less(
            self, Dirichlet, np.zeros)

    def test_log_prob_shape(self):
        def _make_samples(shape):
            samples = np.ones(shape, dtype=np.float32)
            return samples / samples.sum(axis=-1, keepdims=True)

        # TODO: This failed with a bug in Tensorflow, waiting fix.
        # https://github.com/tensorflow/tensorflow/issues/8391
        # _test_static([3, None], [3, 2, 1, None], [3, 2, 3])
        utils.test_1parameter_log_prob_shape_one_rank_less(
            self, Dirichlet, np.ones, _make_samples)

    def test_value(self):
        def dirichlet_logpdf(x, alpha):
            # scipy's implementation of dirichlet logpdf doesn't support
            # batch of x, we use this modified version.
            def _lnB(alpha):
                return np.sum(gammaln(alpha)) - gammaln(np.sum(alpha))

            lnB = _lnB(alpha)
            return - lnB + np.sum(np.log(x) * (alpha - 1), -1)

        def dirichlet_pdf(x, alpha):
            return np.exp(dirichlet_logpdf(x, alpha))

        with self.test_session(use_gpu=True):
            def _test_value_alpha_rank1(alpha, given):
                alpha = np.array(alpha, np.float32)
                given = np.array(given, np.float32)
                dist = Dirichlet(alpha)
                log_p = dist.log_prob(given)
                target_log_p = dirichlet_logpdf(given, alpha)
                self.assertAllClose(log_p.eval(), target_log_p)
                p = dist.prob(given)
                target_p = dirichlet_pdf(given, alpha)
                self.assertAllClose(p.eval(), target_p)

            _test_value_alpha_rank1([1., 1., 1.],
                                    [[0.2, 0.5, 0.3], [0.3, 0.4, 0.3]])
            _test_value_alpha_rank1([2., 3., 4.], [0.3, 0.7, 0.])
            # TODO: fix for case when alpha=1, given=0

            def _test_value_alpha_rank2_given_rank2(alpha, given):
                alpha = np.array(alpha, np.float32)
                given = np.array(given, np.float32)
                alpha_b = alpha * np.ones_like(given)
                given_b = given * np.ones_like(alpha)
                dist = Dirichlet(alpha)
                log_p = dist.log_prob(given)
                target_log_p = np.array(
                    [dirichlet_logpdf(given_b[i], alpha_b[i])
                     for i in range(alpha_b.shape[0])])
                self.assertAllClose(log_p.eval(), target_log_p)
                p = dist.prob(given)
                target_p = np.array(
                    [dirichlet_pdf(given_b[i], alpha_b[i])
                     for i in range(alpha_b.shape[0])])
                self.assertAllClose(p.eval(), target_p)

            _test_value_alpha_rank2_given_rank2([[1., 2.], [3., 4.]],
                                                [0.5, 0.5])
            _test_value_alpha_rank2_given_rank2([[5., 6.], [7., 8.]],
                                                [[0.1, 0.9]])
            _test_value_alpha_rank2_given_rank2([[100., 1.], [0.01, 10.]],
                                                [[0., 1.], [1., 0.]])

    def test_check_numerics(self):
        alpha = tf.placeholder(tf.float32, None)
        given = tf.placeholder(tf.float32, None)
        dist = Dirichlet(alpha, check_numerics=True)
        log_p = dist.log_prob(given)
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(given\).*Tensor had Inf"):
                log_p.eval(feed_dict={alpha: np.ones([2]), given: [0., 1.]})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "lbeta\(alpha\).*Tensor had NaN"):
                log_p.eval(feed_dict={alpha: [-1., 1.], given: [0.5, 0.5]})

    def test_dtype(self):
        utils.test_dtype_1parameter_continuous(self, Dirichlet)

class TestExpConcrete(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                ExpConcrete(1., logits=tf.zeros([]))

    def test_init_n_categories(self):
        con = ExpConcrete(1., tf.ones([10]))
        self.assertTrue(isinstance(con.n_categories, int))
        self.assertEqual(con.n_categories, 10)

        with self.test_session(use_gpu=True):
            logits = tf.placeholder(tf.float32, None)
            con2 = ExpConcrete(1., logits)
            self.assertEqual(
                con2.n_categories.eval(feed_dict={logits: np.ones([10])}), 10)
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should have rank"):
                con2.n_categories.eval(feed_dict={logits: 1.})

    def test_init_temperature(self):
        with self.assertRaisesRegexp(ValueError,
                                     "should be a scalar"):
            ExpConcrete([1.], [1., 2.])

        with self.test_session(use_gpu=True):
            temperature = tf.placeholder(tf.float32, None)
            con = ExpConcrete(temperature, [1., 2.])
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should be a scalar"):
                con.temperature.eval(feed_dict={temperature: [1.]})

    def test_value_shape(self):
        # static
        con = ExpConcrete(1., tf.placeholder(tf.float32, [None, 10]))
        self.assertEqual(con.get_value_shape().as_list(), [10])

        # dynamic
        logits = tf.placeholder(tf.float32, None)
        con2 = ExpConcrete(1., logits)
        self.assertTrue(con2._value_shape().dtype is tf.int32)
        with self.test_session(use_gpu=True):
            self.assertEqual(con2._value_shape().eval(
                feed_dict={logits: np.ones([2, 1, 3])}).tolist(), [3])

        self.assertEqual(con._value_shape().dtype, tf.int32)

    def test_batch_shape(self):
        def _proxy_distribution(logits):
            return ExpConcrete(1., logits)
        utils.test_batch_shape_1parameter(
            self, _proxy_distribution, np.zeros, is_univariate=False)

    def test_sample_shape(self):
        def _proxy_distribution(logits):
            return ExpConcrete(1., logits)
        utils.test_1parameter_sample_shape_one_rank_less(
            self, _proxy_distribution, np.zeros)

    def test_log_prob_shape(self):
        def _proxy_distribution(logits):
            return ExpConcrete(1., logits)

        def _make_samples(shape):
            samples = np.ones(shape, dtype=np.float32)
            return np.log(samples / samples.sum(axis=-1, keepdims=True))

        utils.test_1parameter_log_prob_shape_one_rank_less(
            self, _proxy_distribution, np.ones, _make_samples)

    def test_value(self):
        with self.test_session(use_gpu=True):
            def _test_value(given, temperature, logits):
                given = np.array(given, np.float32)
                logits = np.array(logits, np.float32)
                n = logits.shape[-1]

                t = temperature
                target_log_p = gammaln(n) + (n - 1) * np.log(t) + \
                    (logits - t * given).sum(axis=-1) - \
                    n * np.log(np.exp(logits - t * given).sum(axis=-1))

                con = ExpConcrete(temperature, logits=logits)
                log_p = con.log_prob(given)
                self.assertAllClose(log_p.eval(), target_log_p)
                p = con.prob(given)
                self.assertAllClose(p.eval(), np.exp(target_log_p))

            _test_value([np.log(0.25), np.log(0.25), np.log(0.5)],
                        0.1,
                        [1., 1., 1.2])
            _test_value([[np.log(0.25), np.log(0.25), np.log(0.5)],
                        [np.log(0.1), np.log(0.5), np.log(0.4)]],
                        0.5,
                        [[1., 1., 1.], [.5, .5, .4]])

    def test_dtype(self):
        utils.test_dtype_2parameter(self, ExpConcrete)

    def test_sample_reparameterized(self):
        temperature = tf.constant(1.0)
        logits = tf.ones([2, 3])
        con_rep = ExpConcrete(temperature, logits)
        samples = con_rep.sample(tf.placeholder(tf.int32, shape=[]))
        t_grads, logits_grads = tf.gradients(samples, [temperature, logits])
        self.assertTrue(t_grads is not None)
        self.assertTrue(logits_grads is not None)

        con_no_rep = ExpConcrete(temperature, logits, is_reparameterized=False)
        samples = con_no_rep.sample(tf.placeholder(tf.int32, shape=[]))
        t_grads, logits_grads = tf.gradients(samples, [temperature, logits])
        self.assertEqual(t_grads, None)
        self.assertEqual(logits_grads, None)

    def test_path_derivative(self):
        temperature = tf.constant(1.0)
        logits = tf.ones([2, 3])
        n_samples = tf.placeholder(tf.int32, shape=[])

        con_rep = ExpConcrete(temperature, logits, use_path_derivative=True)
        samples = con_rep.sample(n_samples)
        log_prob = con_rep.log_prob(samples)
        t_path_grads, logits_path_grads = tf.gradients(log_prob,
                                                       [temperature, logits])
        sample_grads = tf.gradients(log_prob, samples)
        t_true_grads = tf.gradients(samples, temperature, sample_grads)[0]
        logits_true_grads = tf.gradients(samples, logits, sample_grads)[0]
        with self.test_session(use_gpu=True) as sess:
            outs = sess.run([t_path_grads, t_true_grads,
                             logits_path_grads, logits_true_grads],
                            feed_dict={n_samples: 7})
            t_path, t_true, logits_path, logits_true = outs
            self.assertAllClose(t_path, t_true)
            self.assertAllClose(logits_path, logits_true)

        con_no_rep = ExpConcrete(temperature, logits, is_reparameterized=False,
                                 use_path_derivative=True)
        samples = con_no_rep.sample(n_samples)
        log_prob = con_no_rep.log_prob(samples)
        t_path_grads, logits_path_grads = tf.gradients(log_prob,
                                                       [temperature, logits])
        self.assertTrue(t_path_grads is None)
        self.assertTrue(logits_path_grads is None)

    def test_check_numerics(self):
        tau = tf.placeholder(tf.float32, None)
        logits = tf.placeholder(tf.float32, None)
        given = tf.placeholder(tf.float32, None)
        dist = ExpConcrete(tau, logits, check_numerics=True)
        log_p = dist.log_prob(given)
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(temperature\).*Tensor had Inf"):
                log_p.eval(feed_dict={tau: 0., logits: np.ones([2]),
                                      given: [1., 1.]})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(temperature\).*Tensor had NaN"):
                log_p.eval(feed_dict={tau: -1., logits: np.ones([2]),
                                      given: [1., 1.]})

class TestConcrete(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank"):
                Concrete(1., logits=tf.zeros([]))

    def test_init_n_categories(self):
        con = Concrete(1., tf.ones([10]))
        self.assertTrue(isinstance(con.n_categories, int))
        self.assertEqual(con.n_categories, 10)

        with self.test_session(use_gpu=True):
            logits = tf.placeholder(tf.float32, None)
            con2 = Concrete(1., logits)
            self.assertEqual(
                con2.n_categories.eval(feed_dict={logits: np.ones([10])}), 10)
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should have rank"):
                con2.n_categories.eval(feed_dict={logits: 1.})

    def test_init_temperature(self):
        with self.assertRaisesRegexp(ValueError,
                                     "should be a scalar"):
            Concrete([1.], [1., 2.])

        with self.test_session(use_gpu=True):
            temperature = tf.placeholder(tf.float32, None)
            con = Concrete(temperature, [1., 2.])
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "should be a scalar"):
                con.temperature.eval(feed_dict={temperature: [1.]})

    def test_value_shape(self):
        # static
        con = Concrete(1., tf.placeholder(tf.float32, [None, 10]))
        self.assertEqual(con.get_value_shape().as_list(), [10])

        # dynamic
        logits = tf.placeholder(tf.float32, None)
        con2 = Concrete(1., logits)
        self.assertTrue(con2._value_shape().dtype is tf.int32)
        with self.test_session(use_gpu=True):
            self.assertEqual(con2._value_shape().eval(
                feed_dict={logits: np.ones([2, 1, 3])}).tolist(), [3])

        self.assertEqual(con._value_shape().dtype, tf.int32)

    def test_batch_shape(self):
        def _proxy_distribution(logits):
            return Concrete(1., logits)
        utils.test_batch_shape_1parameter(
            self, _proxy_distribution, np.zeros, is_univariate=False)

    def test_sample_shape(self):
        def _proxy_distribution(logits):
            return Concrete(1., logits)
        utils.test_1parameter_sample_shape_one_rank_less(
            self, _proxy_distribution, np.zeros)

    def test_log_prob_shape(self):
        def _proxy_distribution(logits):
            return Concrete(1., logits)

        def _make_samples(shape):
            samples = np.ones(shape, dtype=np.float32)
            return np.log(samples / samples.sum(axis=-1, keepdims=True))

        utils.test_1parameter_log_prob_shape_one_rank_less(
            self, _proxy_distribution, np.ones, _make_samples)

    def test_value(self):
        with self.test_session(use_gpu=True):
            def _test_value(given, temperature, logits):
                given = np.array(given, np.float32)
                logits = np.array(logits, np.float32)
                n = logits.shape[-1]

                t = temperature
                target_log_p = gammaln(n) + (n - 1) * np.log(t) + \
                    (logits - (t + 1) * np.log(given)).sum(axis=-1) - \
                    n * np.log(np.exp(logits - t * np.log(given)).sum(axis=-1))

                con = Concrete(temperature, logits=logits)
                log_p = con.log_prob(given)
                self.assertAllClose(log_p.eval(), target_log_p)
                p = con.prob(given)
                self.assertAllClose(p.eval(), np.exp(target_log_p))

            _test_value([0.25, 0.25, 0.5],
                        0.1,
                        [1., 1., 1.2])
            _test_value([[0.25, 0.25, 0.5],
                        [0.1, 0.5, 0.4]],
                        0.5,
                        [[1., 1., 1.], [.5, .5, .4]])

    def test_dtype(self):
        utils.test_dtype_2parameter(self, Concrete)

    def test_sample_reparameterized(self):
        temperature = tf.constant(1.0)
        logits = tf.ones([2, 3])
        con_rep = Concrete(temperature, logits)
        samples = con_rep.sample(tf.placeholder(tf.int32, shape=[]))
        t_grads, logits_grads = tf.gradients(samples, [temperature, logits])
        self.assertTrue(t_grads is not None)
        self.assertTrue(logits_grads is not None)

        con_no_rep = Concrete(temperature, logits, is_reparameterized=False)
        samples = con_no_rep.sample(tf.placeholder(tf.int32, shape=[]))
        t_grads, logits_grads = tf.gradients(samples, [temperature, logits])
        self.assertEqual(t_grads, None)
        self.assertEqual(logits_grads, None)

    def test_path_derivative(self):
        temperature = tf.constant(1.0)
        logits = tf.ones([2, 3])
        n_samples = tf.placeholder(tf.int32, shape=[])

        con_rep = Concrete(temperature, logits, use_path_derivative=True)
        samples = con_rep.sample(n_samples)
        log_prob = con_rep.log_prob(samples)
        t_path_grads, logits_path_grads = tf.gradients(log_prob,
                                                       [temperature, logits])
        sample_grads = tf.gradients(log_prob, samples)
        t_true_grads = tf.gradients(samples, temperature, sample_grads)[0]
        logits_true_grads = tf.gradients(samples, logits, sample_grads)[0]
        with self.test_session(use_gpu=True) as sess:
            outs = sess.run([t_path_grads, t_true_grads,
                             logits_path_grads, logits_true_grads],
                            feed_dict={n_samples: 7})
            t_path, t_true, logits_path, logits_true = outs
            self.assertAllClose(t_path, t_true)
            self.assertAllClose(logits_path, logits_true)

        con_no_rep = Concrete(temperature, logits, is_reparameterized=False,
                              use_path_derivative=True)
        samples = con_no_rep.sample(n_samples)
        log_prob = con_no_rep.log_prob(samples)
        t_path_grads, logits_path_grads = tf.gradients(log_prob,
                                                       [temperature, logits])
        self.assertTrue(t_path_grads is None)
        self.assertTrue(logits_path_grads is None)

    def test_check_numerics(self):
        tau = tf.placeholder(tf.float32, None)
        logits = tf.placeholder(tf.float32, None)
        given = tf.placeholder(tf.float32, None)
        dist = Concrete(tau, logits, check_numerics=True)
        log_p = dist.log_prob(given)
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(given\).*Tensor had Inf"):
                log_p.eval(feed_dict={tau: 1., logits: np.ones([2]),
                                      given: [0., 1.]})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(given\).*Tensor had NaN"):
                log_p.eval(feed_dict={tau: 1., logits: np.ones([2]),
                                      given: [1., -1.]})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(temperature\).*Tensor had Inf"):
                log_p.eval(feed_dict={tau: 0., logits: np.ones([2]),
                                      given: [1., 1.]})
            with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                         "log\(temperature\).*Tensor had NaN"):
                log_p.eval(feed_dict={tau: -1., logits: np.ones([2]),
                                      given: [1., 1.]})

class TestMatrixVariateNormalCholesky(tf.test.TestCase):
    def test_init_check_shape(self):
        with self.test_session(use_gpu=True):
            with self.assertRaisesRegexp(ValueError, "should have rank >= 2"):
                MatrixVariateNormalCholesky(
                    tf.zeros([]), tf.zeros([]), tf.zeros([]))
            with self.assertRaisesRegexp(ValueError, "should have rank >= 2"):
                MatrixVariateNormalCholesky(
                    tf.zeros([1, 2]), tf.zeros([1]), tf.zeros([2, 2]))
            with self.assertRaisesRegexp(ValueError, "should have rank >= 2"):
                MatrixVariateNormalCholesky(
                    tf.zeros([1, 2]), tf.zeros([1, 1]), tf.zeros([1]))
            with self.assertRaisesRegexp(ValueError, 'compatible'):
                MatrixVariateNormalCholesky(
                    tf.zeros([1, 2, 3]),
                    tf.placeholder(tf.float32, [1, 3, 3]),
                    tf.placeholder(tf.float32, [1, 3, 3]))
            with self.assertRaisesRegexp(ValueError, 'compatible'):
                MatrixVariateNormalCholesky(
                    tf.zeros([1, 2, 3]),
                    tf.placeholder(tf.float32, [1, 2, 2]),
                    tf.placeholder(tf.float32, [1, 2, 2]))
            with self.assertRaisesRegexp(ValueError, 'compatible'):
                MatrixVariateNormalCholesky(
                    tf.zeros([2, 3]),
                    tf.placeholder(tf.float32, [1, 2, 2]),
                    tf.placeholder(tf.float32, [1, 3, 3]))
            u = tf.placeholder(tf.float32, [None])
            len_u = tf.shape(u)[0]
            v = tf.placeholder(tf.float32, [None])
            len_v = tf.shape(v)[0]
            dst = MatrixVariateNormalCholesky(
                tf.zeros([2, 3]), tf.zeros([len_u, len_u]),
                tf.zeros([len_v, len_v]))
            with self.assertRaisesRegexp(
                    tf.errors.InvalidArgumentError, 'compatible'):
                dst.sample().eval(
                    feed_dict={u: np.ones((3,)), v: np.ones((3,))})
                dst.sample().eval(
                    feed_dict={u: np.ones((2,)), v: np.ones((2,))})

    def test_shape_inference(self):
        with self.test_session(use_gpu=True):
            # Static
            mean = 10 * np.random.normal(size=(10, 11, 2, 3)).astype('d')
            u_tril = np.zeros((10, 11, 2, 2))
            v_tril = np.zeros((10, 11, 3, 3))
            dst = MatrixVariateNormalCholesky(
                tf.constant(mean), tf.constant(u_tril), tf.constant(v_tril))
            self.assertEqual(dst.get_batch_shape().as_list(), [10, 11])
            self.assertEqual(dst.get_value_shape().as_list(), [2, 3])
            # Dynamic
            unk_mean = tf.placeholder(tf.float32, None)
            unk_u_tril = tf.placeholder(tf.float32, None)
            unk_v_tril = tf.placeholder(tf.float32, None)
            dst = MatrixVariateNormalCholesky(unk_mean, unk_u_tril, unk_v_tril)
            self.assertEqual(dst.get_value_shape().as_list(), [None, None])
            feed_dict = {unk_mean: np.ones((2, 3)), unk_u_tril: np.eye(2),
                         unk_v_tril: np.eye(3)}
            self.assertEqual(list(dst.batch_shape.eval(feed_dict)), [])
            self.assertEqual(list(dst.value_shape.eval(feed_dict)), [2, 3])

    def _gen_test_params(self, seed):
        np.random.seed(seed)
        mean = 10 * np.random.normal(size=(10, 11, 2, 3)).astype('d')
        u = np.zeros((10, 11, 2, 2))
        v = np.zeros((10, 11, 3, 3))
        u_chol = np.zeros_like(u)
        v_chol = np.zeros_like(v)
        for i in range(10):
            for j in range(11):
                u[i, j] = stats.invwishart.rvs(2, np.eye(2))
                u[i, j] /= np.max(np.diag(u[i, j]))
                u_chol[i, j, :, :] = np.linalg.cholesky(u[i, j])

                v[i, j] = stats.invwishart.rvs(3, np.eye(3))
                v[i, j] /= np.max(np.diag(v[i, j]))
                v_chol[i, j, :, :] = np.linalg.cholesky(v[i, j])
        return mean, u, u_chol, v, v_chol

    @contextmanager
    def fixed_randomness_session(self, seed):
        with tf.Graph().as_default() as g:
            with self.test_session(use_gpu=True, graph=g):
                tf.set_random_seed(seed)
                yield

    def test_sample(self):
        with self.fixed_randomness_session(233):
            def test_sample_with(seed):
                mean, u, u_chol, v, v_chol = self._gen_test_params(seed)
                dst = MatrixVariateNormalCholesky(
                    tf.constant(mean), tf.constant(u_chol),
                    tf.constant(v_chol))
                n_exp = 20000
                samples = dst.sample(n_exp)
                sample_shape = (n_exp, 10, 11, 2, 3)
                self.assertEqual(samples.shape.as_list(), list(sample_shape))
                samples = dst.sample(n_exp).eval()
                self.assertEqual(samples.shape, sample_shape)
                self.assertAllClose(
                    np.mean(samples, axis=0), mean, rtol=5e-2, atol=5e-2)
                samples = np.reshape(samples.transpose([0, 1, 2, 4, 3]),
                                     [n_exp, 10, 11, -1])
                for i in range(10):
                    for j in range(11):
                        for k in range(3):
                            self.assertAllClose(np.cov(samples[:, i, j, :].T),
                                                np.kron(v[i, j], u[i, j]),
                                                rtol=1e-1, atol=1e-1)

            for seed in [23, 233, 2333]:
                test_sample_with(seed)

    def test_prob(self):
        with self.fixed_randomness_session(233):
            def test_prob_with(seed):
                mean, u, u_chol, v, v_chol = self._gen_test_params(seed)
                dst = MatrixVariateNormalCholesky(
                    tf.constant(mean), tf.constant(u_chol),
                    tf.constant(v_chol), check_numerics=True)
                n_exp = 200
                samples = dst.sample(n_exp).eval()
                log_pdf = dst.log_prob(tf.constant(samples))
                pdf_shape = (n_exp, 10, 11)
                self.assertEqual(log_pdf.shape.as_list(), list(pdf_shape))
                log_pdf = log_pdf.eval()
                self.assertEqual(log_pdf.shape, pdf_shape)
                for i in range(10):
                    for j in range(11):
                        log_pdf_exact = stats.matrix_normal.logpdf(
                            samples[:, i, j, :], mean[i, j],
                            u[i, j], v[i, j])
                        self.assertAllClose(
                            log_pdf_exact, log_pdf[:, i, j])
                self.assertAllClose(
                    np.exp(log_pdf), dst.prob(tf.constant(samples)).eval())

            for seed in [23, 233, 2333]:
                test_prob_with(seed)

    def test_sample_reparameterized(self):
        mean, u, u_chol, v, v_chol = self._gen_test_params(23)
        mean = tf.constant(mean)
        u_chol = tf.constant(u_chol)
        v_chol = tf.constant(v_chol)
        mvn_rep = MatrixVariateNormalCholesky(mean, u_chol, v_chol)
        samples = mvn_rep.sample(tf.placeholder(tf.int32, shape=[]))
        mean_grads, u_grads, v_grads = tf.gradients(
            samples, [mean, u_chol, v_chol])
        self.assertTrue(mean_grads is not None)
        self.assertTrue(u_grads is not None)
        self.assertTrue(v_grads is not None)
        mvn_rep = MatrixVariateNormalCholesky(mean, u_chol, v_chol,
                                              is_reparameterized=False)
        samples = mvn_rep.sample(tf.placeholder(tf.int32, shape=[]))
        mean_grads, u_grads, v_grads = tf.gradients(
            samples, [mean, u_chol, v_chol])
        self.assertEqual(mean_grads, None)
        self.assertEqual(u_grads, None)
        self.assertEqual(v_grads, None)

max_stars_repo_name: Li-Scottech/zhusuan

max_issues_repo_name: Li-Scottech/zhusuan

max_forks_repo_name: Li-Scottech/zhusuan


text:
# -*- coding: utf-8 -*-
"""
Created on Thu Dec  9 17:25:44 2021

@author: s1253
"""

import os
import tensorflow as tf
import numpy as np
import tensorflow_datasets as tfds
from tensorflow.keras.models import Model
from tensorflow.keras.utils import plot_model
import tensorflow.keras.backend as k

from ops import EncoderLayer, DecoderLayer, preprocess, create_masks
from tensorflow.keras.layers import Dense, Embedding, Dropout, Input

output = "transformer"

output_dir = os.path.join(output, "data")
en_vocab_file = os.path.join("data", "en_vocab")
zh_vocab_file = os.path.join("data", "zh_vocab")

if not os.path.exists(output_dir):
  os.makedirs(output_dir)

config = tfds.translate.wmt.WmtConfig(
  version=tfds.core.Version('0.0.1'),
  language_pair=("zh", "en"),
  subsets={
    tfds.Split.TRAIN: ["newscommentary_v14"]
  }
)

# loading data

builder = tfds.builder("wmt_translate", config=config)
builder.download_and_prepare(download_dir=output_dir)

# using 70% of data in the dataset
examples = builder.as_dataset(split=['train[:70%]', 'train[70%:71%]','train[71%:]'], as_supervised=True)

train_examples, val_examples, _ = examples

# building english encode table

try:
  subword_encoder_en = tfds.deprecated.text.SubwordTextEncoder.load_from_file(en_vocab_file)
  print(f"載入已建立的字典： {en_vocab_file}")
except:
  print("沒有已建立的字典，從頭建立。")
  subword_encoder_en = tfds.deprecated.text.SubwordTextEncoder.build_from_corpus(
      (en.numpy() for en, _ in train_examples),
      target_vocab_size=2**15)

  subword_encoder_en.save_to_file(en_vocab_file)

print(f"字典大小：{subword_encoder_en.vocab_size}")

# building chinese encode table

try:
  subword_encoder_zh = tfds.deprecated.text.SubwordTextEncoder.load_from_file(zh_vocab_file)
  print(f"載入已建立的字典： {zh_vocab_file}")
except:
  print("沒有已建立的字典，從頭建立。")
  subword_encoder_zh = tfds.deprecated.text.SubwordTextEncoder.build_from_corpus(
      (zh.numpy() for _, zh in train_examples),
      target_vocab_size=2**15,
      max_subword_length=1)

  subword_encoder_zh.save_to_file(zh_vocab_file)

print(f"字典大小：{subword_encoder_zh.vocab_size}")

def encode(en_t, zh_t):
  en_indices = [subword_encoder_en.vocab_size] + subword_encoder_en.encode(
      en_t.numpy()) + [subword_encoder_en.vocab_size + 1]

  zh_indices = [subword_encoder_zh.vocab_size] + subword_encoder_zh.encode(
      zh_t.numpy()) + [subword_encoder_zh.vocab_size + 1]

  return en_indices, zh_indices

def tf_encode(en, zh):
    return tf.py_function(encode, [en, zh], [tf.int64 , tf.int64])

MAX_LENGTH = 35
BATCH_SIZE = 64
BUFFER_SIZE = 200

def filter_max_length(en, zh, max_length=MAX_LENGTH):

  return tf.logical_and(tf.size(en) <= max_length,
                        tf.size(zh) <= max_length)

def map_inp(inp, tar):
    tar_inp = tar[:, :-1]
    tar_real = tar[:, 1:]

    return (inp, tar_inp), tar_real

train_dataset = (train_examples
                 .map(tf_encode)
                 .cache()
                 .filter(filter_max_length)
                 .shuffle(BUFFER_SIZE)
                 .padded_batch(batch_size=BATCH_SIZE,padded_shapes=([-1], [-1]))
                 .prefetch(tf.data.experimental.AUTOTUNE)
                 .map(map_inp)
                 )

val_dataset = (val_examples
               .map(tf_encode)
               .filter(filter_max_length)
               .padded_batch(BATCH_SIZE, padded_shapes=([-1], [-1]))
               .map(map_inp)
               )

#%%

# build model

num_layers = 4
d_model = 512
dff = 2048
num_heads = 8
dropout_rate = 0.2

input_vocab_size = subword_encoder_en.vocab_size + 2
target_vocab_size = subword_encoder_zh.vocab_size + 2

encoder_inp = Input((None,))
decoder_inp = Input((None,))
''' ---------------------------------------------------------------- '''
#encoder

enc_padding_mask, combined_mask, dec_padding_mask = create_masks()(encoder_inp, decoder_inp)

enc_len = tf.shape(encoder_inp)[1]

x = Embedding(input_vocab_size, d_model)(encoder_inp)
x = preprocess()(x, enc_len, d_model)
x = Dropout(dropout_rate)(x)

for _ in range(num_layers):
    enc_lay = EncoderLayer(d_model, num_heads, dff, dropout_rate)
    x = enc_lay(x, mask=enc_padding_mask)

enc_out = x

''' ---------------------------------------------------------------- '''
# decoder

attention_weights = {}

dec_len = tf.shape(decoder_inp)[1]

x = Embedding(input_vocab_size, d_model)(decoder_inp)
x = preprocess()(x, dec_len, d_model)
x = Dropout(dropout_rate)(x)

for i in range(num_layers):
    dec_lay = DecoderLayer(d_model, num_heads, dff, dropout_rate)
    x, block1, block2= dec_lay(x, enc_out, combined_mask, dec_padding_mask)

    attention_weights['decoder_layer{}_block1'.format(i + 1)] = block1
    attention_weights['decoder_layer{}_block2'.format(i + 1)] = block2

dec_out = x

''' ---------------------------------------------------------------- '''

Transformer_out = Dense(target_vocab_size)(x)

Transformer = Model([encoder_inp, decoder_inp], Transformer_out)

plot_model(Transformer, r"result/model_structure.png", show_shapes=True, show_layer_names=True)

#%%

# training

from tensorflow.keras.callbacks import ModelCheckpoint

# save model

checkpoint = ModelCheckpoint(r"model/origin/weights{epoch:02d}-{loss:.2f}.h5", monitor='loss', verbose=1, save_best_only=True,
mode='min')

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')

def loss_function(real, pred):
  mask = tf.math.logical_not(tf.math.equal(real, 0))

  loss_ = loss_object(real, pred)
  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask

  return tf.reduce_mean(loss_)

optimizer = tf.keras.optimizers.Adam(0.0001, beta_1=0.9, beta_2=0.98,
                                     epsilon=1e-9)

Transformer.compile(loss=loss_function, optimizer=optimizer, metrics=["accuracy"])

History = Transformer.fit(train_dataset
                        , validation_data=val_dataset
                        , epochs=50
                        , callbacks=checkpoint
                        , shuffle=True)

attention_weights_fn = k.function([encoder_inp, decoder_inp],[attention_weights])

import matplotlib.pyplot as plt

# visualize loss and accuracy

arr = History.history['accuracy'][-1]
val = History.history['val_accuracy'][-1]

plt.figure(figsize = (15,5))
plt.subplot(1,2,1)
plt.plot(History.history['accuracy'])
plt.plot(History.history['val_accuracy'])
plt.title('model accuracy')
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')

plt.subplot(1,2,2)
plt.plot(History.history['loss'])
plt.plot(History.history['val_loss'])
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')

plt.savefig(r"result/accracy_loss.png")

#%%

# evaluate

Transformer.load_weights(r"model/origin/weights50-0.41.h5")
Transformer.summary()

def evaluate(inp_sentence):

    start_token = [subword_encoder_en.vocab_size]
    end_token = [subword_encoder_en.vocab_size + 1]

    inp_sentence = start_token + subword_encoder_en.encode(inp_sentence) + end_token
    encoder_input = tf.expand_dims(inp_sentence, 0)

    decoder_input = [subword_encoder_zh.vocab_size]
    output = tf.expand_dims(decoder_input, 0)

    for i in range(MAX_LENGTH):

        predictions = Transformer.predict((encoder_input, output))
        attention_weights = attention_weights_fn((encoder_input, output))[0]

        predictions = predictions[: , -1:, :]  # (batch_size, 1, vocab_size)

        predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)

        if tf.equal(predicted_id, subword_encoder_zh.vocab_size + 1):
          return tf.squeeze(output, axis=0), attention_weights

        output = tf.concat([output, predicted_id], axis=-1)

    return tf.squeeze(output, axis=0), attention_weights

sentence = "In recent years, Taiwan’s economy is growing rapidly"

predicted_seq,  attention_weights_after = evaluate(sentence)

attention_weights_after = attention_weights_after[0]

target_vocab_size = subword_encoder_zh.vocab_size

predicted_seq = [idx for idx in predicted_seq if idx < target_vocab_size]

predicted_sentence = subword_encoder_zh.decode(predicted_seq)

print("sentence:", sentence)
print("-" * 20)

print("predicted_seq:", predicted_seq)
print("-" * 20)

print("predicted_sentence:", predicted_sentence)

import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']

# visualize maltihead attention

def plot_attention_weights(attention_weights, sentence, predicted_seq, layer_name, max_len_tar=None):

    fig = plt.figure(figsize=(17, 7))

    sentence = subword_encoder_en.encode(sentence)

    max_len_tar = len(predicted_seq)

    attention_weights = tf.squeeze(attention_weights[layer_name], axis=0)
    # (num_heads, tar_seq_len, inp_seq_len)

    for head in range(attention_weights.shape[0]):
        ax = fig.add_subplot(2, 4, head + 1)

        attn_map = np.transpose(attention_weights[head][:max_len_tar, :])
        ax.matshow(attn_map, cmap='viridis')  # (inp_seq_len, tar_seq_len)

        ax.set_xticks(range(len(predicted_seq)))
        ax.set_xlim(-0.5, max_len_tar -1.5)

        ax.set_yticks(range(len(sentence) + 2))
        ax.set_xticklabels([subword_encoder_zh.decode([i]) for i in predicted_seq
                            if i < subword_encoder_zh.vocab_size], fontsize=12)

        ax.set_yticklabels(
            ['<start>'] + [subword_encoder_en.decode([i]) for i in sentence] + ['<end>'],
            )

        ax.set_xlabel('Head {}'.format(head + 1))
        ax.tick_params(axis="x", labelsize=12)
        ax.tick_params(axis="y", labelsize=12)

    plt.tight_layout()
    plt.savefig(r"result/attention_weights.png")
    plt.show()
    plt.close(fig)

layer_name = f"decoder_layer{num_layers}_block2"

plot_attention_weights(attention_weights_after, sentence, predicted_seq, layer_name, max_len_tar=18)

ext: py

reason: import numpy


text:
'''
Copyright (C) 2010-2021 Alibaba Group Holding Limited.
'''

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from collections import OrderedDict

import numpy as np

import os
import sys
from collections import OrderedDict

sys.path.append(['../../', '../'])
from utils import load_pretrained_checkpoint, load_checkpoint
import logging
from .DSN_Fusion import DSNNet

class LabelSmoothingCrossEntropy(torch.nn.Module):
    def __init__(self, smoothing: float = 0.1,
                 reduction="mean", weight=None):
        super(LabelSmoothingCrossEntropy, self).__init__()
        self.smoothing = smoothing
        self.reduction = reduction
        self.weight = weight

    def reduce_loss(self, loss):
        return loss.mean() if self.reduction == 'mean' else loss.sum() \
         if self.reduction == 'sum' else loss

    def linear_combination(self, x, y):
        return self.smoothing * x + (1 - self.smoothing) * y

    def forward(self, preds, target):
        assert 0 <= self.smoothing < 1

        if self.weight is not None:
            self.weight = self.weight.to(preds.device)

        n = preds.size(-1)
        log_preds = F.log_softmax(preds, dim=-1)
        loss = self.reduce_loss(-log_preds.sum(dim=-1))
        nll = F.nll_loss(
            log_preds, target, reduction=self.reduction, weight=self.weight
        )
        return self.linear_combination(loss / n, nll)

class Encoder(nn.Module):
    def __init__(self, C_in, C_out, dilation=2):
        super(Encoder, self).__init__()
        self.enconv = nn.Sequential(
            nn.Conv2d(C_in, C_in, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(C_in),
            nn.ReLU(inplace=False),

            nn.Conv2d(C_in, C_in // 2, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(C_in // 2),
            nn.ReLU(inplace=False),

            nn.Conv2d(C_in // 2, C_in // 4, kernel_size=1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(C_in // 4),
            nn.ReLU(inplace=False),

            nn.Conv2d(C_in // 4, C_out, kernel_size=1, stride=1, padding=0, bias=False),
        )

    def forward(self, x1, x2):
        b, c = x1.shape
        x = torch.cat((x1, x2), dim=1).view(b, -1, 1, 1)
        x = self.enconv(x)
        return x

class Decoder(nn.Module):
    def __init__(self, C_in, C_out, dilation=2):
        super(Decoder, self).__init__()
        self.deconv = nn.Sequential(
            nn.Conv2d(C_in, C_out // 4, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(C_out // 4),
            nn.ReLU(),

            nn.Conv2d(C_out // 4, C_out // 2, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(C_out // 2),
            nn.ReLU(),
        )

    def forward(self, x):
        x = self.deconv(x)
        return x

class FusionModule(nn.Module):
    def __init__(self, channel_in=1024, channel_out=256, num_classes=60):
        super(FusionModule, self).__init__()
        self.encoder = Encoder(channel_in, channel_out)
        self.decoder = Decoder(channel_out, channel_in)
        self.efc = nn.Conv2d(channel_out, num_classes, kernel_size=1, padding=0, bias=False)

    def forward(self, r, d):
        en_x = self.encoder(r, d)  # [4, 256, 1, 1]
        de_x = self.decoder(en_x)
        en_x = self.efc(en_x)
        return en_x.squeeze(), de_x

class CrossFusionNet(nn.Module):
    def __init__(self, args, num_classes, pretrained, spatial_interact=True, temporal_interact=True):
        super(CrossFusionNet, self).__init__()
        self._MES = torch.nn.MSELoss()
        self._BCE = torch.nn.BCELoss()
        self._CE = LabelSmoothingCrossEntropy()
        self.spatial_interact = spatial_interact
        self.temporal_interact = temporal_interact

        self.fusion_model = FusionModule(channel_out=256, num_classes=num_classes)
        self.avg_pool = nn.AdaptiveAvgPool3d(1)
        self.fc = nn.Conv2d(512, 1, kernel_size=1, padding=0, bias=False)
        self.dropout = nn.Dropout(0.5)

        assert args.rgb_checkpoint and args.depth_checkpoint
        self.Modalit_rgb = DSNNet(args, num_classes=num_classes,
                                     pretrained=args.rgb_checkpoint)

        self.Modalit_depth = DSNNet(args, num_classes=num_classes,
                                       pretrained=args.depth_checkpoint)

        if self.spatial_interact:
            self.crossFusion = nn.Sequential(
                nn.Conv2d(512 * 2, 512, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(512),
                nn.ReLU(),
                nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, bias=False),
                nn.Dropout(0.4)

            )
        if self.temporal_interact:
            self.crossFusionT = nn.Sequential(
                nn.Conv2d(512 * 2, 512, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(512),
                nn.ReLU(),
                nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0, bias=False),
                nn.Dropout(0.4)
            )

        self.classifier1 = nn.Sequential(
            nn.LayerNorm(512),
            nn.Linear(512, num_classes)
        )
        self.classifier2 = nn.Sequential(
            nn.LayerNorm(512),
            nn.Linear(512, num_classes)
        )

        if pretrained:
            load_pretrained_checkpoint(self, pretrained)
            logging.info("Load Pre-trained model state_dict Done !")

    def forward(self, inputs, garrs, target):
        rgb, depth = inputs
        rgb_garr, depth_garr = garrs

        spatial_M = self.Modalit_rgb(rgb, rgb_garr, endpoint='spatial')
        spatial_K = self.Modalit_depth(depth, depth_garr, endpoint='spatial')

        if self.spatial_interact:
            b, t, c = spatial_M.shape
            spatial_fusion_features = self.crossFusion(F.normalize(torch.cat((spatial_M, spatial_K), dim=-1), p = 2, dim=-1).view(b, c*2, t, 1)).squeeze()

        (temporal_M, M_xs, M_xm, M_xl), distillationM, _ = self.Modalit_rgb(x=spatial_M + spatial_fusion_features.view(spatial_M.shape) if self.spatial_interact else spatial_M,
                                                                         endpoint='temporal')  # size[4, 512]
        (temporal_K, K_xs, K_xm, K_xl), distillationK, _ = self.Modalit_depth(x=spatial_K + spatial_fusion_features.view(spatial_M.shape) if self.spatial_interact else spatial_K,
                                                                           endpoint='temporal')
        logit_r = self.classifier1(temporal_M)
        logit_d = self.classifier2(temporal_K)

        if self.temporal_interact:
            b, c = temporal_M.shape
            temporal_fusion_features = self.crossFusionT(F.normalize(torch.cat((temporal_M, temporal_K), dim=-1), p = 2, dim=-1).view(b, c*2, 1, 1)).squeeze()
            temporal_M, temporal_K = temporal_M+temporal_fusion_features, temporal_K+temporal_fusion_features

        en_x, de_x = self.fusion_model(temporal_M, temporal_K)
        b, c = temporal_M.shape
        bce_r = torch.sigmoid(self.fc(self.dropout(temporal_M).view(b, c, 1, 1))).view(b, -1)
        bce_d = torch.sigmoid(self.fc(self.dropout(temporal_K).view(b, c, 1, 1))).view(b, -1)

        BCE_loss = self._BCE(bce_r, torch.ones(bce_r.size(0), 1).cuda()) + self._BCE(bce_d, torch.zeros(bce_d.size(0),
                                                                                1).cuda())
        MSE_loss = self._MES(de_x.view(b, c), temporal_M) + self._MES(de_x.view(b, c), temporal_K)
        CE_loss = self._CE(en_x, target) + self._CE(logit_r, target) + self._CE(logit_d, target)
        distillation = distillationM + distillationK

        return (en_x, logit_r, logit_d), (CE_loss, BCE_loss, MSE_loss, distillation)

max_stars_repo_name: damo-cv/MotionRGBD

max_issues_repo_name: zhoubenjia/MotionRGBD

max_forks_repo_name: zhoubenjia/MotionRGBD


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Воспроизведение видеоданных из сенсора Kinect 2
"""

# ######################################################################################################################
# Заглушка
# ######################################################################################################################

import sys  # Доступ к некоторым переменным и функциям Python

MIN_OS_WINDOWS = (8, 0)

# linux или OS X
if sys.platform == "linux" or sys.platform == "linux2" or sys.platform == "darwin":
    raise RuntimeError("Требуется операционная система Windows >= {}.{}".format(MIN_OS_WINDOWS[0], MIN_OS_WINDOWS[1]))

# ######################################################################################################################
# Импорт необходимых инструментов
# ######################################################################################################################
import time         # Работа со временем
import numpy as np  # Научные вычисления
import cv2          # Алгоритмы компьютерного зрения

from datetime import datetime  # Работа со временем
from _testcapi import USHRT_MAX  # Максимально доступное число для формата ushort

# Персональные
from core2pkgs import config as cfg  # Глобальный файл настроек
from kinect2pvv.core import PyKinectRuntime, PyKinectV2  # Работа с Kinect 2

# ######################################################################################################################
# Сообщения
# ######################################################################################################################
class Messages(cfg.Messages):
    """Класс для сообщений"""

    # ------------------------------------------------------------------------------------------------------------------
    # Конструктор
    # ------------------------------------------------------------------------------------------------------------------

    def __init__(self):
        super().__init__()  # Выполнение конструктора из суперкласса

        self._run_kinect = '[{}] Запуск сенсора Kinect 2 ...'
        self._kinect_not_found = '[{}{}{}] Достигнут лимит ожидания запуска сенсора Kinect 2 в {} секунд ...'

# ######################################################################################################################
# Воспроизведение видеоданных из сенсора Kinect 2
# ######################################################################################################################
class KinectViewer(Messages):
    """Класс для Воспроизведения видеоданных из сенсора Kinect 2"""

    # ------------------------------------------------------------------------------------------------------------------
    # Конструктор
    # ------------------------------------------------------------------------------------------------------------------

    def __init__(self):
        super().__init__()  # Выполнение конструктора из суперкласса

        self._kinect = None  # Kinect 2

        self._wait = 10  # Количество секунд для ожидания включения Kinect 2

        self._cnt_bodies = 0  # Количество найденных скелетов

        self._skeleton_landmarks_color = {}  # Ориентиры скелетов для соединения линиями
        self._skeleton_landmarks_depth = {}  # Ориентиры скелетов карт глубины

    # ------------------------------------------------------------------------------------------------------------------
    # Свойства
    # ------------------------------------------------------------------------------------------------------------------

    # Получение Kinect 2
    @property
    def kinect(self):
        return self._kinect

    # ------------------------------------------------------------------------------------------------------------------
    #  Внутренние методы
    # ------------------------------------------------------------------------------------------------------------------

    # Извлечение кординат скелетной точки
    def __draw_body_bone_color(self, joint_points, joints, joint):
        """
        Извлечение кординат скелетной точки

        TODO: Описание
        """

        # TODO: Проверка аргументов

        joint_state = joints[joint].TrackingState  # Отслеживание состояние сустава

        # Координаты сустава не отслежены
        if joint_state == PyKinectV2.TrackingState_NotTracked or joint_state == PyKinectV2.TrackingState_Inferred:
            return None

        # Координаты сустава
        return {
            'x': int(joint_points[joint].x),
            'y': int(joint_points[joint].y)
        }

    # Извлечение кординат скелетной точки из карты глубины
    def __draw_body_bone_depth(self, joint_points, joints, joint):
        """
        Извлечение кординат скелетной точки из карты глубины

        TODO: Описание
        """

        # TODO: Проверка аргументов

        joint_state = joints[joint].TrackingState  # Отслеживание состояние сустава

        # Координаты сустава не отслежены
        if joint_state == PyKinectV2.TrackingState_NotTracked or joint_state == PyKinectV2.TrackingState_Inferred:
            return None

        x = int(joint_points[joint].x)
        y = int(joint_points[joint].y)
        z = int(self._depth_frame[y * self.kinect.depth_frame_desc.Width + x])

        # Координаты сустава
        return {
            'x': x,
            'y': y,
            'z': z
        }

    # ------------------------------------------------------------------------------------------------------------------
    #  Внешние методы
    # ------------------------------------------------------------------------------------------------------------------

    # Запуск
    def start(self, out = True):
        """
        Запуск Kinect 2

        ([bool, bool]) -> bool

        Аргументы:
           out - Печатать процесс выполнения

        Возвращает: True если запуск Kinect произведен, в обратном случае False
        """

        # Проверка аргументов
        if type(out) is not bool:
            return False

        # Вывод сообщения
        if out is True:
            print(self._run_kinect.format(datetime.now().strftime(self._format_time)))

        # Инициализация Kinect 2 на получение цветного изображения, карты глубины и скелетных моделей
        # noinspection PyUnreachableCode
        self._kinect = PyKinectRuntime.PyKinectRuntime(
            PyKinectV2.FrameSourceTypes_Color
            | PyKinectV2.FrameSourceTypes_Body
            | PyKinectV2.FrameSourceTypes_Depth
            | PyKinectV2.FrameSourceTypes_Infrared
            | PyKinectV2.FrameSourceTypes_BodyIndex
        )

        start_time = time.time()  # Отсчет времени выполнения

        # Ожидаем получение информации из Kinect 2
        while True:
            if self._kinect.has_new_color_frame() and \
                    self._kinect.has_new_depth_frame() and \
                    self._kinect.has_new_infrared_frame():
                return True

            end_time = round(time.time() - start_time, 2)  # Конец времени выполнения

            if end_time >= self._wait:
                # Вывод сообщения
                if out is True:
                    print(self._kinect_not_found.format(
                        self.red, datetime.now().strftime(self._format_time), self.end, self._wait
                    ))

                return False

    # Получение цветного кадра из Kinect 2
    def get_color_frame(self):
        """
        Получение цветного кадра из Kinect 2

        () -> numpy.ndarray

        Возвращает: Цветной кадр
        """

        out_frame = self.kinect.get_last_color_frame()  # Получение цветного кадра с Kinect

        # Преобразование кадра в необходимый формат (1920, 1080, RGB)
        out_frame = cv2.cvtColor(
            cv2.cvtColor(out_frame.reshape(1080, 1920, -1).astype(np.uint8), cv2.COLOR_RGBA2RGB), cv2.COLOR_BGR2RGB
        )

        return out_frame  # Результат

    # Получение карты глубины из Kinect 2
    def get_depth_frame(self):
        """
        Получение карты глубины из Kinect 2

        () -> numpy.ndarray

        Возвращает: Карту глубины в цветном формате
        """

        out_frame = self.kinect.get_last_depth_frame()  # Получение карты глубины с Kinect

        # Преобразование карты глубины в необходимый формат (512, 424)
        out_frame = out_frame.reshape(((424, 512))).astype(np.uint16)

        out_frame = cv2.applyColorMap(cv2.convertScaleAbs(out_frame, alpha = 255 / (4500 - 500)), cv2.COLORMAP_JET)

        return out_frame  # Результат

    # Получение инфракрасного кадра из Kinect 2
    def get_infrared_frame(self, norm = 0.32, out = True):
        """
        Получение инфракрасного кадра из Kinect 2

        ([float]) -> numpy.ndarray

        Аргументы:
           norm - Нормализация значений инфракрасной камеры

        Возвращает: Инфракрасный кадр преобразованный под цветной формат
        """

        # Проверка аргументов
        if type(norm) is not float or norm < 0.01 or norm > 1 or type(out) is not bool:
            # Вывод сообщения
            if out is True:
                print(self._invalid_arguments.format(
                    self.red, datetime.now().strftime(self._format_time), self.end,
                    __class__.__name__ + '.' + self.get_infrared_frame.__name__
                ))

            return None

        out_frame = self.kinect.get_last_infrared_frame()  # Получение инфракрасного кадра с Kinect

        # Преобразование инфракрасного кадра в необходимый формат (512, 424)
        out_frame = out_frame.reshape(((424, 512))).astype(np.uint16)

        # Нормализация значений инфракрасной камеры
        if norm < 1.0:
            out_frame = cv2.max(
                0.01, cv2.min(
                    1.0, (out_frame / USHRT_MAX) / norm
                )
            ) * USHRT_MAX

        out_frame = cv2.cvtColor(cv2.convertScaleAbs(out_frame, alpha = 255 / USHRT_MAX), cv2.COLOR_GRAY2RGB)

        return out_frame  # Результат

    # Получение ориентиров скелета
    def get_bodies(self):
        """
        Получение ориентиров скелета

        () ->

        Возвращает: Словари с ориентирами скелетов в цветном формате и карте глубины
        """

        bodies = self.kinect.get_last_body_frame()  # Получение ориентиров скелета

        out_joints_color = {}  # Словарь с найдеными скелетами
        out_joints_depth = {}  # Словарь с найдеными скелетами карты глубины

        self._cnt_bodies = 0  # Обнуление количества найденных скелетов

        # Ориентиры скелета возможно присутствуют
        if bodies is not None:
            # Пройтись по всем возможным скелетам
            for i in range(0, self.kinect.max_body_count):
                body = bodies.bodies[i]  # Скелетная модель
                # Скелетная модель не найдена
                if not body.is_tracked:
                    continue

                self._cnt_bodies += 1  # Увеличение счетчика количества найденных скелетов

                temp_joints_color = {}  # Словарь с найденным скелетом
                temp_joints_depth = {}  # Словарь с найденным скелетом карты глубины

                joints = body.joints  # Скелетная модель

                # Получение скелетной модели для цветного кадра
                joint_points_color = self.kinect.body_joints_to_color_space(joints)

                # Получение скелетной модели для карты глубины
                joint_points_depth = self.kinect.body_joints_to_depth_space(joints)

                # Суставы скелетной модели
                joints_type = {
                    'neck': PyKinectV2.JointType_Neck,                     # Шея
                    'spine_shoulder': PyKinectV2.JointType_SpineShoulder,  # Позвоночник на уровне плеч
                    'spine_mid': PyKinectV2.JointType_SpineMid,            # Позвоночник по центру тела
                    'spine_base': PyKinectV2.JointType_SpineBase,          # Позвоночник на уровне ног (центр таза)
                    'shoulder_right': PyKinectV2.JointType_ShoulderRight,  # Правое плечо
                    'elbow_right': PyKinectV2.JointType_ElbowRight,        # Правый локоть
                    'wrist_right': PyKinectV2.JointType_WristRight,        # Правое запястье
                    'shoulder_left': PyKinectV2.JointType_ShoulderLeft,    # Левое плечо
                    'elbow_left': PyKinectV2.JointType_ElbowLeft,          # Левый локоть
                    'wrist_left': PyKinectV2.JointType_WristLeft,          # Левое запястье
                    'hand_right': PyKinectV2.JointType_HandRight,          # Правая рука
                    'hand_left': PyKinectV2.JointType_HandLeft,            # Левая рука
                    'trumb_right': PyKinectV2.JointType_ThumbRight,        # Большой палец правой руки
                    'hand_tip_right': PyKinectV2.JointType_HandTipRight,   # 4 пальца правой руки
                    'trumb_left': PyKinectV2.JointType_ThumbLeft,          # Большой палец левой руки
                    'hand_tip_left': PyKinectV2.JointType_HandTipLeft,     # 4 пальца левой руки
                    'hip_right': PyKinectV2.JointType_HipRight,            # Правое бедро
                    'knee_right': PyKinectV2.JointType_KneeRight,          # Правое колено
                    'ankle_right': PyKinectV2.JointType_AnkleRight,        # Правая лодыжка
                    'foot_right': PyKinectV2.JointType_FootRight,          # Правая нога
                    'hip_left': PyKinectV2.JointType_HipLeft,              # Левое бедро
                    'knee_left': PyKinectV2.JointType_KneeLeft,            # Левое колено
                    'ankle_left': PyKinectV2.JointType_AnkleLeft,          # Левая лодыжка
                    'foot_left': PyKinectV2.JointType_FootLeft,            # Левая нога
                    'head': PyKinectV2.JointType_Head                      # Голова
                }

                # Проход по всем возможным суставам скелетной модели
                for key, val in joints_type.items():
                    # Извлечение кординат скелетной точки
                    temp_joints_color[key] = self.__draw_body_bone_color(joint_points_color, joints, val)

                    # Извлечение кординат скелетной точки карт глубины
                    temp_joints_depth[key] = self.__draw_body_bone_depth(joint_points_depth, joints, val)

                # Добавление скелета в словарь со всеми найденными скелетами
                out_joints_color[self._cnt_bodies] = temp_joints_color  # Цветное изображение
                out_joints_depth[self._cnt_bodies] = temp_joints_depth  # Карта глубины

        # Ориентиры скелетов
        self._skeleton_landmarks_color = out_joints_color  # Цветная камера
        self._skeleton_landmarks_depth = out_joints_depth  # Карта глубины

        print(type(self._skeleton_landmarks_color), self._skeleton_landmarks_depth)

max_stars_repo_name: DmitryRyumin/pkgs

max_issues_repo_name: DmitryRyumin/pkgs

max_forks_repo_name: DmitryRyumin/pkgs