text:
# coding: utf-8

"""
    Tests of the solar radiation model, based on Duffie, J.A., and
    Beckman, W. A., 1974, "Solar energy thermal processes"
"""

from solarpy import *
from numpy import sin, cos, deg2rad, rad2deg, array
from numpy.testing import (assert_equal, assert_almost_equal,
                           assert_array_almost_equal)
from datetime import datetime
import unittest as ut

class Test_b_nday(ut.TestCase):
    """
    Tests B(n) values.
    """
    def test_single_days(self):
        date = datetime(2019, 1, 1)  # Jan 1
        expected_value = 0
        assert_equal(b_nday(date), expected_value)

        date = datetime(2019, 12, 31)  # Dec 31
        expected_value = 6.2659711
        assert_almost_equal(b_nday(date), expected_value, 6)

    def test_exception(self):
        self.assertRaises(TypeError, b_nday, 6)

class Test_gon(ut.TestCase):
    """
    Tests radiation on a plane normal values.
    """
    def test_min(self):
        # array with all days
        gon_ = array([gon(datetime(datetime.now().year, 1, 1) +
                      timedelta(days=i)) for i in range(365)])
        self.assertTrue((gon_ > 1320).all())

    def test_max(self):
        # array with all days
        gon_ = array([gon(datetime(datetime.now().year, 1, 1) +
                      timedelta(days=i)) for i in range(365)])
        self.assertTrue((gon_ < 1420).all())

    def test_Jan1(self):
        date = datetime(2019, 1, 1)  # January 1
        expected_value = 1415
        self.assertAlmostEqual(gon(date), expected_value, delta=1)

    def test_summerSolstice(self):
        date = datetime(2019, 6, 20)  # June 20
        expected_value = 1322
        self.assertAlmostEqual(gon(date), expected_value, delta=1)

    def test_exception(self):
        self.assertRaises(TypeError, gon, 'a')

class Test_eq_time(ut.TestCase):
    """
    Tests equation of time values.
    """
    def test_min(self):
        # array with all days
        eq_ = array([eq_time(datetime(datetime.now().year, 1, 1) +
                     timedelta(days=i)) for i in range(365)])
        self.assertTrue((eq_ > -15).all())

    def test_max(self):
        # array with all days
        eq_ = array([eq_time(datetime(datetime.now().year, 1, 1) +
                     timedelta(days=i)) for i in range(365)])
        self.assertTrue((eq_ < 17).all())

    def test_Jan1(self):
        date = datetime(2019, 1, 1)  # January 1
        expected_value = -3
        self.assertAlmostEqual(eq_time(date), expected_value, delta=1)

    def test_summerSolstice(self):
        date = datetime(2019, 6, 20)  # June 20
        expected_value = -1
        self.assertAlmostEqual(eq_time(date), expected_value, delta=1)

    def test_exception(self):
        self.assertRaises(TypeError, eq_time, 'a')

class Test_declination(ut.TestCase):
    """
    Tests equation of time values.
    """
    def test_min(self):
        # array with all days
        dec_ = array([declination(datetime(datetime.now().year, 1, 1) +
                      timedelta(days=i)) for i in range(365)])
        self.assertTrue((rad2deg(dec_) > -24).all())

    def test_max(self):
        # array with all days
        dec_ = array([declination(datetime(datetime.now().year, 1, 1) +
                      timedelta(days=i)) for i in range(365)])
        self.assertTrue((rad2deg(dec_) < 24).all())

    def test_Jan1(self):
        date = datetime(2019, 1, 1)  # January 1
        expected_value = deg2rad(-23)
        self.assertAlmostEqual(declination(date), expected_value, 1)

    def test_summerSolstice(self):
        date = datetime(2019, 6, 20)  # June 20
        expected_value = deg2rad(23)
        self.assertAlmostEqual(declination(date), expected_value, 1)

    def test_Feb13(self):  # Example 1.6.1
        date = datetime(2019, 2, 13)  # Feb 13
        expected_value = deg2rad(-14)
        self.assertAlmostEqual(declination(date), expected_value, 1)

    def test_Mar16(self):  # Example 1.6.3
        date = datetime(2019, 3, 16)  # March 16
        expected_value = deg2rad(-2.4)
        self.assertAlmostEqual(declination(date), expected_value, 1)

    def test_exception(self):
        self.assertRaises(TypeError, declination, 12)

class Test_solar_time(ut.TestCase):
    """
    Tests solar time function. Values from Duffie and Beckman example 1.5.1
    """
    def test_Feb3(self):
        lng = 89.4
        date = datetime(2019, 2, 3, 10, 30)  # standard time 10:30
        expected_value = datetime(2019, 2, 3, 10, 18, 54)
        # compare ignoring microseconds (as in the example)
        self.assertEqual(standard2solar_time(date, lng).replace(microsecond=0),
                         expected_value)

    def test_exception_date(self):
        self.assertRaises(TypeError, standard2solar_time, 12, 8.3)

    def test_exception_lng(self):
        date = datetime(2019, 2, 13)
        self.assertRaises(ValueError, standard2solar_time, date, -181)
        self.assertRaises(TypeError, standard2solar_time, date, '122')

class Test_hour_angle(ut.TestCase):
    """
    Tests hour angle function. Values from Duffie and Beckman
    """
    def test_examples(self):
        # noon
        date = datetime(2019, 1, 1, 12, 0)
        expected_value = deg2rad(0)
        self.assertEqual(hour_angle(date), expected_value)

        # Example 1.6.1
        date = datetime(2019, 1, 1, 10, 30)
        expected_value = deg2rad(-22.5)
        self.assertEqual(hour_angle(date), expected_value)

        # Example 1.6.2a
        date = datetime(2019, 1, 1, 9, 30)
        expected_value = deg2rad(-37.5)
        self.assertEqual(hour_angle(date), expected_value)

        # Example 1.6.2b
        date = datetime(2019, 1, 1, 18, 30)
        expected_value = deg2rad(97.5)
        self.assertEqual(hour_angle(date), expected_value)

        # Example 1.6.3
        date = datetime(2019, 1, 1, 16, 0)
        expected_value = deg2rad(60)
        self.assertEqual(hour_angle(date), expected_value)

        # Example 1.7.1
        date = datetime(2019, 1, 1, 14, 0)
        expected_value = deg2rad(30)
        self.assertEqual(hour_angle(date), expected_value)

    def test_exception_date(self):
        self.assertRaises(TypeError, hour_angle, 121)

class Test_angle_of_incidence(ut.TestCase):
    """
    Tests angle of incidence function. Values from Duffie and Beckman
    example 1.6.1
    """
    def test_example(self):
        date = datetime(2019, 2, 13, 10, 30)  # Feb 13, 10:30 am (solar)
        lat = 43
        beta = 45
        surf_az = 15
        expected_value = deg2rad(35)
        self.assertAlmostEqual(theta(date, lat, beta, surf_az),
                               expected_value, 2)

    def test_exception_date(self):
        self.assertRaises(TypeError, theta, 121, 1, 1, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, theta, date, 91, 1, 1)
        self.assertRaises(TypeError, theta, date, '91', 0, 0)

class Test_zenith_angle(ut.TestCase):
    """
    Tests zenith angle function. Values from Duffie and Beckman
    """
    def test_examples(self):
        # noon at summer solstice and lat = 23.4º
        date = datetime(2019, 6, 20, 12, 0)  # Jun 20, 12:00 am (solar)
        lat = 23.45
        expected_value = deg2rad(0)
        self.assertAlmostEqual(theta_z(date, lat), expected_value, 2)

        # Example 1.6.2a
        date = datetime(2019, 2, 13, 9, 30)  # Feb 13, 9:30 am (solar)
        lat = 43
        expected_value = deg2rad(66.5)
        self.assertAlmostEqual(theta_z(date, lat), expected_value, 2)

        # Example 1.6.2b
        date = datetime(2019, 7, 1, 18, 30)  # Jul 1, 18:30 am (solar)
        lat = 43
        expected_value = deg2rad(79.6)
        self.assertAlmostEqual(theta_z(date, lat), expected_value, 2)

        # Example 1.6.3
        date = datetime(2019, 3, 16, 16, 0)  # Mar 16, 16:00 am (solar)
        lat = 43
        expected_value = deg2rad(70.3)
        self.assertAlmostEqual(theta_z(date, lat), expected_value, 2)

    def test_exception_date(self):
        self.assertRaises(TypeError, theta_z, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, theta_z, date, -91)
        self.assertRaises(TypeError, theta_z, date, '91')

class Test_solar_azimuth(ut.TestCase):
    """
    Tests solar azimuth angle function. Values from Duffie and Beckman
    """
    def test_noon(self):
        # different values of (date, lat) at noon
        hour, minute = 12, 0  # noon

        # northern hemisphere
        expected_value = deg2rad(0)

        date = datetime(2019, 1, 1, hour, minute)  # Jan 1
        lat = 0
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        date = datetime(2019, 4, 1, hour, minute)  # Apr 1
        lat = 30
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        date = datetime(2019, 8, 1, hour, minute)  # Aug 1
        lat = 60
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        date = datetime(2019, 10, 1, hour, minute)  # Oct 1
        lat = 90
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        # northern hemisphere
        expected_value = deg2rad(180)

        date = datetime(2019, 4, 1, hour, minute)  # Apr 1
        lat = -30
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        date = datetime(2019, 8, 1, hour, minute)  # Aug 1
        lat = -60
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        date = datetime(2019, 10, 1, hour, minute)  # Oct 1
        lat = -90
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

    def test_examples(self):
        # Example 1.6.2a
        date = datetime(2019, 2, 13, 9, 30)  # Feb 13, 9:30am
        lat = 43
        expected_value = deg2rad(-40.0)
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        # Example 1.6.2b
        date = datetime(2019, 7, 1, 18, 30)  # Jul 1, 1:30am
        lat = 43
        expected_value = deg2rad(112.0)
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

        # Example 1.6.3
        date = datetime(2019, 3, 16, 16, 0)  # Mar 16, 16:00am
        lat = 43
        expected_value = deg2rad(66.8)
        self.assertAlmostEqual(solar_azimuth(date, lat), expected_value, 2)

    def test_exception_date(self):
        self.assertRaises(TypeError, solar_azimuth, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, solar_azimuth, date, -91)
        self.assertRaises(TypeError, solar_azimuth, date, '91')

class Test_solar_altitude(ut.TestCase):
    """
    Tests solar azimuth angle function. Values from Duffie and Beckman
    """
    def test_examples(self):
        # Example 1.6.3
        date = datetime(2019, 3, 16, 16, 0)  # Mar 16, 16:00am
        lat = 43
        expected_value = deg2rad(19.7)
        self.assertAlmostEqual(solar_altitude(date, lat), expected_value, 2)

    def test_exception_date(self):
        self.assertRaises(TypeError, solar_altitude, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, solar_altitude, date, -91)
        self.assertRaises(TypeError, solar_altitude, date, '91')

class Test_sunset_hour_angle(ut.TestCase):
    """
    Tests sunset hour angle function. Values from Duffie and Beckman
    """
    def test_errors(self):
        date = datetime(2019, 1, 1)
        self.assertRaises(NoSunsetNoSunrise, sunset_hour_angle, date, 80)

        date = datetime(2019, 6, 20)
        self.assertRaises(NoSunsetNoSunrise, sunset_hour_angle, date, -75)

    def test_examples(self):
        # Example 1.6.3
        date = datetime(2019, 3, 16)
        lat = 43
        expected_value = deg2rad(87.8)
        self.assertAlmostEqual(sunset_hour_angle(date, lat), expected_value, 1)

    def test_exception_date(self):
        self.assertRaises(TypeError, sunset_hour_angle, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, sunset_hour_angle, date, 156)
        self.assertRaises(TypeError, sunset_hour_angle, date, '91')

class Test_sunset_time(ut.TestCase):
    """
    Tests sunset time function. Values from Duffie and Beckman
    """
    def test_examples(self):
        # Example 1.6.3
        date = datetime(2019, 3, 16, 16, 0)  # Mar 16, 16:00am
        lat = 43

        expected_value = 17
        self.assertAlmostEqual(sunset_time(date, lat).hour,
                               expected_value, 2)
        expected_value = 52  # diferent year than boook!
        self.assertAlmostEqual(sunset_time(date, lat).minute,
                               expected_value, 2)

    def test_NoSunsetNoSunrise(self):
        date = datetime(2019, 8, 1)  # summer
        lat = 89  # North-Pole
        self.assertRaises(NoSunsetNoSunrise, sunset_time, date, lat)

    def test_exception_date(self):
        self.assertRaises(TypeError, sunset_time, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, sunset_time, date, 156)
        self.assertRaises(TypeError, sunset_time, date, '91')

class Test_sunrise_hour_angle(ut.TestCase):
    """
    Tests sunrise hour angle function. Values from Duffie and Beckman
    """
    def test_errors(self):
        date = datetime(2019, 1, 1)
        self.assertRaises(NoSunsetNoSunrise, sunrise_hour_angle, date, 80)

        date = datetime(2019, 6, 20)
        self.assertRaises(NoSunsetNoSunrise, sunrise_hour_angle, date, -75)

    def test_examples(self):
        # Example 1.6.3
        date = datetime(2019, 3, 16)
        lat = 43
        expected_value = deg2rad(-87.8)
        self.assertAlmostEqual(sunrise_hour_angle(date, lat),
                               expected_value, 1)

    def test_exception_date(self):
        self.assertRaises(TypeError, sunrise_hour_angle, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, sunrise_hour_angle, date, 156)
        self.assertRaises(TypeError, sunrise_hour_angle, date, '91')

class Test_sunrise_time(ut.TestCase):
    """
    Tests sunrise time function. Values from Duffie and Beckman
    """
    def test_examples(self):
        # Example 1.6.3
        date = datetime(2019, 3, 16)
        lat = 43
        expected_value = 6
        self.assertAlmostEqual(sunrise_time(date, lat).hour,
                               expected_value, 2)
        expected_value = 7  # diferent year than boook!
        self.assertAlmostEqual(sunrise_time(date, lat).minute,
                               expected_value, 2)

    def test_NoSunsetNoSunrise(self):
        date = datetime(2019, 8, 1)  # summer
        lat = 89  # North-Pole
        self.assertRaises(NoSunsetNoSunrise, sunrise_time, date, lat)

    def test_exception_date(self):
        self.assertRaises(TypeError, sunrise_time, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, sunrise_time, date, 1526)
        self.assertRaises(TypeError, sunrise_time, date, '91')

class Test_daylight_hours(ut.TestCase):
    """
    Tests daylight hours function
    """
    def test_south_pole(self):
        # South Pole in the summer
        date = datetime(2019, 1, 1)
        lat = -80
        expected_value = 24
        self.assertAlmostEqual(daylight_hours(date, lat), expected_value)

        # South Pole in the winter
        date = datetime(2019, 8, 1)
        lat = -85
        expected_value = 0
        self.assertAlmostEqual(daylight_hours(date, lat), expected_value)

    def test_north_pole(self):
        # North Pole in the winter
        date = datetime(2019, 1, 1)
        lat = 82
        expected_value = 0
        self.assertAlmostEqual(daylight_hours(date, lat), expected_value)

        # North Pole in the summer
        date = datetime(2019, 8, 1)
        lat = 78
        expected_value = 24
        self.assertAlmostEqual(daylight_hours(date, lat), expected_value)

    def test_equator(self):
        # Equator in the august
        date = datetime(2019, 8, 15)
        lat = 0
        expected_value = 12
        self.assertAlmostEqual(daylight_hours(date, lat), expected_value)

        # Equator in the winter
        date = datetime(2019, 12, 15)
        lat = 0
        expected_value = 12
        self.assertAlmostEqual(daylight_hours(date, lat), expected_value)

    def test_exception_date(self):
        self.assertRaises(TypeError, daylight_hours, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, daylight_hours, date, 1526)
        self.assertRaises(TypeError, daylight_hours, date, '91')

class Test_solar_vector_ned(ut.TestCase):
    """
    Test function that calculates solar vector in ned frame
    """
    def test_permanent_darkness(self):
        # permanent darkness: south pole in winter
        date = datetime(2019, 6, 15, 12, 0)
        lat = -80
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

        date = datetime(2019, 6, 1, 17, 0)
        lat = -70
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

        # permanent darkness: north pole in winter
        date = datetime(2019, 1, 1, 10, 0)
        lat = 83
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

        date = datetime(2019, 11, 1, 19, 0)
        lat = 76
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

    def test_night(self):
        # night: north hemisphere
        date = datetime(2019, 1, 5)
        lat = 33
        ss_t = sunset_time(date, lat)
        hour, minute = ss_t.hour, ss_t.minute + 1  # 1min after sunset
        date = datetime(2019, 1, 5, hour, minute)
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

        date = datetime(2019, 9, 1)
        lat = 15
        sr_t = sunrise_time(date, lat)
        hour, minute = sr_t.hour, sr_t.minute - 1  # 1min before sunrise
        date = datetime(2019, 9, 1, hour, minute)
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

        # night: south hemisphere
        date = datetime(2019, 2, 3)
        lat = -63
        ss_t = sunset_time(date, lat)
        hour, minute = ss_t.hour, ss_t.minute + 1  # 1min after sunset
        date = datetime(2019, 2, 3, hour, minute)
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

        date = datetime(2019, 10, 1)
        n = 264
        lat = -15
        sr_t = sunrise_time(date, lat)
        hour, minute = sr_t.hour, sr_t.minute - 1  # 1min before sunrise
        date = datetime(2019, 10, 1, hour, minute)
        expected_value = array([0, 0, 0])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

    def test_summer_solstice(self):
        # summer solstice, solar noon, lat=declination
        date = datetime(2019, 6, 21, 12, 0)
        lat = 23 + 26/60 + 14/3600  # obliquity in 2019
        expected_value = array([0, 0, -1])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

    def test_permanent_day(self):
        # permanent light
        date = datetime(2019, 6, 20, 12, 0)  # summer solstice, solar noon
        lat = 90
        alt = deg2rad(23 + 26/60 + 14/3600)
        expected_value = array([-cos(alt), 0, -sin(alt)])
        assert_array_almost_equal(solar_vector_ned(date, lat),
                                  expected_value, 3)

    def test_exception_date(self):
        self.assertRaises(TypeError, solar_vector_ned, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, solar_vector_ned, date, 1526)
        self.assertRaises(TypeError, solar_vector_ned, date, '91')

class Test_air_mass_KY1989(ut.TestCase):
    """
    Tests air mass function based on the work of Kasten and Young (1989)
    """
    def test_errors(self):
        self.assertRaises(ValueError, air_mass_kastenyoung1989, 0, -1)

    def test_limit_values(self):
        # air mass through zenit direction at sea level
        theta_z = 0
        h = 0
        expected_value = 1
        self.assertAlmostEqual(air_mass_kastenyoung1989(theta_z, h),
                               expected_value, 3)

        # air mass through zenit direction at exosphere
        # requieres deactivation of limit parameter
        theta_z = 0
        h = 1e8  # 10.000 km
        expected_value = 0
        self.assertAlmostEqual(air_mass_kastenyoung1989(theta_z, h, False),
                               expected_value)

        # model limits (zenith=91.5)
        theta_z = 94
        h = 0
        expected_value = air_mass_kastenyoung1989(91.5, h)
        self.assertEqual(air_mass_kastenyoung1989(theta_z, h),
                         expected_value)

class Test_air_mass_Y1994(ut.TestCase):
    """
    Tests air mass function based on the work of Young (1994)
    """
    def test_limit_values(self):
        # air mass through zenit direction at sea level
        theta_z = 0
        expected_value = 1
        self.assertAlmostEqual(air_mass_young1994(theta_z), expected_value, 4)

class Test_beam_irradiance(ut.TestCase):
    """
    Tests beam_irradiance function based on the work of Aglietti et al (2009)
    """
    def test_errors(self):
        # erroneus altitud
        h = -10
        date = datetime(2019, 6, 20, 10, 0)
        lat = -63
        self.assertRaises(ValueError, beam_irradiance, h, date, lat)

        # erroneus latitude
        h = 0
        date = datetime(2019, 6, 20, 10, 0)
        lat = -91
        self.assertRaises(ValueError, beam_irradiance, h, date, lat)

    def test_limit_values(self):
        # sun below the horizon
        h = 0
        date = datetime(2019, 6, 20, 12, 0)
        lat = -69
        expected_value = 0
        self.assertEqual(beam_irradiance(h, date, lat), expected_value)

        h = 10000
        date = datetime(2019, 1, 1, 12, 0)
        lat = 87
        expected_value = 0
        self.assertEqual(beam_irradiance(h, date, lat), expected_value)

        # north pole winter nigth
        h = 1000
        date = datetime(2019, 11, 15, 5, 0)
        lat = 80
        expected_value = 0
        self.assertEqual(beam_irradiance(h, date, lat), expected_value)

        # south pole winter nigth
        h = 5000
        date = datetime(2019, 5, 20, 22, 0)
        lat = -85
        expected_value = 0
        self.assertEqual(beam_irradiance(h, date, lat), expected_value)

        # TODO: more test!

    def test_exception_alt(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, beam_irradiance, -1, date, 0)

    def test_exception_date(self):
        self.assertRaises(TypeError, beam_irradiance, 0, 121, 1)

    def test_exception_lat(self):
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, beam_irradiance, 0, date, 1526)
        self.assertRaises(TypeError, beam_irradiance, 0, date, '91')

class Test_irradiance_on_plane(ut.TestCase):
    """
    Test function that calculates solar irradiance in a plane defined
    by its normal vector
    """
    def tests_northern_hemisphere(self):
        # summer solstice, solar noon, lat=declination, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 6, 20, 12, 0)
        lat = 23 + 26/60 + 14/3600
        expected_value = beam_irradiance(h, date, lat)
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # summer solstice, solar noon, lat=declination, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 6, 20, 12, 0)
        lat = 23 + 26/60 + 14/3600
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # summer solstice, night, lat=declination, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 6, 20, 0, 0)
        lat = 23 + 26/60 + 14/3600
        expected_value = beam_irradiance(h, date, lat)
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # summer solstice, night, lat=declination, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 6, 20, 0, 0)
        lat = 23 + 26/60 + 14/3600
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, permanent darkness, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 12, 22, 12, 0)
        lat = 70
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, permanent darkness, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 12, 22, 12, 0)
        lat = 70
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, night, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 12, 22, 3, 0)
        lat = 40
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, night, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 12, 22, 3, 0)
        lat = 40
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # any day, solar noon, any latitude, plane sideways
        vnorm = array([0, 1, 0])
        h = 0
        date = datetime(2019, 4, 1, 12, 0)
        lat = 47.3
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

    def tests_southern_hemisphere(self):
        # summer solstice, solar noon, lat=declination, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 12, 22, 12, 0)
        lat = -(23 + 26/60 + 14/3600)
        expected_value = beam_irradiance(h, date, lat)
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # summer solstice, solar noon, lat=declination, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 12, 22, 12, 0)
        lat = -(23 + 26/60 + 14/3600)
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # summer solstice, night, lat=declination, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 12, 22, 0, 0)
        lat = -(23 + 26/60 + 14/3600)
        expected_value = beam_irradiance(h, date, lat)
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # summer solstice, night, lat=declination, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 12, 22, 0, 0)
        lat = -(23 + 26/60 + 14/3600)
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, permanent darkness, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 6, 20, 12, 0)
        lat = -70
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, permanent darkness, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 6, 20, 12, 0)
        lat = -70
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, night, plane right-side-up
        vnorm = array([0, 0, -1])
        h = 20000
        date = datetime(2019, 6, 20, 3, 0)
        lat = -40
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # winter solstice, night, plane upside-down
        vnorm = array([0, 0, 1])
        h = 20000
        date = datetime(2019, 6, 20, 3, 0)
        lat = -40
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

        # any day, solar noon, any latitude, plane sideways
        vnorm = array([0, 1, 0])
        h = 0
        date = datetime(2019, 10, 5, 12, 0)
        lat = -13.1
        expected_value = 0
        self.assertAlmostEqual(irradiance_on_plane(vnorm, h, date, lat),
                               expected_value, 3)

    def test_exception_vector(self):
        v = 'a'
        date = datetime(2019, 12, 13, 12, 0)
        self.assertRaises(ValueError, irradiance_on_plane, v, 0, date, 0)

    def test_exception_alt(self):
        v = array([0, 0, -1])
        date = datetime(2019, 12, 13, 12, 0)
        self.assertRaises(ValueError, irradiance_on_plane, v, -1, date, 0)

    def test_exception_date(self):
        v = array([0, 1, 0])
        self.assertRaises(TypeError, irradiance_on_plane, v, 0, 121, 1)

    def test_exception_lat(self):
        v = array([0, 1, 0])
        date = datetime(2019, 12, 13)
        self.assertRaises(ValueError, irradiance_on_plane, v, 0, date, 1526)
        self.assertRaises(TypeError, irradiance_on_plane, v, 0, date, '91')

max_stars_repo_name: aqreed/pysolarRadiation

max_issues_repo_name: aqreed/pysolarRadiation

max_forks_repo_name: aqreed/pysolarRadiation

reason: from numpy


text:
import h5py
import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
from torch.nn.functional import cosine_similarity as cos
from torch.nn.functional import pairwise_distance as euc
from torch.nn.functional import cross_entropy
from torch.nn.functional import margin_ranking_loss as rank_loss

def order_sim(hypers, hypos):
    sub = hypers - hypos
    act = nn.functional.relu(sub)
    partial_order_dis = act.norm(p=2, dim=1)
    partial_order_sim = -partial_order_dis
    return partial_order_sim

class RelationEmbedding(nn.Module):
    def __init__(self, input_len, output_len, label_vec_path):
        super(RelationEmbedding, self).__init__()
        self.hidden = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(input_len, output_len))
        self.output = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(output_len, output_len))

        label_vec_file = h5py.File(label_vec_path, 'r')
        gt_label_vecs = np.array(label_vec_file['label_vec'])
        self._gt_label_vecs = Variable(torch.from_numpy(gt_label_vecs)).float().cuda()

    def forward(self, sbj_vec, obj_vec):
        sbj_obj_vec = torch.cat([sbj_vec, obj_vec], 1)
        hidden = self.hidden(sbj_obj_vec)
        embedding = self.output(hidden)
        score_stack = Variable(torch.zeros(len(embedding), len(self._gt_label_vecs))).cuda()
        for i in range(len(embedding)):
            order_sims = order_sim(self._gt_label_vecs, embedding[i])
            score_stack[i] = order_sims
        return score_stack

def relation_embedding_loss(sbj_vec1, pre_vec1, obj_vec1, pre_emb1,
                            sbj_vec2, pre_vec2, obj_vec2, pre_emb2):
    emb_sim = order_sim(pre_emb1, pre_emb2)
    sbj_sim = order_sim(sbj_vec1, sbj_vec2)
    obj_sim = order_sim(obj_vec1, obj_vec2)
    pre_sim = order_sim(pre_vec1, pre_vec2)
    target = emb_sim / (sbj_sim + pre_sim + obj_sim)
    return torch.var(target)

def order_rank_loss(pos_sim, neg_sim):
    y = Variable(torch.ones(len(pos_sim))).cuda()
    # expect: positive sample get higher score, min margin = 1
    loss = rank_loss(pos_sim, neg_sim, y, margin=1)
    return loss

def order_softmax_test(batch_scores, pos_neg_inds):
    loss_scores = Variable(torch.zeros(len(batch_scores), len(pos_neg_inds[0]))).float().cuda()
    for i in range(len(batch_scores)):
        loss_scores[i] = batch_scores[i, pos_neg_inds[i]]
    y = Variable(torch.zeros(len(batch_scores))).long().cuda()
    acc = 0.0
    for scores in loss_scores:
        p_score = scores[0]
        n_score_max = torch.max(scores[1:])
        if p_score > n_score_max:
            acc += 1
    acc = acc / len(batch_scores)
    return acc, loss_scores, y

def order_rank_eval(pos_vecs, neg_vecs, gt_vecs):
    pos_sim = order_sim(gt_vecs, pos_vecs)
    neg_sim = order_sim(gt_vecs, neg_vecs)
    diff = pos_sim - neg_sim
    true_count = np.where(diff > 0)[0].shape[0] * 1.0
    acc = true_count / pos_sim.shape[0]
    return acc, pos_sim, neg_sim

def order_rank_test(pred_scores, gt_label_vecs):
    ranks = []
    for scores in pred_scores:
        ranked_inds = np.argsort(scores).tolist()
        ranked_inds.reverse()
        ranks.append(ranked_inds)
    return ranks

ext: py

reason: import numpy


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

# Parallel Processing Teaching Toolkit
# PyOpenCL - Example 07
# Multiple Kernels Execution
# https://github.com/javierip/parallel-processing-teaching-toolkit

import pyopencl as cl
import numpy as np

import time   # For measure the running times

VECTOR_SIZE = 50000  # Elements of vector

# Create four random vectors
a_host = np.random.rand(VECTOR_SIZE).astype(np.float32)
b_host = np.random.rand(VECTOR_SIZE).astype(np.float32)
c_host = np.random.rand(VECTOR_SIZE).astype(np.float32)
d_host = np.random.rand(VECTOR_SIZE).astype(np.float32)

# Create empty vectors for results
res_host_1= np.zeros(VECTOR_SIZE).astype(np.float32)
res_host_2= np.zeros(VECTOR_SIZE).astype(np.float32)
res_host= np.zeros(VECTOR_SIZE).astype(np.float32)
res_gpu_host= np.zeros(VECTOR_SIZE).astype(np.float32)

# Create CL context
platform = cl.get_platforms()[0]
device = platform.get_devices()[0] #get first gpu available

print "Running: ", platform
print "On GPU: ", device

ctx = cl.Context([device])
queue = cl.CommandQueue(ctx)

# Transfer host (CPU) memory to device (GPU) memory
mf = cl.mem_flags
a_gpu = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=a_host)
b_gpu = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=b_host)
c_gpu = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=c_host)
d_gpu = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=d_host)

# Kernel code
prg = cl.Program(ctx, """
__kernel void sum(__global const float *a_gpu, __global const float *b_gpu, __global float *res_gpu) {
  int gid = get_global_id(0);
  res_gpu[gid] = a_gpu[gid] + b_gpu[gid];
}
__kernel void multi(__global const float *a_gpu, __global const float *b_gpu, __global float *res_gpu) {
  int gid = get_global_id(0);
  res_gpu[gid] = a_gpu[gid] * b_gpu[gid];
}
__kernel void div(__global const float *a_gpu, __global const float *b_gpu, __global float *res_gpu) {
  int gid = get_global_id(0);
  res_gpu[gid] = a_gpu[gid] / b_gpu[gid];
}
""").build()

# Create empty gpu array for the result
res_gpu_1 = cl.Buffer(ctx, mf.WRITE_ONLY, a_host.nbytes)
res_gpu_2 = cl.Buffer(ctx, mf.WRITE_ONLY, a_host.nbytes)
res_gpu = cl.Buffer(ctx, mf.WRITE_ONLY, a_host.nbytes)

tic=time.time()
#Operation using the GPU - call the kernel on the card
prg.multi(queue, a_host.shape, None, a_gpu, b_gpu, res_gpu_1)
prg.div(queue, a_host.shape, None, c_gpu, d_gpu, res_gpu_2)
prg.sum(queue, a_host.shape, None, res_gpu_1, res_gpu_2, res_gpu)
time_gpu=time.time()-tic

#Clear GPU resources
res_gpu_host = np.empty_like(a_host)
cl.enqueue_copy(queue, res_gpu_host, res_gpu)

tic=time.time()

#Operation using the cpu
for i in range(0,VECTOR_SIZE):
	res_host_1[i]=a_host[i]*b_host[i]
for i in range(0,VECTOR_SIZE):
	res_host_2[i]=c_host[i]/d_host[i]
for i in range(0,VECTOR_SIZE):
	res_host[i]=res_host_1[i]+res_host_2[i]

time_cpu=time.time()-tic

# Print the results
print "-" * 80
print "CHECK :" #0 = GOOD
print "-" * 80
print res_gpu_host-res_host
print "-" * 80
print "Vector (a*b+c/d)"
print "Vector Size:", VECTOR_SIZE
print "Time CPU:", time_cpu
print "Time GPU:", time_gpu


text:
# code based on:
#YAD2K https://github.com/allanzelener/YAD2K
#darkflow https://github.com/thtrieu/darkflow
#Darknet.keras https://github.com/sunshineatnoon/Darknet.keras

import numpy as np
import cv2

def load_weights(model,yolo_weight_file):

    data = np.fromfile(yolo_weight_file,np.float32)
    data=data[4:]

    index = 0
    for layer in model.layers:
        shape = [w.shape for w in layer.get_weights()]
        if shape != []:
            kshape,bshape = shape
            bia = data[index:index+np.prod(bshape)].reshape(bshape)
            index += np.prod(bshape)
            ker = data[index:index+np.prod(kshape)].reshape(kshape)
            index += np.prod(kshape)
            layer.set_weights([ker,bia])

class Box:
    def __init__(self):
        self.x, self.y = float(), float()
        self.w, self.h = float(), float()
        self.c = float()
        self.prob = float()

def overlap(x1,w1,x2,w2):
    l1 = x1 - w1 / 2.;
    l2 = x2 - w2 / 2.;
    left = max(l1, l2)
    r1 = x1 + w1 / 2.;
    r2 = x2 + w2 / 2.;
    right = min(r1, r2)
    return right - left;

def box_intersection(a, b):
    w = overlap(a.x, a.w, b.x, b.w);
    h = overlap(a.y, a.h, b.y, b.h);
    if w < 0 or h < 0: return 0;
    area = w * h;
    return area;

def box_union(a, b):
    i = box_intersection(a, b);
    u = a.w * a.h + b.w * b.h - i;
    return u;

def box_iou(a, b):
    return box_intersection(a, b) / box_union(a, b);

def yolo_net_out_to_car_boxes(net_out, threshold = 0.2, sqrt=1.8,C=20, B=2, S=7):
    class_num = 6
    boxes = []
    SS        =  S * S # number of grid cells
    prob_size = SS * C # class probabilities
    conf_size = SS * B # confidences for each grid cell

    probs = net_out[0 : prob_size]
    confs = net_out[prob_size : (prob_size + conf_size)]
    cords = net_out[(prob_size + conf_size) : ]
    probs = probs.reshape([SS, C])
    confs = confs.reshape([SS, B])
    cords = cords.reshape([SS, B, 4])

    for grid in range(SS):
        for b in range(B):
            bx   = Box()
            bx.c =  confs[grid, b]
            bx.x = (cords[grid, b, 0] + grid %  S) / S
            bx.y = (cords[grid, b, 1] + grid // S) / S
            bx.w =  cords[grid, b, 2] ** sqrt
            bx.h =  cords[grid, b, 3] ** sqrt
            p = probs[grid, :] * bx.c

            if p[class_num] >= threshold:
                bx.prob = p[class_num]
                boxes.append(bx)

    # combine boxes that are overlap
    boxes.sort(key=lambda b:b.prob,reverse=True)
    for i in range(len(boxes)):
        boxi = boxes[i]
        if boxi.prob == 0: continue
        for j in range(i + 1, len(boxes)):
            boxj = boxes[j]
            if box_iou(boxi, boxj) >= .4:
                boxes[j].prob = 0.
    boxes = [b for b in boxes if b.prob > 0.]

    return boxes

def draw_box(boxes,im,crop_dim):
    imgcv = im
    [xmin,xmax] = crop_dim[0]
    [ymin,ymax] = crop_dim[1]
    for b in boxes:
        h, w, _ = imgcv.shape
        left  = int ((b.x - b.w/2.) * w)
        right = int ((b.x + b.w/2.) * w)
        top   = int ((b.y - b.h/2.) * h)
        bot   = int ((b.y + b.h/2.) * h)
        left = int(left*(xmax-xmin)/w + xmin)
        right = int(right*(xmax-xmin)/w + xmin)
        top = int(top*(ymax-ymin)/h + ymin)
        bot = int(bot*(ymax-ymin)/h + ymin)

        if left  < 0    :  left = 0
        if right > w - 1: right = w - 1
        if top   < 0    :   top = 0
        if bot   > h - 1:   bot = h - 1
        thick = int((h + w) // 150)
        cv2.rectangle(imgcv, (left, top), (right, bot), (255,0,0), thick)

    return imgcv

max_stars_repo_name: Wentaobi/Vehicle-Detection

max_issues_repo_name: Wentaobi/Vehicle-Detection

max_forks_repo_name: Wentaobi/Vehicle-Detection


text:
from OpenGL.GLUT import *
from OpenGL.GLU import *
from OpenGL.GL import *
from OpenGL.raw.GL.ARB.vertex_array_object import glGenVertexArrays, \
                          glBindVertexArray
import numpy as np

import sys

NumVertices = 6
VAOs = GLuint(0)

def ReadShader(filename):
    with open(filename, 'rb') as file:
        return file.read()
    return ''

def LoadShaders(shaders):
    program = glCreateProgram()

    for shaderEntry in shaders:
        shader = glCreateShader(shaderEntry[0])
        source = ReadShader(shaderEntry[1])

        glShaderSource(shader, source)
        glCompileShader(shader)

        log = glGetShaderInfoLog(shader)
        if log:
            print("Error in shader: {0}".format(shaderEntry[1]))
            print(log)

        glAttachShader(program, shader)

    glLinkProgram(program)

    return program

def init():
    global VAOs

    glGenVertexArrays(1, VAOs)
    glBindVertexArray(VAOs.value)

    vertices = [
        [-0.90, -0.90],
        [ 0.85, -0.90],
        [-0.90,  0.85],
        [ 0.90, -0.85],
        [ 0.90,  0.90],
        [-0.85,  0.90],
    ]

    vertices = np.array(vertices, dtype="f")

    buffers = glGenBuffers(1)
    glBindBuffer(GL_ARRAY_BUFFER, buffers)
    glBufferData(GL_ARRAY_BUFFER, vertices, GL_STATIC_DRAW)

    shaders = [
        (GL_VERTEX_SHADER, "triangles.vert"),
        (GL_FRAGMENT_SHADER, "triangles.frag"),
    ]

    program = LoadShaders(shaders)
    glUseProgram(program)
    glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, None)

    glEnableVertexAttribArray(0)

def display():
    glClear(GL_COLOR_BUFFER_BIT)

    glBindVertexArray(VAOs);
    glDrawArrays(GL_TRIANGLES, 0, NumVertices)

    glFlush()

if __name__ == '__main__':
    glutInit(sys.argv)
    glutInitDisplayMode(GLUT_RGBA)
    glutInitWindowSize(512, 512)
    glutInitContextVersion(3, 3)
    glutInitContextProfile(GLUT_CORE_PROFILE)
    glutCreateWindow(sys.argv[0])

    init()

    glutDisplayFunc(display)
    glutMainLoop()


text:
import base64
# from typing import Tuple, List, Union, Dict, Iterable

import numpy as np
# import torch.nn as nn

def rle_encode(mask: np.ndarray) -> dict:
    """Perform Run-Length Encoding (RLE) on a binary mask.
    """

    assert mask.dtype == bool and mask.ndim == 2, 'RLE encoding requires a binary mask (dtype=bool).'
    pixels = mask.flatten()
    pixels = np.concatenate([[0], pixels, [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    return dict(data=base64.b64encode(runs.astype(np.uint32).tobytes()).decode('utf-8'), shape=mask.shape)

def rle_decode(rle: dict) -> np.ndarray:
    """Decode a Run-Length Encoding (RLE).
    """

    runs = np.frombuffer(base64.b64decode(rle['data']), np.uint32)
    shape = rle['shape']
    starts, lengths = [np.asarray(x, dtype=int) for x in (runs[0:][::2], runs[1:][::2])]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(shape)

max_stars_repo_name: FinchZHU/ultra-thin-PRM

max_issues_repo_name: FinchZHU/ultra-thin-PRM

max_forks_repo_name: FinchZHU/ultra-thin-PRM


text:
import numpy as np
import argparse
import logging
import time

SUFFIX = '_shuffle.txt'

def load_input(filename):
    t = time.time()
    with open(filename) as f:
        data = f.readlines()
    logging.info('load %d lines in %.4f s', len(data), time.time() - t)
    t = time.time()
    np.random.shuffle(data)
    logging.info('shuffle in %.4f s', time.time() - t)
    return data

def get_output(filename):
    return filename + SUFFIX

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('input')
    parser.add_argument('-v', action='store_true')
    parser.add_argument('-d', action='store_true', help='dry run')
    args = parser.parse_args()
    if args.v:
        logging.basicConfig(level=logging.INFO)

    output = get_output(args.input)

    if not args.d:
        data = load_input(args.input)
        t = time.time()
        with open(output, 'w') as f:
            f.writelines(data)
        logging.info('write to %s in %.4f s', output, time.time() - t)

    print(output)

max_stars_repo_name: cgsdfc/Neural-Dialogue-Generation

max_issues_repo_name: cgsdfc/Neural-Dialogue-Generation

max_forks_repo_name: cgsdfc/Neural-Dialogue-Generation


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

"""A demo that runs hand tracking and object detection on camera frames using OpenCV. 2 EDGETPU
"""
import argparse
import collections
import common
import cv2
import numpy as np
import os
import math
from PIL import Image
import re
from edgetpu.detection.engine import DetectionEngine

import gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk

import time
import svgwrite
import gstreamer
from pose_engine import PoseEngine
import tflite_runtime.interpreter as tflite

Object = collections.namedtuple('Object', ['id', 'score', 'bbox'])

#==============================
EDGES = (
    ('nose', 'left eye'),
    ('nose', 'right eye'),
    ('nose', 'left ear'),
    ('nose', 'right ear'),
    ('left ear', 'left eye'),
    ('right ear', 'right eye'),
    ('left eye', 'right eye'),
    ('left shoulder', 'right shoulder'),
    ('left shoulder', 'left elbow'),
    ('left shoulder', 'left hip'),
    ('right shoulder', 'right elbow'),
    ('right shoulder', 'right hip'),
    ('left elbow', 'left wrist'),
    ('right elbow', 'right wrist'),
    ('left hip', 'right hip'),
    ('left hip', 'left knee'),
    ('right hip', 'right knee'),
    ('left knee', 'left ankle'),
    ('right knee', 'right ankle'),
)

HEADCHECK = ('nose', 'left eye','right eye' ,'left ear', 'right ear')
SHOULDERCHECK = ('left shoulder', 'right shoulder')
HIPCHECK = ('left hip','right hip')
KNEECHECK = ('left knee','right knee')
ANKLECHECK = ('left ankle','right ankle')

def shadow_text(cv2_im, x, y, text, font_size=16):
    cv2_im = cv2.putText(cv2_im, text, (x + 1, y + 1),
                             cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 0), 2)
    #dwg.add(dwg.text(text, insert=, fill='black',
    #                 font_size=font_size, style='font-family:sans-serif'))
    #dwg.add(dwg.text(text, insert=(x, y), fill='white',
    #                 font_size=font_size, style='font-family:sans-serif'))

def draw_pose(cv2_im, cv2_sodidi, pose, numobject, src_size, color='yellow', threshold=0.2):
    box_x = 0
    box_y = 0
    box_w = 641
    box_h = 480
    scale_x, scale_y = src_size[0] / box_w, src_size[1] / box_h
    xys = {}
    #==bien dung de tinh khoang cach giua cac bo phan trong co the ============
    pts_sodidi = []
    headarea={}
    shoulderarea={}
    elbow={}
    lengbackbone=60
    lengleg= 86
    lengface = 30
    #=======================================================
    for label, keypoint in pose.keypoints.items():
        if keypoint.score < threshold: continue
        # Offset and scale to source coordinate space.
        kp_y = int((keypoint.yx[0] - box_y) * scale_y)
        kp_x = int((keypoint.yx[1] - box_x) * scale_x)
        cv2_im = cv2.putText(cv2_im, str(numobject),(kp_x + 1, kp_y + 1), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 0), 2)
        xys[label] = (numobject,kp_x, kp_y)

        cv2.circle(cv2_im,(int(kp_x),int(kp_y)),5,(0,255,255),-1)

    return xys

    '''
    for a, b in EDGES:
        if a not in xys or b not in xys: continue
        anum,ax, ay = xys[a]
        bnum,bx, by = xys[b]
        print(numobject,a,xys[a],b,xys[b])
        cv2.line(cv2_im,(ax, ay), (bx, by),(0,255,255))
    '''

def mapcamto2dplane(pts_in):
    # provide points from image 1
    pts_src = np.array([[154, 174], [702, 349], [702, 572],[1, 572], [1, 191]])
    # corresponding points from image 2 (i.e. (154, 174) matches (212, 80))
    pts_dst = np.array([[154, 174], [702, 349], [702, 572],[1, 572], [1, 191]])#np.array([[212, 80],[489, 80],[505, 180],[367, 235], [144,153]])

    # calculate matrix H
    h, status = cv2.findHomography(pts_src, pts_dst)

    # provide a point you wish to map from image 1 to image 2
    #pts_in = np.array([[154, 174]], dtype='float32')
    #pts_in = np.array([pts_in])

    # finally, get the mapping
    pointsOut = cv2.perspectiveTransform(pts_in, h)
    pointsOut = np.array([pointsOut])
    point_out = [b for a in pointsOut for b in a]
    return point_out

def check_distance(x1,y1,x2,y2):
    dist = math.sqrt((x2-x1)**2 + (y2-y1)**2)
    return dist

def avg_fps_counter(window_size):
    window = collections.deque(maxlen=window_size)
    prev = time.monotonic()
    yield 0.0  # First fps value.

    while True:
        curr = time.monotonic()
        window.append(curr - prev)
        prev = curr
        yield len(window) / sum(window)

#==============================
def load_labels(path):
    p = re.compile(r'\s*(\d+)(.+)')
    with open(path, 'r', encoding='utf-8') as f:
       lines = (p.match(line).groups() for line in f.readlines())
       return {int(num): text.strip() for num, text in lines}

class BBox(collections.namedtuple('BBox', ['xmin', 'ymin', 'xmax', 'ymax'])):
    """Bounding box.
    Represents a rectangle which sides are either vertical or horizontal, parallel
    to the x or y axis.
    """
    __slots__ = ()

def get_output(interpreter, score_threshold, top_k, image_scale=1.0):
    """Returns list of detected objects."""
    boxes = common.output_tensor(interpreter, 0)
    class_ids = common.output_tensor(interpreter, 1)
    scores = common.output_tensor(interpreter, 2)
    count = int(common.output_tensor(interpreter, 3))

    def make(i):
        ymin, xmin, ymax, xmax = boxes[i]
        return Object(
            id=int(class_ids[i]),
            score=scores[i],
            bbox=BBox(xmin=np.maximum(0.0, xmin),
                      ymin=np.maximum(0.0, ymin),
                      xmax=np.minimum(1.0, xmax),
                      ymax=np.minimum(1.0, ymax)))

    return [make(i) for i in range(top_k) if scores[i] >= score_threshold]

def main():
    default_model_dir = '../all_models'
    default_model = 'posenet/posenet_mobilenet_v1_075_481_641_quant_decoder_edgetpu.tflite'
    default_labels = 'hand_label.txt'
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', help='.tflite model path',
                        default=os.path.join(default_model_dir,default_model))
    parser.add_argument('--labels', help='label file path',
                        default=os.path.join(default_model_dir, default_labels))
    parser.add_argument('--top_k', type=int, default=1,
                        help='number of categories with highest score to display')
    parser.add_argument('--camera_idx', type=str, help='Index of which video source to use. ', default = 0)
    parser.add_argument('--threshold', type=float, default=0.5,
                        help='classifier score threshold')
    args = parser.parse_args()

    #print('Loading Handtracking model {} with {} labels.'.format(args.model, args.labels))

    #engine = DetectionEngine(args.model)
    #labels = load_labels(args.labels)
    #=====================================================================
    src_size = (640, 480)
    print('Loading Pose model {}'.format(args.model))
    engine = PoseEngine(args.model)
    #engine = PoseEngine('models/mobilenet/posenet_mobilenet_v1_075_481_641_quant_decoder_edgetpu.tflite')
    #=====================================================================
    # for detection
    print('Loading Detection model {} with {} labels.'.format('../all_models/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite', '../all_models/coco_labels.txt'))
    #interpreter2 = common.make_interpreter('../all_models/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite')
    #interpreter2.allocate_tensors()
    engine2 = DetectionEngine('../all_models/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite')
    labels2 = load_labels('../all_models/coco_labels.txt')
    cap = cv2.VideoCapture(args.camera_idx)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        cv2_im = frame

        cv2_im_rgb = cv2.cvtColor(cv2_im, cv2.COLOR_BGR2RGB)
        pil_im = Image.fromarray(cv2_im_rgb)
        #declare new window for show pose in 2d plane========================
        h_cap, w_cap, _ = cv2_im.shape
        cv2_sodidi = np.zeros((h_cap,w_cap,3), np.uint8)
        #======================================pose processing=================================

        poses, inference_time = engine.DetectPosesInImage(np.uint8(pil_im.resize((641, 481), Image.NEAREST)))
        #print('Posese is',poses)

        n = 0
        sum_process_time = 0
        sum_inference_time = 0
        ctr = 0
        fps_counter  = avg_fps_counter(30)

        input_shape = engine.get_input_tensor_shape()

        inference_size = (input_shape[2], input_shape[1])

        #print('Shape is',input_shape)
        #print('inference size is:',inference_size)
        start_time = time.monotonic()

        end_time = time.monotonic()
        n += 1
        sum_process_time += 1000 * (end_time - start_time)
        sum_inference_time += inference_time

        avg_inference_time = sum_inference_time / n
        text_line = 'PoseNet: %.1fms (%.2f fps) TrueFPS: %.2f' % (
            avg_inference_time, 1000 / avg_inference_time, next(fps_counter)
        )

        shadow_text(cv2_im, 10, 20, text_line)
        numobject = 0
        xys={}
        pts_sodidi_arr=[]
        pts_xys_arr=[]
        listwarning=[]
        #draw_pose(cv2_im, poses, dis, src_size)
        for pose in poses:
           for i in range(len(poses)-1):
            pose = poses[i]
            for label, keypoint in pose.keypoints.items():
                #print(label)
                #print(keypoint)
                if keypoint.score < 0.2: continue
                if (label=='nose'):
                    print('yx0,',keypoint.yx)

            for j in range(len(poses)):
                pose1 = poses[j]
                #print(pose.keypoints.items())
                for label, keypoint in pose1.keypoints.items():
                    if keypoint.score < 0.2: continue
                    if (label=='nose'):
                        print('yx1,',keypoint.yx)
            '''
            pts_sodidi, xys = draw_pose(cv2_im,cv2_sodidi, pose, numobject, src_size)
            #print(pts_sodidi)
            pts_sodidi_arr.append(pts_sodidi)
            pts_xys_arr.append(xys)

            numobject += 1
            #print('len coor_av',coor_ave)
            #print(xys,coor_ave)kghkkgkgkgerg.hbjbbsbdbs
        pts_sodidi_arr = np.array([pts_sodidi_arr])
        v2 = [b for a in pts_sodidi_arr for b in a]
        print(v2)
        print(xys)
        print(numobject)

        #leng = coor_ave.length
        #print(leng)

        #for a in pts_sodidi_arr:
        #    for b in a:
        #        print(b[0])

        for i in range(0,len(v2)):
            a,x1,y1 = v2[i]
            for j in range(1,len(v2)):
                if i == j:
                    break
                b,x2,y2 = v2[j]
                distance = check_distance(x1,y1,x2,y2)
                print('distance',distance)
                if distance > 100:
                    cv2_sodidi = cv2.circle(cv2_sodidi,(int(x1),int(y1)),5,(0,0,255),-1)
                    cv2_sodidi = cv2.circle(cv2_sodidi,(int(x2),int(y2)),5,(0,0,255),-1)
                    listwarning.append(i)
                    listwarning.append(j)
                else:
                    cv2_sodidi = cv2.circle(cv2_sodidi,(int(x1),int(y1)),5,(255,0,0),-1)
                    cv2_sodidi = cv2.circle(cv2_sodidi,(int(x2),int(y2)),5,(255,0,0),-1)
        print('listwarning',listwarning)
        for a, b in EDGES:
            if a not in xys or b not in xys: continue
            num,ax, ay = xys[a]
            num,bx, by = xys[b]
            if num in listwarning:
            #print(numobject,a,xys[a],b,xys[b])
                cv2.line(cv2_im,(ax, ay), (bx, by),(0,0,255))
            else:
                cv2.line(cv2_im,(ax, ay), (bx, by),(255,0,0))
        #==============================================================================================
        #cv2_im = append_objs_to_img(cv2_im, objs, labels)

        # detection
        #common.set_input(interpreter2, pil_im)
        #interpreter2.invoke()
        #objs = get_output(interpreter2, score_threshold=0.2, top_k=3)
        objs = engine2.detect_with_image(pil_im,
                                  threshold=0.2,
                                  keep_aspect_ratio=True,
                                  relative_coord=True,
                                  top_k=3)

        cv2_im = append_objs_to_img(cv2_im, objs, labels2)

        cv2.imshow('frame', cv2_im)
        cv2.imshow('1', cv2_sodidi)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

def append_objs_to_img(cv2_im, objs, labels):

    height, width, channels = cv2_im.shape
    for obj in objs:
        x0, y0, x1, y1 = obj.bounding_box.flatten().tolist() #list(obj.bbox)
        x0, y0, x1, y1 = int(x0*width), int(y0*height), int(x1*width), int(y1*height)
        percent = int(100 * obj.score)
        label = '{}% {}'.format(percent, labels.get(obj.label_id, obj.label_id))
        if(labels.get(obj.label_id, obj.label_id)=='person'):
            cv2_im = cv2.rectangle(cv2_im, (x0, y0), (x1, y1), (0, 255, 0), 2)
            cv2_im = cv2.putText(cv2_im, label, (x0, y0+30),
                                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 0), 2)
    return cv2_im

    '''
    height, width, channels = cv2_im.shape

    boxes_ob = []
    confidences = []
    classIDs = []
    for obj in objs:
        x0, y0, x1, y1 = obj.bounding_box.flatten().tolist() #list(obj.bbox)
        x0, y0, x1, y1 = int(x0*width), int(y0*height), int(x1*width), int(y1*height)
        percent = int(100 * obj.score)
        confidence = obj.score
        label = '{}% {}'.format(percent, labels.get(obj.label_id, obj.label_id))
        if(labels.get(obj.label_id, obj.label_id)=='person'):
            classIDs.append(0)
            confidences.append(float(confidence))
            boxes_ob.append([x0, y0, x1,y1])
            #cv2_im = cv2.rectangle(cv2_im, (x0, y0), (x1, y1), (0, 255, 0), 2)
            #cv2_im = cv2.putText(cv2_im, label, (x0, y0+30),
            #                 cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 0), 2)
    idxs = cv2.dnn.NMSBoxes(boxes_ob, confidences, 0.5,0.3)
    #print('idxs',idxs)
    #print('classID',classIDs)
    ind = []
    for i in range(0,len(classIDs)):
        if(classIDs[i]==0):
            ind.append(i)
    a = []
    b = []

    if len(idxs) > 0:
            for i in idxs.flatten():
                (x, y) = (boxes_ob[i][0], boxes_ob[i][1])
                (w, h) = (boxes_ob[i][2], boxes_ob[i][3])
                a.append(x)
                b.append(y)
    distance=[]
    nsd = []
    for i in range(0,len(a)-1):
        for k in range(1,len(a)):
            if(k==i):
                break
            else:
                x_dist = (a[k] - a[i])
                y_dist = (b[k] - b[i])
                d = math.sqrt(x_dist * x_dist + y_dist * y_dist)
                distance.append(d)
                if(d <=1000):
                    nsd.append(i)
                    nsd.append(k)
                nsd = list(dict.fromkeys(nsd))
                print(nsd)
    color = (0, 0, 255)
    for i in nsd:
        (x, y) = (boxes_ob[i][0], boxes_ob[i][1])
        (w, h) = (boxes_ob[i][2], boxes_ob[i][3])
        cv2_im=cv2.rectangle(cv2_im, (x, y), (x + w, y + h), color, 2)
        text = "Alert"
        cv2_im=cv2.putText(cv2_im, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,0.5, color, 2)
    color = (0, 255, 0)
    if len(idxs) > 0:
        for i in idxs.flatten():
            if (i in nsd):
                break
            else:
                (x, y) = (boxes_ob[i][0], boxes_ob[i][1])
                (w, h) = (boxes_ob[i][2], boxes_ob[i][3])
                cv2_im=cv2.rectangle(cv2_im, (x, y), (x + w, y + h), color, 2)
                text = 'OK'
                cv2_im=cv2.putText(cv2_im, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,0.5, color, 2)

    #cv2.imshow("Social Distancing Detector", image)
    return cv2_im
    '''

if __name__ == '__main__':
    main()

max_stars_repo_name: haithienld/tommy-coral

max_issues_repo_name: haithienld/tommy-coral

max_forks_repo_name: haithienld/tommy-coral


text:
import pandas as pd
import json
import re
import numpy as np
import glob

class EmptySolutionError(Exception):
    pass

def get_max_xrate(o):
    if o["isSellOrder"]:
        return o["maxSellAmount"] / o["minBuyAmount"]
    else:
        return o["maxBuyAmount"] / o["maxSellAmount"]

def compute_avg_eth_price_usd(orders):
    return sum(
        [o["sellTokenDailyPriceUSD"] for o in orders if o["sellToken"] == "WETH"]
        + [o["buyTokenDailyPriceUSD"] for o in orders if o["buyToken"] == "WETH"]
    ) / len([o for o in orders if "WETH" in {o["sellToken"], o["buyToken"]}])

# O(n) iterator that for every element A in the first iterable returns
# the largest element B on the second iterable that satisfies test.
# Assumes that both lists are sorted.
# Example: for this input
# 1 2 4 5 9
# 1 3 6 10
# test(a,b) = a >= b
# returns:
# 1 1 3 3 6
def get_largest_element_sequence(a, b, test):
    idx_b = 0
    for idx_a in range(len(a)):
        while idx_b < len(b) - 1 and test(a[idx_a], b[idx_b + 1]):
            idx_b += 1
        if not test(a[idx_a], b[idx_b]):
            raise ValueError("Found no element satisfying test.")
        yield b[idx_b]

def load_block_data_file_to_df(fname):
    with open(fname, "r") as f:
        d = json.load(f)
    eth_price_usd = compute_avg_eth_price_usd(d["orders"])
    d = [
        {
            "block": o["uniswap"]["block"],
            "index": o["uniswap"]["index"],
            "sell_token": o["sellToken"],
            "buy_token": o["buyToken"],
            "max_buy_amount": o["maxBuyAmount"] if not o["isSellOrder"] else None,
            "max_sell_amount": o["maxSellAmount"] if o["isSellOrder"] else None,
            "sell_token_price_eth": o["sellTokenPriceETH"],
            "buy_token_price_eth": o["buyTokenPriceETH"],
            "sell_token_price_usd": o["sellTokenPriceETH"] * eth_price_usd,
            "buy_token_price_usd": o["buyTokenPriceETH"] * eth_price_usd,
            "timestamp": o["uniswap"]["timestamp"],
            "exec_sell_amount": o["uniswap"]["amounts"][0],
            "exec_buy_amount": o["uniswap"]["amounts"][-1],
            "nr_pools": len(o["uniswap"]["amounts"]) - 1,
            "is_sell_order": o["isSellOrder"],
            "address": o["address"],
            "sell_reserve": float(o["uniswap"]["balancesSellToken"][0]),
            "buy_reserve": float(o["uniswap"]["balancesBuyToken"][-1]),
            #'max_xrate': get_max_xrate(o)
        }
        for o in d["orders"]
    ]
    df = pd.DataFrame.from_records(d)
    df["xrate"] = df.exec_sell_amount / df.exec_buy_amount
    df["block_index"] = df.apply(
        lambda r: "_".join(r[["block", "index"]].astype(str).values), axis=1
    )
    df["token_pair"] = df.apply(
        lambda r: "-".join(sorted([r["sell_token"], r["buy_token"]])), axis=1
    )
    df["exec_vol"] = df.exec_sell_amount * df.sell_token_price_usd
    df["max_vol_usd"] = df.apply(
        lambda r: r.max_sell_amount * r.sell_token_price_usd
        if r.is_sell_order
        else r.max_buy_amount * r.buy_token_price_usd,
        axis=1,
    )
    df["max_vol_eth"] = df.apply(
        lambda r: r.max_sell_amount * r.sell_token_price_eth
        if r.is_sell_order
        else r.max_buy_amount * r.buy_token_price_eth,
        axis=1,
    )

    return df.set_index("block_index")

def remove_most_active_users(df_exec, fraction_to_remove):
    nr_addresses = df_exec.address.nunique()
    addresses = (
        df_exec.address.value_counts()
        .iloc[round(nr_addresses * fraction_to_remove) :]
        .index
    )
    return df_exec[df_exec.address.isin(addresses)]

def load_solver_solution(fname):
    with open(fname, "r") as f:
        d = json.load(f)
    d = [
        {
            "block": int(oid.split("-")[0]),
            "index": int(oid.split("-")[1]),
            "sell_token": o["sell_token"],
            "buy_token": o["buy_token"],
            "exec_sell_amount": int(o["exec_sell_amount"]) * 1e-18,
            "exec_buy_amount": int(o["exec_buy_amount"]) * 1e-18,
            "is_sell_order": o["is_sell_order"],
        }
        for oid, o in d["orders"].items()
    ]
    if len(d) == 0:
        raise EmptySolutionError()
    df = pd.DataFrame.from_records(d)
    df["xrate"] = df.exec_sell_amount / df.exec_buy_amount
    df["block_index"] = df.apply(
        lambda r: "_".join(r[["block", "index"]].astype(str).values), axis=1
    )
    return df.set_index("block_index")

def merge_exec_and_solved(fname, df_exec, from_timestamp, to_timestamp):
    df_sol = load_solver_solution(fname)
    df = df_exec[
        (df_exec.timestamp >= from_timestamp) & (df_exec.timestamp <= to_timestamp)
    ].merge(
        df_sol[["exec_sell_amount", "exec_buy_amount", "xrate"]],
        how="inner",
        on="block_index",
        suffixes=("_uni", "_gp"),
    )
    df["batch_start_time"] = from_timestamp
    df["batch_end_time"] = to_timestamp
    df["surplus"] = df.xrate_uni / df.xrate_gp
    savings_buy = df.exec_buy_amount_gp - df.exec_buy_amount_uni
    savings_sell = df.exec_sell_amount_uni - df.exec_sell_amount_gp
    df["savings_vol_usd"] = (
        savings_buy * df["buy_token_price_usd"]
        + savings_sell * df["sell_token_price_usd"]
    )
    return df

def create_batch_table(solution_fname, df_exec):
    m = re.search(r"_([0-9]+)\-([0-9]+)(\-[0-9]+)*\.json$", solution_fname)
    from_timestamp, to_timestamp = int(m[1]), int(m[2])
    return merge_exec_and_solved(solution_fname, df_exec, from_timestamp, to_timestamp)

def compute_savings_per_token(df):
    savings_buy_per_token = df.groupby("buy_token").savings_buy.sum()
    savings_sell_per_token = df.groupby("sell_token").savings_sell.sum()
    return savings_buy_per_token.add(savings_sell_per_token, fill_value=0)

def compute_mean_gp_rel_surplus(df):
    return np.exp(np.mean(np.log(df.xrate_uni) - np.log(df.xrate_gp)))

def create_batches_table(solution_dir, df_exec):
    dfs = []
    for fname in glob.glob(f"{solution_dir}/*.json"):
        try:
            dfs.append(create_batch_table(fname, df_exec))
        except EmptySolutionError:
            pass
    return pd.concat(dfs, axis=0).sort_index()

def compute_orig_batch(batchdf, df_exec):
    batch_start_time = batchdf.batch_start_time.iloc[0]
    batch_end_time = batchdf.batch_end_time.iloc[0]
    return df_exec[
        (df_exec.timestamp >= batch_start_time) & (df_exec.timestamp <= batch_end_time)
    ]

def compute_orig_batch_size(batchdf, df_exec):
    batch_start_time = batchdf.batch_start_time.iloc[0]
    batch_end_time = batchdf.batch_end_time.iloc[0]
    return (
        (df_exec.timestamp >= batch_start_time) & (df_exec.timestamp <= batch_end_time)
    ).sum()

def remove_batches_not_fully_executed(df_sol, df_exec):
    problem_batch_sizes = df_sol.groupby(["batch_start_time", "batch_end_time"]).apply(
        compute_orig_batch_size, df_exec=df_exec
    )
    solution_batch_sizes = (
        df_sol.groupby(["batch_start_time", "batch_end_time"]).count().block
    )
    batch_start_times = [
        b[0]
        for b in solution_batch_sizes[solution_batch_sizes == problem_batch_sizes].index
    ]
    return df_sol[df_sol.batch_start_time.isin(batch_start_times)]

def compute_orig_total_orders(df_sol, df_exec):
    df = df_sol.groupby(["batch_start_time", "batch_end_time"]).apply(
        compute_orig_batch, df_exec=df_exec
    )
    tokens = pd.concat([df_sol.sell_token, df_sol.buy_token], axis=0).unique()
    return (df.sell_token.isin(tokens) & df.buy_token.isin(tokens)).sum()

def compute_orig_total_users(df_sol, df_exec):
    df = df_sol.groupby(["batch_start_time", "batch_end_time"]).apply(
        compute_orig_batch, df_exec=df_exec
    )
    tokens = pd.concat([df_sol.sell_token, df_sol.buy_token], axis=0).unique()
    return df[df.sell_token.isin(tokens) & df.buy_token.isin(tokens)].address.nunique()

def filter_batches_with_large_liquidity_updates(df_sol):
    # remove batches for which there was a liquidity update to some used pool
    # that resulted in a change of more or less CUTOFF fraction of its liquidity
    CUTOFF = 0.3

    def large_liquidity_update_occurred_in_batch(batch_df):
        def occurred_in_token_pair(batch_df):
            if batch_df.shape[0] == 1:
                return False

            def occurred_between_consecutive_trades(r):
                n1 = r.sell_reserve.iloc[0] + r.exec_sell_amount_uni.iloc[0]
                if r.sell_token.iloc[0] == r.sell_token.iloc[1]:
                    n2 = r.sell_reserve.iloc[1]
                else:
                    assert r.sell_token.iloc[0] == r.buy_token.iloc[1]
                    n2 = r.buy_reserve.iloc[1]
                return abs(n1 - n2) / max(n1, n2) >= CUTOFF

            df = pd.concat([batch_df, batch_df.shift(-1)], axis=1).iloc[:-1]
            return np.any(df.apply(occurred_between_consecutive_trades, axis=1))

        return np.any(batch_df.groupby("token_pair").apply(occurred_in_token_pair))

    df = df_sol[
        [
            "batch_start_time",
            "token_pair",
            "sell_token",
            "buy_token",
            "sell_reserve",
            "buy_reserve",
            "exec_sell_amount_uni",
            "exec_buy_amount_uni",
        ]
    ].groupby("batch_start_time")
    m = df.apply(large_liquidity_update_occurred_in_batch)
    bad_batches = m[m].index
    return df_sol[~df_sol.batch_start_time.isin(bad_batches)]

def get_dfs(
    instance_path, batch_duration, nr_tokens, user_frac, limit_xrate_relax_frac
):
    data_path = f"{instance_path}/s{batch_duration}-t{nr_tokens}-u{user_frac}-l{limit_xrate_relax_frac}/"
    df_exec = load_block_data_file_to_df(f"{data_path}/per_block.json")
    df_sol = create_batches_table(f"{data_path}/solutions/", df_exec)

    # remove batches where there were untouched orders
    df_sol = remove_batches_not_fully_executed(df_sol, df_exec)

    # remove outliers (bottom and top OUTLIER_FRAC quantile of surplus variable)
    # OUTLIER_FRAC = 0.01
    # not_outlier = (df_sol.surplus > df_sol.surplus.quantile(OUTLIER_FRAC)) & (df_sol.surplus < df_sol.surplus.quantile(1-OUTLIER_FRAC))
    # df_sol = df_sol[not_outlier]

    # v = df_sol.max_vol_usd.quantile(.99)
    # df_sol = df_sol[df_sol.max_vol_usd <= v]

    df_sol = filter_batches_with_large_liquidity_updates(df_sol)

    # remove batches with weird results
    # df_sol = df_sol[~df_sol.batch_start_time.isin([1603206524])]

    return (df_sol, df_exec)

def get_block_data_file(
    instance_path, batch_duration, nr_tokens, user_frac, limit_xrate_relax_frac
):
    data_path = f"{instance_path}/s{batch_duration}-t{nr_tokens}-u{user_frac}-l{limit_xrate_relax_frac}/"
    return load_block_data_file_to_df(f"{data_path}/per_block.json")

def get_prices_at_blocks(data_path, blocks, tokens):
    with open(f"{data_path}/per_block.json", "r") as f:
        d = json.load(f)
    prices_in_file = {int(k): v for k, v in d["spot_prices"].items()}
    blocks_in_file = list(prices_in_file.keys())

    prices = {b: {t: None} for b in blocks for t in tokens}
    for t in tokens:
        blocks_with_prices_for_t = list(
            get_largest_element_sequence(
                blocks,
                blocks_in_file,
                lambda a, b: b <= a and t in prices_in_file[b].keys(),
            )
        )
        for bi in range(len(blocks)):
            prices[blocks[bi]][t] = prices_in_file[blocks_with_prices_for_t[bi]][t]
    # prices = {blocks[bi]: {t: prices[blocks_in_file[bi]][t]} for bi in range(len(blocks)) for t in prices[blocks_in_file[bi]].keys()}
    assert set(prices.keys()) == set(blocks)
    return prices

max_stars_repo_name: The-Makers-of-things/dune-snippets

max_issues_repo_name: The-Makers-of-things/dune-snippets

max_forks_repo_name: The-Makers-of-things/dune-snippets


text:
import torch
import torch.nn as nn
import numpy as np

class Mish(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x * (torch.tanh(torch.nn.functional.softplus(x)))
        return x

class Upsample(nn.Module):
    def __init__(self):
        super(Upsample, self).__init__()

    def forward(self, x, target_size, inference=False):
        assert (x.data.dim() == 4)
        # _, _, tH, tW = target_size

        if inference:

            # B = x.data.size(0)
            # C = x.data.size(1)
            # H = x.data.size(2)
            # W = x.data.size(3)

            return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1). \
                expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3),
                       target_size[3] // x.size(3)). \
                contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
        else:
            return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')

class Conv_Bn_Activation(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, activation, bn=True, bias=False):
        super().__init__()
        pad = (kernel_size - 1) // 2

        self.conv = nn.ModuleList()
        if bias:
            self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad))
        else:
            self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad, bias=False))
        if bn:
            self.conv.append(nn.BatchNorm2d(out_channels))
        if activation == "mish":
            self.conv.append(Mish())
        elif activation == "relu":
            self.conv.append(nn.ReLU(inplace=True))
        elif activation == "leaky":
            self.conv.append(nn.LeakyReLU(0.1, inplace=True))
        elif activation == "linear":
            pass
        else:
            print("activate error !!! {} {} {}".format(sys._getframe().f_code.co_filename,
                                                       sys._getframe().f_code.co_name, sys._getframe().f_lineno))

    def forward(self, x):
        for l in self.conv:
            x = l(x)
        return x

class ResBlock(nn.Module):
    """
    Sequential residual blocks each of which consists of \
    two convolution layers.
    Args:
        ch (int): number of input and output channels.
        nblocks (int): number of residual blocks.
        shortcut (bool): if True, residual tensor addition is enabled.
    """

    def __init__(self, ch, nblocks=1, shortcut=True):
        super().__init__()
        self.shortcut = shortcut
        self.module_list = nn.ModuleList()
        for i in range(nblocks):
            resblock_one = nn.ModuleList()
            resblock_one.append(Conv_Bn_Activation(ch, ch, 1, 1, 'mish'))
            resblock_one.append(Conv_Bn_Activation(ch, ch, 3, 1, 'mish'))
            self.module_list.append(resblock_one)

    def forward(self, x):
        for module in self.module_list:
            h = x
            for res in module:
                h = res(h)
            x = x + h if self.shortcut else h
        return x

class DownSample1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv_Bn_Activation(3, 32, 3, 1, 'mish')

        self.conv2 = Conv_Bn_Activation(32, 64, 3, 2, 'mish')
        self.conv3 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
        # [route]
        # layers = -2
        self.conv4 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')

        self.conv5 = Conv_Bn_Activation(64, 32, 1, 1, 'mish')
        self.conv6 = Conv_Bn_Activation(32, 64, 3, 1, 'mish')
        # [shortcut]
        # from=-3
        # activation = linear

        self.conv7 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
        # [route]
        # layers = -1, -7
        self.conv8 = Conv_Bn_Activation(128, 64, 1, 1, 'mish')

    def forward(self, input):
        x1 = self.conv1(input)
        x2 = self.conv2(x1)
        x3 = self.conv3(x2)
        # route -2
        x4 = self.conv4(x2)
        x5 = self.conv5(x4)
        x6 = self.conv6(x5)
        # shortcut -3
        x6 = x6 + x4

        x7 = self.conv7(x6)
        # [route]
        # layers = -1, -7
        x7 = torch.cat([x7, x3], dim=1)
        x8 = self.conv8(x7)
        return x8

class DownSample2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv_Bn_Activation(64, 128, 3, 2, 'mish')
        self.conv2 = Conv_Bn_Activation(128, 64, 1, 1, 'mish')
        # r -2
        self.conv3 = Conv_Bn_Activation(128, 64, 1, 1, 'mish')

        self.resblock = ResBlock(ch=64, nblocks=2)

        # s -3
        self.conv4 = Conv_Bn_Activation(64, 64, 1, 1, 'mish')
        # r -1 -10
        self.conv5 = Conv_Bn_Activation(128, 128, 1, 1, 'mish')

    def forward(self, input):
        x1 = self.conv1(input)
        x2 = self.conv2(x1)
        x3 = self.conv3(x1)

        r = self.resblock(x3)
        x4 = self.conv4(r)

        x4 = torch.cat([x4, x2], dim=1)
        x5 = self.conv5(x4)
        return x5

class DownSample3(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv_Bn_Activation(128, 256, 3, 2, 'mish')
        self.conv2 = Conv_Bn_Activation(256, 128, 1, 1, 'mish')
        self.conv3 = Conv_Bn_Activation(256, 128, 1, 1, 'mish')

        self.resblock = ResBlock(ch=128, nblocks=8)
        self.conv4 = Conv_Bn_Activation(128, 128, 1, 1, 'mish')
        self.conv5 = Conv_Bn_Activation(256, 256, 1, 1, 'mish')

    def forward(self, input):
        x1 = self.conv1(input)
        x2 = self.conv2(x1)
        x3 = self.conv3(x1)

        r = self.resblock(x3)
        x4 = self.conv4(r)

        x4 = torch.cat([x4, x2], dim=1)
        x5 = self.conv5(x4)
        return x5

class DownSample4(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv_Bn_Activation(256, 512, 3, 2, 'mish')
        self.conv2 = Conv_Bn_Activation(512, 256, 1, 1, 'mish')
        self.conv3 = Conv_Bn_Activation(512, 256, 1, 1, 'mish')

        self.resblock = ResBlock(ch=256, nblocks=8)
        self.conv4 = Conv_Bn_Activation(256, 256, 1, 1, 'mish')
        self.conv5 = Conv_Bn_Activation(512, 512, 1, 1, 'mish')

    def forward(self, input):
        x1 = self.conv1(input)
        x2 = self.conv2(x1)
        x3 = self.conv3(x1)

        r = self.resblock(x3)
        x4 = self.conv4(r)

        x4 = torch.cat([x4, x2], dim=1)
        x5 = self.conv5(x4)
        return x5

class DownSample5(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv_Bn_Activation(512, 1024, 3, 2, 'mish')
        self.conv2 = Conv_Bn_Activation(1024, 512, 1, 1, 'mish')
        self.conv3 = Conv_Bn_Activation(1024, 512, 1, 1, 'mish')

        self.resblock = ResBlock(ch=512, nblocks=4)
        self.conv4 = Conv_Bn_Activation(512, 512, 1, 1, 'mish')
        self.conv5 = Conv_Bn_Activation(1024, 1024, 1, 1, 'mish')

    def forward(self, input):
        x1 = self.conv1(input)
        x2 = self.conv2(x1)
        x3 = self.conv3(x1)

        r = self.resblock(x3)
        x4 = self.conv4(r)

        x4 = torch.cat([x4, x2], dim=1)
        x5 = self.conv5(x4)
        return x5

class Neck(nn.Module):
    def __init__(self, inference=False):
        super().__init__()
        self.inference = inference

        self.conv1 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
        self.conv2 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
        self.conv3 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
        # SPP
        self.maxpool1 = nn.MaxPool2d(kernel_size=5, stride=1, padding=5 // 2)
        self.maxpool2 = nn.MaxPool2d(kernel_size=9, stride=1, padding=9 // 2)
        self.maxpool3 = nn.MaxPool2d(kernel_size=13, stride=1, padding=13 // 2)

        # R -1 -3 -5 -6
        # SPP
        self.conv4 = Conv_Bn_Activation(2048, 512, 1, 1, 'leaky')
        self.conv5 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
        self.conv6 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
        self.conv7 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        # UP
        self.upsample1 = Upsample()
        # R 85
        self.conv8 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        # R -1 -3
        self.conv9 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        self.conv10 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
        self.conv11 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        self.conv12 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
        self.conv13 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        self.conv14 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
        # UP
        self.upsample2 = Upsample()
        # R 54
        self.conv15 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
        # R -1 -3
        self.conv16 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
        self.conv17 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
        self.conv18 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')
        self.conv19 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
        self.conv20 = Conv_Bn_Activation(256, 128, 1, 1, 'leaky')

    def forward(self, input, downsample4, downsample3, inference=False):
        x1 = self.conv1(input)
        x2 = self.conv2(x1)
        x3 = self.conv3(x2)
        # SPP
        m1 = self.maxpool1(x3)
        m2 = self.maxpool2(x3)
        m3 = self.maxpool3(x3)
        spp = torch.cat([m3, m2, m1, x3], dim=1)
        # SPP end
        x4 = self.conv4(spp)
        x5 = self.conv5(x4)
        x6 = self.conv6(x5)
        x7 = self.conv7(x6)
        # UP
        up = self.upsample1(x7, downsample4.size(), self.inference)
        # R 85
        x8 = self.conv8(downsample4)
        # R -1 -3
        x8 = torch.cat([x8, up], dim=1)

        x9 = self.conv9(x8)
        x10 = self.conv10(x9)
        x11 = self.conv11(x10)
        x12 = self.conv12(x11)
        x13 = self.conv13(x12)
        x14 = self.conv14(x13)

        # UP
        up = self.upsample2(x14, downsample3.size(), self.inference)
        # R 54
        x15 = self.conv15(downsample3)
        # R -1 -3
        x15 = torch.cat([x15, up], dim=1)

        x16 = self.conv16(x15)
        x17 = self.conv17(x16)
        x18 = self.conv18(x17)
        x19 = self.conv19(x18)
        x20 = self.conv20(x19)
        return x20, x13, x6

class Yolov4s(nn.Module):
    def __init__(self, yolov4conv137weight=None, n_classes=80, inference=False):
        super().__init__()

        output_ch = (4 + 1 + n_classes) * 3

        # backbone
        self.down1 = DownSample1()
        self.down2 = DownSample2()
        self.down3 = DownSample3()

        self.down4 = DownSample4()
        self.down5 = DownSample5()
        # neck
        self.neek = Neck(inference)
        # yolov4conv137
        if yolov4conv137weight:
            _model = nn.Sequential(self.down1, self.down2, self.down3, self.down4, self.down5, self.neek)
            pretrained_dict = torch.load(yolov4conv137weight)

            model_dict = _model.state_dict()
            # 1. filter out unnecessary keys
            pretrained_dict = {k1: v for (k, v), k1 in zip(pretrained_dict.items(), model_dict)}
            # 2. overwrite entries in the existing state dict
            model_dict.update(pretrained_dict)
            _model.load_state_dict(model_dict)

        # head
        self.head = Yolov4Heads(output_ch, n_classes, inference)

    def forward(self, input):
        d1 = self.down1(input)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)

        x20, x13, x6 = self.neek(d5, d4, d3)

        output = self.head(x20, x13, x6)

        return output

def get_region_boxes(boxes_and_confs):
    # print('Getting boxes from boxes and confs ...')

    boxes_list = []
    confs_list = []
    confs_w_list = []

    for item in boxes_and_confs:
        boxes_list.append(item[0])

        confs_w_list.append(item[1])

    # boxes: [batch, num1 + num2 + num3, 1, 4]
    # confs: [batch, num1 + num2 + num3, num_classes]
    boxes = torch.cat(boxes_list, dim=1)
    confs_w = torch.cat(confs_w_list, dim=1)

    return [boxes,  confs_w]

class Yolov4Heads(nn.Module):
    def __init__(self, output_ch, n_classes, inference=False):
        super().__init__()
        self.inference = inference

        self.conv1 = Conv_Bn_Activation(128, 256, 3, 1, 'leaky')
        self.conv2 = Conv_Bn_Activation(256, output_ch, 1, 1, 'linear', bn=False, bias=True)

        self.yolo1 = YoloLayers(
            anchor_mask=[0, 1, 2], num_classes=n_classes,
            anchors=[12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401],
            num_anchors=9, stride=8)

        # R -4
        self.conv3 = Conv_Bn_Activation(128, 256, 3, 2, 'leaky')

        # R -1 -16
        self.conv4 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        self.conv5 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
        self.conv6 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        self.conv7 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
        self.conv8 = Conv_Bn_Activation(512, 256, 1, 1, 'leaky')
        self.conv9 = Conv_Bn_Activation(256, 512, 3, 1, 'leaky')
        self.conv10 = Conv_Bn_Activation(512, output_ch, 1, 1, 'linear', bn=False, bias=True)

        self.yolo2 = YoloLayers(
            anchor_mask=[3, 4, 5], num_classes=n_classes,
            anchors=[12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401],
            num_anchors=9, stride=16)

        # R -4
        self.conv11 = Conv_Bn_Activation(256, 512, 3, 2, 'leaky')

        # R -1 -37
        self.conv12 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
        self.conv13 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
        self.conv14 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
        self.conv15 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
        self.conv16 = Conv_Bn_Activation(1024, 512, 1, 1, 'leaky')
        self.conv17 = Conv_Bn_Activation(512, 1024, 3, 1, 'leaky')
        self.conv18 = Conv_Bn_Activation(1024, output_ch, 1, 1, 'linear', bn=False, bias=True)

        self.yolo3 = YoloLayers(
            anchor_mask=[6, 7, 8], num_classes=n_classes,
            anchors=[12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401],
            num_anchors=9, stride=32)

    def forward(self, input1, input2, input3):
        x1 = self.conv1(input1)
        x2 = self.conv2(x1)

        x3 = self.conv3(input1)
        # R -1 -16
        x3 = torch.cat([x3, input2], dim=1)
        x4 = self.conv4(x3)
        x5 = self.conv5(x4)
        x6 = self.conv6(x5)
        x7 = self.conv7(x6)
        x8 = self.conv8(x7)
        x9 = self.conv9(x8)
        x10 = self.conv10(x9)

        # R -4
        x11 = self.conv11(x8)
        # R -1 -37
        x11 = torch.cat([x11, input3], dim=1)

        x12 = self.conv12(x11)
        x13 = self.conv13(x12)
        x14 = self.conv14(x13)
        x15 = self.conv15(x14)
        x16 = self.conv16(x15)
        x17 = self.conv17(x16)
        x18 = self.conv18(x17)

        if self.inference:
            y1 = self.yolo1(x2)
            y2 = self.yolo2(x10)
            y3 = self.yolo3(x18)

            return get_region_boxes([y1, y2, y3])

        else:
            return [x2, x10, x18]

class YoloLayers(nn.Module):
    ''' Yolo layer
    model_out: while inference,is post-processing inside or outside the model
        true:outside
    '''
    def __init__(self, anchor_mask=[], num_classes=0, anchors=[], num_anchors=1, stride=32, model_out=False):
        super(YoloLayers, self).__init__()
        self.anchor_mask = anchor_mask
        self.num_classes = num_classes
        self.anchors = anchors
        self.num_anchors = num_anchors
        self.anchor_step = len(anchors) // num_anchors
        self.coord_scale = 1
        self.noobject_scale = 1
        self.object_scale = 5
        self.class_scale = 1
        self.thresh = 0.6
        self.stride = stride
        self.seen = 0
        self.scale_x_y = 1

        self.model_out = model_out

    def forward(self, output, target=None):
        if self.training:
            return output
        masked_anchors = []
        for m in self.anchor_mask:
            masked_anchors += self.anchors[m * self.anchor_step:(m + 1) * self.anchor_step]
        masked_anchors = [anchor / self.stride for anchor in masked_anchors]

        return yolo_forward_dynamics(output, self.thresh, self.num_classes, masked_anchors, len(self.anchor_mask),scale_x_y=self.scale_x_y)

def yolo_forward_dynamics(output, conf_thresh, num_classes, anchors, num_anchors, scale_x_y, only_objectness=1,
                         validation=False):
    # Output would be invalid if it does not satisfy this assert
    # assert (output.size(1) == (5 + num_classes) * num_anchors)

    # Slice the second dimension (channel) of output into:
    # [ 2, 2, 1, num_classes, 2, 2, 1, num_classes, 2, 2, 1, num_classes ]
    # And then into
    # bxy = [ 6 ] bwh = [ 6 ] det_conf = [ 3 ] cls_conf = [ num_classes * 3 ]
    # batch = output.size(0)
    # H = output.size(2)
    # W = output.size(3)

    bxy_list = []
    bwh_list = []
    det_confs_list = []
    cls_confs_list = []

    for i in range(num_anchors):
        begin = i * (5 + num_classes)
        end = (i + 1) * (5 + num_classes)

        bxy_list.append(output[:, begin: begin + 2])
        bwh_list.append(output[:, begin + 2: begin + 4])
        det_confs_list.append(output[:, begin + 4: begin + 5])
        cls_confs_list.append(output[:, begin + 5: end])

    # Shape: [batch, num_anchors * 2, H, W]
    bxy = torch.cat(bxy_list, dim=1)
    # Shape: [batch, num_anchors * 2, H, W]
    bwh = torch.cat(bwh_list, dim=1)

    # Shape: [batch, num_anchors, H, W]
    det_confs = torch.cat(det_confs_list, dim=1)
    # Shape: [batch, num_anchors * H * W]
    det_confs = det_confs.view(output.size(0), num_anchors * output.size(2) * output.size(3))

    # Shape: [batch, num_anchors * num_classes, H, W]
    cls_confs = torch.cat(cls_confs_list, dim=1)
    # Shape: [batch, num_anchors, num_classes, H * W]
    cls_confs = cls_confs.view(output.size(0), num_anchors, num_classes, output.size(2) * output.size(3))
    # Shape: [batch, num_anchors, num_classes, H * W] --> [batch, num_anchors * H * W, num_classes]
    cls_confso = cls_confs.permute(0, 1, 3, 2).reshape(output.size(0), num_anchors * output.size(2) * output.size(3),
                                                       num_classes)

    # Apply sigmoid(), exp() and softmax() to slices
    #
    bxy = torch.sigmoid(bxy) * scale_x_y - 0.5 * (scale_x_y - 1)
    bwh = torch.exp(bwh)
    det_confs = torch.sigmoid(det_confs)
    cls_confs = torch.sigmoid(cls_confso)

    # Prepare C-x, C-y, P-w, P-h (None of them are torch related)
    grid_x = np.expand_dims(np.expand_dims(
        np.expand_dims(np.linspace(0, output.size(3) - 1, output.size(3)), axis=0).repeat(output.size(2), 0), axis=0),
                            axis=0)
    grid_y = np.expand_dims(np.expand_dims(
        np.expand_dims(np.linspace(0, output.size(2) - 1, output.size(2)), axis=1).repeat(output.size(3), 1), axis=0),
                            axis=0)
    # grid_x = torch.linspace(0, W - 1, W).reshape(1, 1, 1, W).repeat(1, 1, H, 1)
    # grid_y = torch.linspace(0, H - 1, H).reshape(1, 1, H, 1).repeat(1, 1, 1, W)

    anchor_w = []
    anchor_h = []
    for i in range(num_anchors):
        anchor_w.append(anchors[i * 2])
        anchor_h.append(anchors[i * 2 + 1])

    device = None
    cuda_check = output.is_cuda
    if cuda_check:
        device = output.get_device()

    bx_list = []
    by_list = []
    bw_list = []
    bh_list = []

    # Apply C-x, C-y, P-w, P-h
    for i in range(num_anchors):
        ii = i * 2
        # Shape: [batch, 1, H, W]
        bx = bxy[:, ii: ii + 1] + torch.tensor(grid_x, device=device,
                                               dtype=torch.float32)  # grid_x.to(device=device, dtype=torch.float32)
        # Shape: [batch, 1, H, W]
        by = bxy[:, ii + 1: ii + 2] + torch.tensor(grid_y, device=device,
                                                   dtype=torch.float32)  # grid_y.to(device=device, dtype=torch.float32)
        # Shape: [batch, 1, H, W]
        bw = bwh[:, ii: ii + 1] * anchor_w[i]
        # Shape: [batch, 1, H, W]
        bh = bwh[:, ii + 1: ii + 2] * anchor_h[i]

        bx_list.append(bx)
        by_list.append(by)
        bw_list.append(bw)
        bh_list.append(bh)

    ########################################
    #   Figure out bboxes from slices     #
    ########################################

    # Shape: [batch, num_anchors, H, W]
    bx = torch.cat(bx_list, dim=1)
    # Shape: [batch, num_anchors, H, W]
    by = torch.cat(by_list, dim=1)
    # Shape: [batch, num_anchors, H, W]
    bw = torch.cat(bw_list, dim=1)
    # Shape: [batch, num_anchors, H, W]
    bh = torch.cat(bh_list, dim=1)

    # Shape: [batch, 2 * num_anchors, H, W]
    bx_bw = torch.cat((bx, bw), dim=1)
    # Shape: [batch, 2 * num_anchors, H, W]
    by_bh = torch.cat((by, bh), dim=1)

    # normalize coordinates to [0, 1]
    bx_bw /= output.size(3)
    by_bh /= output.size(2)

    # Shape: [batch, num_anchors * H * W, 1]
    bx = bx_bw[:, :num_anchors].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
    by = by_bh[:, :num_anchors].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
    bw = bx_bw[:, num_anchors:].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
    bh = by_bh[:, num_anchors:].view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)

    bx1 = bx - bw * 0.5
    by1 = by - bh * 0.5
    bx2 = bx1 + bw
    by2 = by1 + bh

    # Shape: [batch, num_anchors * h * w, 4] -> [batch, num_anchors * h * w, 1, 4]
    boxes = torch.cat((bx1, by1, bx2, by2), dim=2).view(output.size(0), num_anchors * output.size(2) * output.size(3),
                                                        1, 4)
    # boxes = boxes.repeat(1, 1, num_classes, 1)

    # boxes:     [batch, num_anchors * H * W, 1, 4]
    # cls_confs: [batch, num_anchors * H * W, num_classes]
    # det_confs: [batch, num_anchors * H * W]

    det_confs = det_confs.view(output.size(0), num_anchors * output.size(2) * output.size(3), 1)
    confs = cls_confs * det_confs
    confs_w = torch.cat((confs, det_confs), dim = -1)
    #print("confs_w-confs", confs_w.size() ,confs.size())

    # boxes: [batch, num_anchors * H * W, 1, 4]
    # confs: [batch, num_anchors * H * W, num_classes]

    return boxes, confs_w