from betamax import Betamax
from github_contents import GithubContents
import pathlib
import os
import json

TOKEN = os.environ.get("TEST_GITHUB_TOKEN") or "XXX"

# When recording new tests, run the test suite with
# VCR_RECORD=new_episodes TEST_GITHUB_TOKEN=... pytest -k name_of_test
# https://betamax.readthedocs.io/en/latest/record_modes.html
VCR_RECORD = os.environ.get("VCR_RECORD") or "none"

with Betamax.configure() as config:
    config.cassette_library_dir = "cassettes"

def test_read_small_file():
    github = GithubContents("simonw", "github-contents-demo", TOKEN, branch="master")
    with Betamax(github.session) as vcr:
        vcr.use_cassette("get-file")
        content, sha = github.read("hello.txt")
    assert b"hello world 3" == content
    assert "3c840b722385abe67a2cfadac6a8eaab8429a45c" == sha

def test_write_small_file():
    github = GithubContents("simonw", "github-contents-demo", TOKEN, branch="master")
    with Betamax(github.session) as vcr:
        vcr.use_cassette("write-file")
        assert (
            "3c840b722385abe67a2cfadac6a8eaab8429a45c",
            "ba288b9c0a35409bf655bc85b5b6b04db7bb7742",
        ) == github.write(
            "hello.txt",
            b"hello world 3",
            commit_message="updated by test",
            committer={"name": "Test", "email": "test"},
        )

def test_write_large():
    github = GithubContents("simonw", "github-contents-demo", TOKEN, branch="master")
    with Betamax(github.session) as vcr:
        vcr.use_cassette("write-file-large")
        assert (
            "659974e29c1fd07322c9adf7de6636c6c0c8f9d8",
            "16b273de7090b618be4924e05b5763a9f4c4c69b",
        ) == github.write_large(
            "write_large.txt",
            b"not actually large but I did use the .write_large() method",
            commit_message="written by test",
            committer={"name": "Test", "email": "test"},
        )

def test_branch_exists_returns_false_for_nonexistant_branch():
    github = GithubContents(
        "simonw", "github-contents-demo", TOKEN, branch="not-a-branch"
    )
    with Betamax(github.session) as vcr:
        vcr.use_cassette("branch-exists-false.json", record=VCR_RECORD)
        assert not github.branch_exists()

def test_branch_exists_returns_true_for_existing_branch():
    github = GithubContents(
        "simonw", "github-contents-demo", TOKEN, branch="demo-branch"
    )
    with Betamax(github.session) as vcr:
        vcr.use_cassette("branch-exists-true.json", record=VCR_RECORD)
        assert github.branch_exists()

def test_no_accidental_oauth_token_in_cassettes():
    # Need to make sure I never accidentally check my OAuth token into GitHub
    for filepath in pathlib.Path("cassettes").glob("*.json"):
        data = json.loads(filepath.read_text())
        for interaction in data["http_interactions"]:
            for header in interaction["request"]["headers"].get("Authorization") or []:
                assert "token XXX" == header


#!/usr/bin/python3
"""
This module includes the base class
"""
import json
import csv
import turtle

class Base:
    """
    The base class
    """
    __nb_objects = 0

    def __init__(self, id=None):
        """
        initialization function
        assign the id of the object if provided
        otherwise assign with the next object number
        """

        if id:
            self.id = id
        else:
            Base.nb_objects = self.nb_objects + 1
            self.id = self.nb_objects

    @property
    def nb_objects(self):
        """
        getter function of nb_objects
        private attribute (number of objects)
        """

        return self.__nb_objects

    @nb_objects.setter
    def nb_objects(self, value):
        """
        number of objects setter
        """

        if not isinstance(value, int):
            raise TypeError("Error: value must be an integer")
        else:
            self.__nb_objects = value

    @staticmethod
    def to_json_string(list_dictionaries):
        """
        return json representation
        of list_dictionaries
        """

        if list_dictionaries and len(list_dictionaries):
            return json.dumps(list_dictionaries)
        else:
            return "[]"

    @classmethod
    def save_to_file(cls, list_objs):
        """
        write json representation of instances
        to file
        """

        file_name = f"{cls.__name__}.json"
        list_dictionaries = []
        if list_objs is not None:
            list_dictionaries = [obj.to_dictionary() for obj in list_objs]

        list_json = Base.to_json_string(list_dictionaries)
        with open(file_name, 'w', encoding='utf-8') as file:
            file.write(list_json)

    @staticmethod
    def from_json_string(json_string):
        """
        return json representation of
        json string
        """

        if json_string and len(json_string):
            return json.loads(json_string)
        else:
            return []

    @classmethod
    def create(cls, **dictionary):
        """
        return Rectangle instance
        from dictionary
        """

        instance = None
        if cls.__name__ == "Rectangle":
            instance = cls(2, 1)
        else:
            instance = cls(1)

        instance.update(**dictionary)
        return instance

    @classmethod
    def load_from_file(cls):
        """
        return an instances of cls
        from json file
        """

        file_string = None
        file_name = f"{cls.__name__}.json"
        try:
            with open(file_name, 'r', encoding='utf-8') as file:
                file_string = file.read()
        except FileNotFoundError:
            return []

        list_objs = cls.from_json_string(file_string)

        list_instances = list()
        for obj in list_objs:
            list_instances.append(cls.create(**obj))

        return list_instances

    @classmethod
    def save_to_file_csv(cls, list_objs):
        """
        saves list of objects
        to csv file
        """

        file_name = f"{cls.__name__}.csv"

        list_dictionaries = []
        if list_objs is not None:
            list_dictionaries = [obj.to_dictionary() for obj in list_objs]

        keys = list_dictionaries[0].keys()
        with open(file_name, 'w', encoding="utf-8", newline='\n') as file:
            writer = csv.DictWriter(file, fieldnames=keys)
            writer.writeheader()
            writer.writerows(list_dictionaries)

    @classmethod
    def load_from_file_csv(cls):
        """
        load objects from csv file
        """

        file_name = f"{cls.__name__}.csv"
        list_instances = []
        try:
            with open(file_name, 'r', encoding='utf-8', newline='\n') as file:
                reader = csv.DictReader(file)
                for row in reader:
                    for key, value in row.items():
                        row[key] = int(value)
                    list_instances.append(cls.create(**row))
        except FileNotFoundError:
            return list_instances

        return list_instances

    @staticmethod
    def draw(list_rectangles, list_squares):
        """
        draw a list of rectangles and squares
        """

        turtle.penup()
        for shape in list_rectangles + list_squares:
            turtle.goto(shape.x, shape.y)
            turtle.pendown()
            for _ in range(2):
                turtle.forward(shape.width)
                turtle.right(90)
                turtle.forward(shape.height)
                turtle.right(90)
            turtle.penup()

        turtle.done()


import os
import time

import numpy as np
import cv2
import vtk
import argparse

from scale_recovery import project_to_3d
from visual.vtk_wrapper.vtk_point_cloud_glyph import VtkPointCloudGlyph
from visual.vtk_wrapper import vtk_utils

def setup_vtk_renderer():
    # Setup renderer
    vtk_renderer = vtk.vtkRenderer()
    vtk_renderer.SetBackground(0.2, 0.3, 0.4)

    return vtk_renderer

def setup_vtk_render_window(window_name, window_size, vtk_renderer):
    # Setup Render Window
    vtk_render_window = vtk.vtkRenderWindow()
    vtk_render_window.SetWindowName(window_name)
    vtk_render_window.SetSize(*window_size)
    vtk_render_window.AddRenderer(vtk_renderer)
    return vtk_render_window

def get_velo_points(raw_data, frame_idx):
    velo_points = raw_data.get_velo(frame_idx)

    # Filter points to certain area
    points = velo_points[:, 0:3]
    area_extents = np.asarray([[0, 100], [-50, 50], [-5, 1]], dtype=np.float32)
    area_filter = \
        (points[:, 0] > area_extents[0, 0]) & \
        (points[:, 0] < area_extents[0, 1]) & \
        (points[:, 1] > area_extents[1, 0]) & \
        (points[:, 1] < area_extents[1, 1]) & \
        (points[:, 2] > area_extents[2, 0]) & \
        (points[:, 2] < area_extents[2, 1])
    points = points[area_filter]

    return points

def visual_depth_real_scale(depth_dir, img_dir, out_dir, width=736, height=416):

    scale_x, scale_y = 736/1920, 416/1080
    cx,cy,fx,fy = 851.8287*scale_x, 424.0041*scale_y, 1998.7356*scale_x, 1991.8909*scale_y

    vtk_window_size = (1280, 720)
    max_fps = 10000.0
    file_list = sorted(os.listdir(depth_dir))
    frame_range = (0, len(file_list))
    ##########################################
    min_loop_time = 1.0 / max_fps
    vtk_renderer = setup_vtk_renderer()
    vtk_render_window = setup_vtk_render_window('Overlaid Point Cloud', vtk_window_size, vtk_renderer)
    current_cam = vtk_renderer.GetActiveCamera()
    current_cam.SetViewUp(0, 0, 1)    # 0, -1, 0
    current_cam.SetPosition(0, -10, -20)  #  0.0, 0, 0
    current_cam.SetFocalPoint(0.0, 0.0, 20)
    current_cam.SetViewAngle(50)
    current_cam.Zoom(0.9)

    # # Create VtkAxes
    vtk_axes = vtk.vtkAxesActor()
    vtk_axes.SetTotalLength(1, 1, 1)

    # Setup interactor
    vtk_interactor = vtk.vtkRenderWindowInteractor()
    vtk_interactor.SetRenderWindow(vtk_render_window)
    vtk_interactor.SetInteractorStyle(vtk_utils.ToggleActorsInteractorStyle(None, vtk_renderer, current_cam, vtk_axes))
    vtk_interactor.Initialize()

    # Point cloud
    vtk_pc = VtkPointCloudGlyph()

    # # Add actors
    vtk_renderer.AddActor(vtk_pc.vtk_actor)

   ###########################################
    for frame_idx in range(*frame_range):
        loop_start_time = time.time()
        print('{} / {}'.format(frame_idx, len(file_list) - 1))

        # Load next frame data
        load_start_time = time.time()

        frame = file_list[frame_idx]
        depth_map_path = os.path.join(depth_dir, frame)
        depth_map = np.load(depth_map_path)  # 736*416

        img_name = frame.split('.')[0]
        img_path = img_name +'.jpg'
        raw_image_path = os.path.join(img_dir, img_path)
        bgr_image = cv2.imread(raw_image_path)
        bgr_image = cv2.resize(bgr_image, (width,height))
        # print(raw_image_path)
        cam0_curr_pc = project_to_3d(depth_map, cx, cy, fx, fy, width, height)
        point_colours = bgr_image.reshape(-1, 3)

        print('load\t\t', time.time() - load_start_time)

        # VtkPointCloud
        vtk_pc_start_time = time.time()
        vtk_pc.set_points(cam0_curr_pc, point_colours)
        print('vtk_pc\t\t', time.time() - vtk_pc_start_time)

        # Reset the clipping range to show all points
        vtk_renderer.ResetCameraClippingRange()

        out_depth_map_path = os.path.join(out_dir, img_name + '.png')
        grabber = vtk.vtkWindowToImageFilter()
        grabber.SetInput(vtk_render_window)
        grabber.Update()
        writer = vtk.vtkPNGWriter()
        writer.SetInputData(grabber.GetOutput())
        writer.SetFileName(out_depth_map_path)
        writer.Write()

        # Render
        render_start_time = time.time()
        vtk_render_window.Render()
        print('render\t\t', time.time() - render_start_time)

        # Pause to keep frame rate under max
        loop_run_time = time.time() - loop_start_time
        print('loop\t\t', loop_run_time)
        if loop_run_time < min_loop_time:
            time.sleep(min_loop_time - loop_run_time)

        print('---')

    print('Done')

    # Keep window open
    vtk_interactor.Start()

def main():
    depth_dir = '../dataset/hs5_0831/depth_map/hs5_35k_3_east_park'
    img_dir = '../dataset/hs5_0831/image_undistort'
    out_dir = './output/hs5_35k_3_east_park'
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    width = 736
    height = 416
    visual_depth_real_scale(depth_dir, img_dir, out_dir, width, height)

if __name__ == '__main__':
    main()


#from django.contrib import admin
from django.conf.urls import url, include
from django.conf import settings
from django.conf.urls.static import serve

urlpatterns = [
    #path('admin/', admin.site.urls),
    url(r'^', include('modulos.inicio.urls', namespace='inicio')),

    url(r'^static/(?P<path>.*)$', serve, {
        'document_root': settings.STATIC_ROOT } ),

    url(r'^media/(?P<path>.*)$', serve, {
    	'document_root': settings.MEDIA_ROOT, }),
]


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

# In[1]:

import pandas as pd
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, plot_tree
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
from sklearn.neighbors import KNeighborsClassifier

# In[2]:

from sklearn.datasets import load_iris
iris_data = load_iris()

# In[3]:

print(iris_data.DESCR)

# In[4]:

df = pd.DataFrame(iris_data.data)
df.columns = iris_data.feature_names
df.head()

# In[5]:

df['Species'] = iris_data.target
df.head()

# In[6]:

df['Species'].unique()

# In[7]:

print("Species name:", iris_data['target_names'])

# In[8]:

df.shape

# In[9]:

df.info()

# In[10]:

df.isnull().sum()

# In[11]:

df.describe()

# In[12]:

df['Species'].value_counts()

# In[13]:

sns.pairplot(df, hue='Species', markers='+', palette='colorblind')
plt.show()

# In[14]:

plt.figure(figsize=(10,7))
sns.heatmap(df.corr(), annot=True)
plt.show()

# In[15]:

df.plot(kind='scatter', x="sepal length (cm)", y="sepal width (cm)")
plt.show()

# In[16]:

fig=df[df.Species==0].plot(kind='scatter', x="sepal length (cm)", y="sepal width (cm)",
                           color='orange', label='Setosa')
df[df.Species==1].plot(kind='scatter', x="sepal length (cm)", y="sepal width (cm)",
                           color='blue', label='Versicolor', ax=fig)
df[df.Species==2].plot(kind='scatter', x="sepal length (cm)", y="sepal width (cm)",
                           color='green', label='Virginica', ax=fig)
fig.set_xlabel("Sepal Length")
fig.set_ylabel("Sepal Width")
fig.set_title("Sepal Length vs width")
plt.show()

# In[17]:

sns.violinplot(x="Species", y='sepal length (cm)', data=df)

# In[18]:

plt.figure(figsize=(12,7))
plt.subplot(2,2,1)
sns.violinplot(x="Species", y='sepal length (cm)', data=df)
plt.subplot(2,2,2)
sns.violinplot(x="Species", y='sepal width (cm)', data=df)
plt.subplot(2,2,3)
sns.violinplot(x="Species", y='petal length (cm)', data=df)
plt.subplot(2,2,4)
sns.violinplot(x="Species", y='petal width (cm)', data=df)
plt.show()

# In[19]:

X = df.drop('Species', axis=1)
y = df['Species']

# In[20]:

X_train, X_test, y_train, y_test = train_test_split(X,y, test_size=0.2, random_state=42)

# In[21]:

print("X_train:", X_train.shape)
print("X_test:", X_test.shape)
print("y_train:", y_train.shape)
print("y_test:", y_test.shape)

# In[22]:

dtree = DecisionTreeClassifier()
dtree.fit(X_train, y_train)
dtree_pred = dtree.predict(X_test)
dtree_acc = accuracy_score(dtree_pred, y_test)
print("Test accuracy: {:.2f}%".format(dtree_acc*100))

# In[23]:

print(classification_report(y_test, dtree_pred))

# In[24]:

print(confusion_matrix(y_test, dtree_pred))

# In[26]:

plt.figure(figsize=(15,10))
plot_tree(dtree, feature_names=iris_data.feature_names, class_names = iris_data['target_names'])

# In[27]:

knn_model = KNeighborsClassifier(n_neighbors=3)
knn_model.fit(X_train, y_train)
knn_pred = knn_model.predict(X_test)
knn_acc = accuracy_score(knn_pred, y_test)
print("Test accuracy: {:.2f}%".format(knn_acc*100))

# In[28]:

print(confusion_matrix(y_test, knn_pred))

# In[29]:

df.head()

# In[30]:

df.columns

# In[31]:

data = {'sepal length (cm)':5.0, 'sepal width (cm)':3.4, 'petal length (cm)':1.4,
       'petal width (cm)':0.2}
index = [0]
new_df = pd.DataFrame(data, index)
new_df

# In[32]:

value_pred = dtree.predict(new_df)
value_pred

Supervised machine learning on iris dataset LIVE.py


from django.forms import ModelForm, FileInput
from accounts.models import Order, Customer
from django.contrib.auth.models import User
from django.contrib.auth.forms import UserCreationForm

class OrderForm(ModelForm):
    class Meta:
        model = Order
        fields = '__all__'

class UpdateOrderForm(ModelForm):
    class Meta:
        model = Order
        fields = ['product','status']

class CustomerForm(ModelForm):
    class Meta:
        model = Customer
        fields = '__all__'
        exclude = ['user']
        widgets = {
            'porfile_pic': FileInput(),
        }

    def __init__(self, *args, **kwargs):
        super(CustomerForm, self).__init__(*args, **kwargs)
        self.fields['porfile_pic'].widget.attrs = {'id':'selectedFile'}

class CreateUserForm(UserCreationForm):
    class Meta:
        model  = User
        fields = ['username','email','password1','password2']


import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
import plotly as py

# interactive map bscc locations

# load data
df = pd.read_csv('data/locations_sorted.tsv', sep='\t')

# change long to lng
df = df.rename(columns={'long': 'lng'})

# change country codes to country names
df['country'] = df['country'].replace({'GB': 'Great Britain', 'CH': 'Switzerland'})

print(df)

fig = px.scatter_geo(df, lat='lat', lon='lng', hover_name='name', color='country', projection='natural earth',
                     size_max=10, opacity=0.5, title='BSCC Locations',
                     labels={'name': 'Name'},
                     hover_data={'Unnamed: 0': False, 'lat': False, 'lng': False, 'id': False, 'country': False, 'count': True})

fig.show()

# interactive map bscc atlas locations

# load data
df = pd.read_csv('data/bscc_atlas_locations.tsv', sep='\t')

# rename columns
df.columns = ['name', 'gnd', 'lat', 'lng', 'address', 'country', 'country_code', 'country_gnd', 'continent_code']

#insert for missing values in continent_code NA
df['continent_code'] = df['continent_code'].fillna('NA')

# change continent codes to continent names
df['continent_code'] = df['continent_code'].replace({'NA': 'North America (3)', 'SA': 'South America (5)', 'EU': 'Europe (15)', 'AS': 'Asia (29)', 'AF': 'Africa (5)', 'OC': 'Oceania'})

# remove unnecessary columns
df2 = df.drop(['gnd', 'address', 'country_code', 'country_gnd'], axis=1)

fig = px.scatter_geo(df2, lat='lat', lon='lng', hover_name='name', color='continent_code', projection='natural earth',
                     size_max=10, opacity=0.5, title='BSCC Atlas Locations',
                     labels={'name': 'Name'},
                     hover_data={'lat': False, 'lng': False, 'continent_code': False, 'country': False})

fig.update_traces(textposition='top center', mode='markers+text', textfont_size=8)

fig.update_traces(marker=dict(size=10))

fig.show()


# -*- coding: utf-8 -*-
"""
Created on Sun Oct  7 23:35:35 2018

@author: Ash
"""

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
from bs4 import BeautifulSoup
import csv
from time import gmtime, strftime

option = webdriver.ChromeOptions()
#option.add_argument(" — incognito")

url = "https://www.facebook.com/Shwapno.ACILL/"

browser = webdriver.Chrome(executable_path=r"F:\software\chromedriver", chrome_options=option)
browser.get(url)
soup = BeautifulSoup(browser.page_source, 'lxml')
#print(soup.prettify())

'''
with open('meenabazar_feature.csv', mode='a',encoding="utf-8",newline='') as meenabazar_feature:
    meenabazar_feature = csv.writer(meenabazar_feature, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
    meenabazar_feature.writerow(['ID','ProductName', 'Price','Date'])
    id=0

    for product,price in zip(productName,productPrice):
        id+=1
        print(product.getText().strip()+' '+price.getText().strip())
        meenabazar_feature.writerow([ str(id),product.getText().strip(),price.getText().strip().split()[1]+'TK',strftime("%Y-%m-%d", gmtime())])
'''

# Wait 20 seconds for page to load
timeout = 60
try:
    #WebDriverWait(browser, timeout).until(EC.visibility_of_element_located((By.XPATH, "//img[@class='avatar width-full rounded-2']")))
    #WebDriverWait(browser, timeout).until(EC.visibility_of_element_located((By.XPATH, "//img[@class='img-responsive']")))
    # find_elements_by_xpath returns an array of selenium objects.
    productName = soup.find_all('div',class_='UFILikeSentenceText')
    id = 0
    productPrice = soup.find_all('div',class_='_5pbx userContent _3576')
    for product in productName:
        print(product.getText().strip())

    for product,price in zip(productName,productPrice):
        id+=1
        print(product.getText().strip()+' '+price.getText().strip())
    titles_element = browser.find_elements_by_xpath("//div[@class='product ng-scope']")

    # use list comprehension to get the actual repo titles and not the selenium objects.
    titles = [x.text for x in titles_element]
    # print out all the titles.
    print('titles:')
    print(titles, '\n')

    language_element = browser.find_elements_by_xpath("//div[@class='price']")
    # same concept as for list-comprehension above.
    languages = [x.text for x in language_element]
    print("languages:")
    print(languages, '\n')

except TimeoutException:
    print("Timed out waiting for page to load")
    browser.quit()


from oidc_component import login_button as oidc_button
from oauth_component import login_button as oauth_button
import streamlit as st
from dotenv import load_dotenv
import os
load_dotenv()

#provider=os.environ['provider']
authority = os.environ['authority'] # 'https://oidc.io/oauth',
clientId = os.environ['clientId']
clientSecret = os.environ['clientSecret']
scope = os.environ['scope']

st.title('Welcome to OIDC-Streamlit')

with st.echo():
    user_info = oidc_button(clientId=clientId , clientSecret=clientSecret, authority = authority ,scope=scope)
    #user_info = oauth_button(clientId=clientId, clientSecret=clientSecret, provider=provider, scope=scope)
    if user_info:
        st.write(f'Hi {user_info["nickname"]}')
        # st.write(user_info) # some private information here

if not user_info:
    st.write("Please login to continue")


#python
from collections import namedtuple

# apetools
from apetools.baseclass import BaseClass
from apetools.commons import enumerations, errors
from parametertree import ParameterTree

AffectorTypes = enumerations.AffectorTypes
IperfDirection = enumerations.IperfDirection
ConfigurationError = errors.ConfigurationError

parameters = ("test_id repetition repetitions output_folder " +
              " receiver sender recovery_time affector").split()

class TestParameter(namedtuple('TestParameter', parameters)):
    """
    A TestParameter holds the settings for a single test-iteration
    """
    __slots__ = ()

    def __str__(self):
        return (self.__class__.__name__ + ":" +
                ','.join(("{f}:{v}".format(f=f, v=getattr(self,f))
                          for f in self._fields)))

# end class TestParameter

class ParameterGenerator(BaseClass):
    """
    A ParameterGenerator is an iterator that generates test-parameters.
    """
    def __init__(self, parameters, *args, **kwargs):
        """
        :param:

         - `parameters`: A list of parameter (namedtuple) lists
        """
        super(ParameterGenerator, self).__init__(*args, **kwargs)
        self.parameters = parameters
        self._tree = None
        return

    @property
    def tree(self):
        """
        :return: parameter-tree populated with parameters (possibly)
        """
        if self._tree is None:
            self._tree = ParameterTree(self.parameters)
        return self._tree

    def __iter__(self):
        """
        :yield: the next namedtuple of parameters
        """
        for parameters in self.tree.paths:
            yield parameters
        return
# end class ParameterGenerator


import random

import numpy as np
from torch.nn import MSELoss
from ignite.engine import Events
from ignite.handlers import ModelCheckpoint, EarlyStopping
from transformers import AdamW, get_linear_schedule_with_warmup

from handler import *
from param import model_param, train_param
from model import Model
from metric import metric
from dataset import train_loader, dev_loader

def train():
    set_seed(train_param.seed)
    model = Model(model_param)
    optimizer = AdamW(model.parameters(), lr=train_param.lr, eps=1e-8)
    update_steps = train_param.epoch * len(train_loader)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=update_steps)
    loss_fn = [translate, MSELoss()]
    device = torch.device(f'cuda:{train_param.device}')
    trainer = create_trainer(model, optimizer, scheduler, loss_fn, train_param.grad_norm, device)
    train_evaluator = create_evaluator(model, metric, device)
    dev_evaluator = create_evaluator(model, metric, device)
    trainer.add_event_handler(Events.ITERATION_COMPLETED(every=train_param.interval), log_training_loss)
    trainer.add_event_handler(Events.EPOCH_COMPLETED, log_results, *(train_evaluator, train_loader, 'Train'))
    trainer.add_event_handler(Events.EPOCH_COMPLETED, log_results, *(dev_evaluator, dev_loader, 'Dev'))
    es_handler = EarlyStopping(patience=train_param.patience, score_function=score_fn, trainer=trainer)
    dev_evaluator.add_event_handler(Events.COMPLETED, es_handler)
    ckpt_handler = ModelCheckpoint(train_param.save_path, '', score_function=score_fn,
                                   score_name='score', require_empty=False)
    dev_evaluator.add_event_handler(Events.COMPLETED, ckpt_handler, {'model': model, 'param': model_param})
    print(f'Start running {train_param.save_path.split("/")[-1]} at device: {train_param.device}\t'
          f'lr: {train_param.lr}')
    trainer.run(train_loader, max_epochs=train_param.epoch)

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.device_count() > 0:
        torch.cuda.manual_seed_all(seed)

def translate(t_inputs, masks, d_outputs):
    t_inputs = t_inputs[:, 1:].unsqueeze(2)
    masks = masks[:, 1:]
    loss = -torch.log(torch.gather(d_outputs.transpose(0, 1), 2, t_inputs)).squeeze(2)
    loss = torch.sum(loss * masks) / torch.sum(masks)
    return loss

if __name__ == '__main__':
    train()


import os
import argparse
import logging
import time
import numpy as np
import numpy.random as npr
import matplotlib

matplotlib.use("agg")
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import datetime

parser = argparse.ArgumentParser()
parser.add_argument("--adjoint", type=eval, default=True)
parser.add_argument("--visualize", type=eval, default=True)
parser.add_argument("--niters", type=int, default=2000)
parser.add_argument("--lr", type=float, default=0.01)
parser.add_argument("--gpu", type=int, default=0)
parser.add_argument("--noise", type=int, default=0.15)
parser.add_argument("--end_time", type=int, default=40)
parser.add_argument("--samples", type=int, default=100)
parser.add_argument("--async_data", type=int, default=True)
parser.add_argument("--vis_freq", type=int, default=10)
args = parser.parse_args()

# Another way of calculating the gradients
if args.adjoint:
    from torchdiffeq import odeint_adjoint as odeint
else:
    from torchdiffeq import odeint

def curve_function(timestamps, func_type):
    if func_type == 0:
        return np.sin(timestamps) * np.exp(-0.1 * timestamps)
    elif func_type == 1:
        return np.cos(0.5 * timestamps) * np.exp(-0.05 * timestamps)

def generate_curve(
    func_type,
    ncurves=1000,
    ntotal=500,
    nsample=100,
    start=0.0,
    stop=40,
    noise_std=0.1,
    savefig=True,
    async_data=True,
):
    """ Generate one curve
    Args:
      nspiral: number of spirals, i.e. batch dimension
      ntotal: total number of datapoints per spiral
      nsample: number of sampled datapoints for model fitting per spiral
      start: spiral starting theta value
      stop: spiral ending theta value
      noise_std: observation noise standard deviation
      a, b: parameters of the Archimedean spiral
      savefig: plot the ground truth for sanity check.

      Returns:
        Tuple where first element is true trajectory of size (nspiral, ntotal, 2),
        second element is noisy observations of size (nspiral, nsample, 2),
        third element is timestamps of size (ntotal,),
        and fourth element is timestamps of size(nsample,) """
    if async_data:
        # Create async. time tensor
        time_rnd = np.random.random_sample(nsample)
        time_rnd = np.sort(time_rnd)
        time_rnd = time_rnd * stop
        time = time_rnd
    else:
        # Or sync. time tensor
        time = np.linspace(0, stop, nsample)

    # Make sure that starting time 0 is the first time
    time[0] = 0.0
    # Create cont. time tensor
    time_cont = np.linspace(0, stop, ntotal)

    # Generate continuous y values
    true_y_cont = curve_function(time_cont, func_type)

    # Generate noisy y values
    dt = time[1:nsample] - time[0 : nsample - 1]
    dt = np.insert(dt, 0, 0)
    true_y_dt = []
    for _ in range(ncurves):
        noise = args.noise * (np.random.rand(nsample) - 0.5) * 2
        new_true_y = curve_function(time, func_type) + noise
        # One of the inputs is the diff. in time between points.
        new_true_dt = dt
        new_true_y_dt = np.stack((new_true_y, new_true_dt), axis=1)
        true_y_dt.append(new_true_y_dt)

    true_y_dt = np.asarray(true_y_dt)
    return true_y_cont, true_y_dt, time_cont, time

def create_plot():
    # Create plot
    if args.visualize:
        figSamp = plt.figure(figsize=(12, 4), facecolor="white")
        ax_traj_samp = figSamp.add_subplot(132, frameon=False)
        ax_traj_samp.cla()  # Clears
        ax_traj_samp.set_title("Ground Truth N' Samples")
        ax_traj_samp.set_xlabel("Time")
        ax_traj_samp.set_ylabel("Curves")
        return figSamp, ax_traj_samp

def plot_data(fig, subplt, y_values, timestamps, line_type, label_name):
    plt.plot(timestamps, y_values, line_type, label=label_name)
    fig.tight_layout()
    plt.legend()
    plt.draw()

# Plot some of the sensors timestamps vs each other
height_async_data = 0

def plot_unsync_sensors(timestamps, label_name):
    global height_async_data
    plt.plot(
        timestamps, height_async_data * np.ones(timestamps.size), "o", label=label_name
    )
    plt.ylim(-0.05, 2 * height_async_data)
    plt.legend()
    plt.draw()
    height_async_data += 0.05

def save_plot(file_name):
    plt.savefig("latent_ode_sync_fig/" + file_name + ".png", dpi=500)

# NODE in the latent space
class LatentNode(nn.Module):
    def __init__(self, latent_dim=4, nhidden=20):
        super(LatentNode, self).__init__()
        self.elu = nn.ELU(inplace=True)
        self.fc1 = nn.Linear(latent_dim, nhidden)
        self.fc2 = nn.Linear(nhidden, nhidden)
        self.fc3 = nn.Linear(nhidden, latent_dim)
        self.nfe = 0

    def forward(self, t, x):
        self.nfe += 1
        out = self.fc1(x)
        out = self.elu(out)
        out = self.fc2(out)
        out = self.elu(out)
        out = self.fc3(out)
        return out

# Bringing input in Y domain to lantent space Z with RNN
class RnnEncoder(nn.Module):
    def __init__(self, latent_dim=4, obs_dim=2, nhidden=25, nbatch=1):
        super(RnnEncoder, self).__init__()
        self.nhidden = nhidden
        self.batch = nbatch
        self.input_to_hidden = nn.Linear(obs_dim + nhidden, nhidden)
        self.hidden_to_output = nn.Linear(nhidden, latent_dim * 2)

    def forward(self, x, h):
        combined = torch.cat((x, h), dim=1)
        hidden_state = torch.tanh(self.input_to_hidden(combined))
        out = self.hidden_to_output(hidden_state)
        return out, hidden_state

    def init_hidden_state(self):
        return torch.zeros(self.batch, self.nhidden)

# Since NODE are space preserving we need to decode.
# Brings new points in latent Z space to Y space.
class Decoder(nn.Module):
    def __init__(self, latent_dim=4, obs_dim=3, nhidden=20):
        super(Decoder, self).__init__()
        self.relu = nn.ReLU(inplace=True)
        self.fc1 = nn.Linear(latent_dim, nhidden)
        self.fc2 = nn.Linear(nhidden, obs_dim)

    def forward(self, z):
        out = self.fc1(z)
        out = self.relu(out)
        out = self.fc2(out)
        return out

# Calculating average loss for better overview of training
class RunningAverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self, momentum=0.99):
        self.momentum = momentum
        self.reset()

    def reset(self):
        self.val = None
        self.avg = 0

    def update(self, val):
        if self.val is None:
            self.avg = val
        else:
            self.avg = self.avg * self.momentum + val * (1 - self.momentum)
        self.val = val

# Creating a pick from logarithmic normal probability distribution
# Basically providing the log of a relative likelihood that the value of the
# random variable would equal that sample.(wikipedia: Probability density function)
def log_normal_prob_dist(x, mean, logvar):
    const = torch.from_numpy(np.array([2.0 * np.pi])).float().to(x.device)
    const = torch.log(const)
    return -0.5 * (const + logvar + (x - mean) ** 2.0 / torch.exp(logvar))

# Kullback-Leibler divergence between a diagonal multivariate normal,
# and a normal distribution with specified mu and logvar.
# (Wikipedia: Kullback-Leibler Divergence, examples: Multivariate normal distribution)
def kullback_leibler_divergence_normaldist(mu1, lv1, mu2, lv2):
    v1 = torch.exp(lv1)
    v2 = torch.exp(lv2)
    lstd1 = lv1 / 2.0
    lstd2 = lv2 / 2.0

    kl = lstd2 - lstd1 + ((v1 + (mu1 - mu2) ** 2.0) / (2.0 * v2)) - 0.5
    return kl

# Combine timestamps into a single tensor
def combine_timestamps(t1, t2):
    t1 = t1.cpu().numpy()
    t2 = t2.cpu().numpy()
    indx1 = []
    indx2 = []
    combine_timestamps = np.concatenate((t1, t2), axis=0)
    combine_timestamps.sort()
    # Remove one of the common timepoints 0.0
    combine_timestamps = np.delete(combine_timestamps, 0)
    # Find index for the both original time tensors t1 and t2
    boolean_match1 = np.in1d(combine_timestamps, t1)
    boolean_match2 = np.in1d(combine_timestamps, t2)
    indx = np.linspace(0, combine_timestamps.size - 1, combine_timestamps.size)
    indx_t1 = indx[boolean_match1]
    indx_t2 = indx[boolean_match2]
    indx_t1 = indx_t1.astype(int)
    indx_t2 = indx_t2.astype(int)
    return combine_timestamps, indx_t1, indx_t2

# Fills in the holes in the original data
def forward_rnn_data(pred_x_A, pred_x_B, time_indx_A, time_indx_B, true_y_A, true_y_B):
    pred_x_A[:, time_indx_A] = true_y_A
    pred_x_B[:, time_indx_B] = true_y_B
    return pred_x_A, pred_x_B

if __name__ == "__main__":
    latent_dim_total = 4
    latent_dim_sensors = 2
    nhidden = 20
    rnn_nhidden = 25
    obs_dim_total = 2  # A and B sensors
    obs_dim_sensors = 2  # Sensor value and dt
    ncurves = 1000
    noise_std = 0.3
    device = torch.device(
        "cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu"
    )
    # Create folder if vis is true
    current_time = datetime.datetime.now()
    str_current_day = (
        str(current_time.year)
        + "_"
        + str(current_time.month)
        + "_"
        + str(current_time.day)
        + "_"
        + str(current_time.hour)
        + "_"
        + str(current_time.minute)
    )
    directory_name = "async_unsync_" + str_current_day
    if not os.path.exists(directory_name):
        os.makedirs(directory_name)

    # Generate two curves with different time vectors
    # Plot ground truth
    true_y_cont_A, true_y_dt_A, time_cont_A, time_A = generate_curve(func_type=0, ncurves=ncurves)
    fig, subplt = create_plot()
    plot_data(fig, subplt, true_y_cont_A, time_cont_A, "-", "Cont. A")
    true_y_cont_B, true_y_dt_B, time_cont_B, time_B = generate_curve(func_type=1, ncurves=ncurves)
    plot_data(fig, subplt, true_y_cont_B, time_cont_B, "-", "Cont. B")
    plot_data(fig, subplt, true_y_dt_A[0, :, 0], time_A, "o", "Samples A")
    plot_data(fig, subplt, true_y_dt_B[0, :, 0], time_B, "o", "Samples B")
    save_plot("ground_truth_samples")

    # Plot a sample to see the unsync. curves
    sample_A = time_A[0:10]
    sample_B = time_B[0:10]
    fig_time, subplt_time = create_plot()
    plot_unsync_sensors(sample_A, "Sensor A")
    plot_unsync_sensors(sample_B, "Sensor B")
    save_plot("unsync_sensors")

    # Make data into torch tensors at GPU, if possible
    true_y_cont_A = torch.from_numpy(true_y_cont_A).float().to(device)
    true_y_dt_A = torch.from_numpy(true_y_dt_A).float().to(device)
    time_cont_A = torch.from_numpy(time_cont_A).float().to(device)
    time_A = torch.from_numpy(time_A).float().to(device)
    true_y_A = true_y_dt_A[:, :, 0]

    true_y_cont_B = torch.from_numpy(true_y_cont_B).float().to(device)
    true_y_dt_B = torch.from_numpy(true_y_dt_B).float().to(device)
    time_cont_B = torch.from_numpy(time_cont_B).float().to(device)
    time_B = torch.from_numpy(time_B).float().to(device)
    true_y_B = true_y_dt_B[:, :, 0]

    # Fetching the common timestamps from sensors for prediction
    # Calculates the dt for forward rnn
    common_timestamps, time_indx_A, time_indx_B = combine_timestamps(time_A, time_B)
    common_timestamps = torch.from_numpy(common_timestamps).float().to(device)
    size = common_timestamps.size(0)
    common_dt = common_timestamps[1:size] - common_timestamps[0 : size - 1]
    common_dt = common_dt.cpu()
    common_dt = np.insert(common_dt, 0, 0)
    common_dt = torch.ones((ncurves, args.samples * 2 - 1)) * common_dt
    common_dt = common_dt.to(device)

    # Create models
    latent_node = LatentNode(latent_dim_total, nhidden).to(device)
    latent_node_fwd = LatentNode(latent_dim_total, nhidden).to(device)
    rnn_encoder_A = RnnEncoder(
        latent_dim_sensors, obs_dim_sensors, rnn_nhidden, ncurves
    ).to(device)
    rnn_encoder_B = RnnEncoder(
        latent_dim_sensors, obs_dim_sensors, rnn_nhidden, ncurves
    ).to(device)
    decoder = Decoder(latent_dim_total, obs_dim_total, nhidden).to(device)
    decoder_fwd = Decoder(latent_dim_total, obs_dim_total, nhidden).to(device)
    rnn_encoder_forw = RnnEncoder(
        latent_dim_total, obs_dim_total + 1, rnn_nhidden, ncurves
    ).to(device)
    params = (
        list(latent_node.parameters())
        + list(rnn_encoder_A.parameters())
        + list(rnn_encoder_B.parameters())
        + list(decoder.parameters())
    )
    optimizer = optim.Adam(params, lr=args.lr)
    # optimizer = optim.RMSprop(params, lr=1e-3)
    loss_meter = RunningAverageMeter()

    for itr in range(1, args.niters + 1):
        optimizer.zero_grad()
        # Run backwards in time through encoder for both A and B sensor
        h_A = rnn_encoder_A.init_hidden_state().to(device)
        h_B = rnn_encoder_B.init_hidden_state().to(device)

        training_times_A = reversed(range(int(np.trunc(true_y_dt_A.size(1)))))
        training_times_B = reversed(range(int(np.trunc(true_y_dt_B.size(1)))))
        # Run through sensor A
        for t in training_times_A:
            obs_A = true_y_dt_A[:, t, :]
            out_A, h_A = rnn_encoder_A.forward(obs_A, h_A)
        # picking out mean and log variance for A sensor
        qz0_mean_A, qz0_logvar_A = (
            out_A[:, :latent_dim_sensors],
            out_A[:, latent_dim_sensors:],
        )

        # Run through sensor B
        for t in training_times_B:
            obs_B = true_y_dt_B[:, t, :]
            out_B, h_B = rnn_encoder_B.forward(obs_B, h_B)
        # picking out mean and log variance for B sensor
        qz0_mean_B, qz0_logvar_B = (
            out_B[:, :latent_dim_sensors],
            out_B[:, latent_dim_sensors:],
        )

        qz0_mean_comb = torch.cat((qz0_mean_A, qz0_mean_B), dim=1)
        qz0_logvar_comb = torch.cat((qz0_logvar_A, qz0_logvar_B), dim=1)
        epsilon = torch.randn(qz0_mean_comb.size()).to(device)
        z0 = epsilon * torch.exp(0.5 * qz0_logvar_comb) + qz0_mean_comb

        # Sending z0 to odeint to integrate 'sample' timesteps forward
        # Permute just permutes the tensor.

        pred_z = odeint(latent_node, z0, common_timestamps).permute(1, 0, 2)
        pred_x = decoder(pred_z)

        # We can only correct our predictions against known values
        # so we pick out the known values from pred_x_A and _B
        pred_x_A = pred_x[:, :, 0]
        pred_x_B = pred_x[:, :, 1]
        pred_x_A_known = pred_x_A[:, time_indx_A]
        pred_x_B_known = pred_x_B[:, time_indx_B]

        # Compute loss
        # First create noise as log var then call log_normal_pdf
        # samp_trajs acts as our x, pred_x as our mean i.o.w mu, noise_logvar as our logvar
        # Create logpx for A and B sensors
        noise_std_A = torch.zeros(pred_x_A_known.size()).to(device) + noise_std
        noise_std_B = torch.zeros(pred_x_B_known.size()).to(device) + noise_std
        noise_logvar_A = 2.0 * torch.log(noise_std_A).to(device)
        noise_logvar_B = 2.0 * torch.log(noise_std_B).to(device)

        logpx_A = log_normal_prob_dist(true_y_A, pred_x_A_known, noise_logvar_A).sum(-1)
        logpx_B = log_normal_prob_dist(true_y_B, pred_x_B_known, noise_logvar_B).sum(-1)

        pz0_mean = pz0_logvar = torch.zeros(z0.size()).to(device)
        analytic_kl = kullback_leibler_divergence_normaldist(
            qz0_mean_comb, qz0_logvar_comb, pz0_mean, pz0_logvar
        ).sum(-1)

        # Take the original data and fill in the gaps with predictions
        pred_true_A, pred_true_B = forward_rnn_data(
            pred_x_A, pred_x_B, time_indx_A, time_indx_B, true_y_A, true_y_B
        )
        # We also need to create a dt vector
        pred_true_AB = torch.stack((pred_true_A, pred_true_B, common_dt), dim=2)

        # #We also need to create a guidance of the rnn_encoder_forw and decoder to generalize better
        # h_forward = rnn_encoder_forw.init_hidden_state().to(device)
        # pred_true_AB_samples = pred_true_AB[:, 0 : args.samples, :]
        # max_time = int(np.trunc(pred_true_AB_samples.size(1)))
        # random_stop_time = np.random.randint(low=0, high=max_time, size=1)[0]
        # training_times_forward = range(random_stop_time)

        # for t in training_times_forward:
        #     obs_forward = pred_true_AB_samples[:, t, :]
        #     out_forward, h_forward = rnn_encoder_forw.forward(obs_forward, h_forward)
        # # picking out mean and log variance
        # qz0_mean_forward, qz0_logvar_forward = (
        #     out_forward[:, :latent_dim_total],
        #     out_forward[:, latent_dim_total:],
        # )

        # epsilon = torch.randn(qz0_mean_forward.size()).to(device)
        # z0 = epsilon * torch.exp(0.5 * qz0_logvar_forward) + qz0_mean_forward

        # y0 = decoder(z0)
        # y0_correct = pred_true_AB[:, random_stop_time-1, 0:2]
        # loss_enc_dec = (y0 - y0_correct).sum(-1)

        loss = torch.mean(
            (-logpx_A) + (-logpx_B) + analytic_kl,
            dim=0,
        )
        #+ loss_enc_dec   - Can be added to loss for training encoder decoder
        loss.backward()
        optimizer.step()
        loss_meter.update(loss.item())
        print("Iter: {}, running avg elbo: {:.4f}".format(itr, -loss_meter.avg))

        #
        if args.visualize and itr % args.vis_freq == 0:
            print("vizing")
            with torch.no_grad():

                h_A = rnn_encoder_A.init_hidden_state().to(device)
                h_B = rnn_encoder_B.init_hidden_state().to(device)

                training_times_A = reversed(range(int(np.trunc(true_y_dt_A.size(1)))))
                training_times_B = reversed(range(int(np.trunc(true_y_dt_B.size(1)))))
                # Run through sensor A
                for t in training_times_A:
                    obs_A = true_y_dt_A[:, t, :]
                    out_A, h_A = rnn_encoder_A.forward(obs_A, h_A)
                # picking out mean and log variance for A sensor
                qz0_mean_A, qz0_logvar_A = (
                    out_A[:, :latent_dim_sensors],
                    out_A[:, latent_dim_sensors:],
                )

                # Run through sensor B
                for t in training_times_B:
                    obs_B = true_y_dt_B[:, t, :]
                    out_B, h_B = rnn_encoder_B.forward(obs_B, h_B)
                # picking out mean and log variance for B sensor
                qz0_mean_B, qz0_logvar_B = (
                    out_B[:, :latent_dim_sensors],
                    out_B[:, latent_dim_sensors:],
                )

                qz0_mean_comb = torch.cat((qz0_mean_A, qz0_mean_B), dim=1)
                qz0_logvar_comb = torch.cat((qz0_logvar_A, qz0_logvar_B), dim=1)
                epsilon = torch.randn(qz0_mean_comb.size()).to(device)
                z0 = qz0_mean_comb

                # Saving prediction to file
                pred_z_all = odeint(latent_node, z0, time_cont_A)
                pred_x_all = decoder(pred_z_all)
                print("saving files")
                torch.save(pred_x_all, directory_name + "/pred_x_all"+str(itr)+".pt")
                torch.save(time_cont_A, directory_name + "/time_all"+str(itr)+".pt")
                print("files saved")

                # Sending z0 to odeint to integrate 'sample' timesteps forward
                # Permute just permutes the tensor.
                # First traject for vis.
                z0 = z0[0]
                pred_z = odeint(latent_node, z0, time_cont_A)
                pred_x = decoder(pred_z)

                # Predicting into the future with rnn forward!
                h_forward = rnn_encoder_forw.init_hidden_state().to(device)
                pred_times_forward = range(args.samples-1)
                pred_true_AB_dt_correct = pred_true_AB[:, args.samples : args.samples * 2 - 1, :]
                for t in pred_times_forward:
                    obs_forward = pred_true_AB_dt_correct[:, t, :]
                    out_forward, h_forward = rnn_encoder_forw.forward(obs_forward, h_forward)
                # picking out mean and log variance
                qz0_mean_forward, qz0_logvar_forward = (
                    out_forward[:, :latent_dim_total],
                    out_forward[:, latent_dim_total:],
                )

                epsilon = torch.randn(qz0_mean_forward.size()).to(device)
                #z0 = epsilon * torch.exp(0.5 * qz0_logvar_forward) + qz0_mean_forward
                z0 = qz0_mean_forward
                z0 = z0[0]
                # Creating time vector of timepoints we want to evaluate in the untrained half
                # Picks out the times from the first timepoint after we stopped the RNN forward
                # and then subtracts the time at which we stopped.
                t_forward_pred = torch.linspace(0,args.end_time,steps=100)

                pred_z = odeint(latent_node, z0, t_forward_pred)
                pred_x_forward = decoder(pred_z)

                plt.figure()
                plt.plot(
                    time_cont_A.cpu().detach().numpy(),
                    pred_x.cpu().detach().numpy(),
                    label="Predictions Back-RNN",
                )
                plt.plot(
                    time_cont_A.cpu().detach().numpy(),
                    true_y_cont_A.cpu().detach().numpy(),
                    label="Ground truth A",
                )
                plt.plot(
                    time_cont_B.cpu().detach().numpy(),
                    true_y_cont_B.cpu().detach().numpy(),
                    label="Ground truth B",
                )
                #plt.plot((t_forward_pred+args.end_time).cpu().detach().numpy(), pred_x_forward.cpu().detach().numpy(), '-', label="RNN forward prediction")
                plt.axvspan(0, 20, facecolor="b", alpha=0.5)
                plt.legend()
                plt.draw()
                plt.savefig(
                    directory_name+"/ground_truth_with_pred" + str(itr) + ".png",
                    dpi=500,
                )
                plt.close()

    params = (
        list(latent_node_fwd.parameters())
        + list(rnn_encoder_forw.parameters())
        + list(decoder_fwd.parameters())
    )
    optimizer = optim.Adam(params, lr=args.lr)
    # optimizer = optim.RMSprop(params, lr=1e-3)
    loss_meter = RunningAverageMeter()
    pred_true_AB = pred_true_AB.detach()
    common_timestamps = common_timestamps.detach()
    print("****Beginning forward training****")
    for itr in range(1, args.niters + 1):
        optimizer.zero_grad()
        # Only take half of the samples to run forward rnn over so
        # we have some space for prediction.
        pred_true_AB_samples = pred_true_AB[:, 0 : args.samples, :]
        # *2-1 cause we have a common point in both sensors, t=0
        rnd_stop = int(np.random.rand(1)*(args.samples-10))+10
        pred_true_AB_correct = pred_true_AB[:, rnd_stop : args.samples * 2 - 1, 0:2]
        h_forward = rnn_encoder_forw.init_hidden_state().to(device)
        training_times_forward = range(rnd_stop)

        for t in training_times_forward:
            obs_forward = pred_true_AB_samples[:, t, :]
            out_forward, h_forward = rnn_encoder_forw.forward(obs_forward, h_forward)
        # picking out mean and log variance
        qz0_mean_forward, qz0_logvar_forward = (
            out_forward[:, :latent_dim_total],
            out_forward[:, latent_dim_total:],
        )

        epsilon = torch.randn(qz0_mean_forward.size()).to(device)
        z0 = epsilon * torch.exp(0.5 * qz0_logvar_forward) + qz0_mean_forward
        # Creating time vector of timepoints we want to evaluate in the untrained half
        # Picks out the times from the first timepoint after we stopped the RNN forward
        # and then subtracts the time at which we stopped.
        t_forward_odeint = common_timestamps[rnd_stop : common_timestamps.size(0)]
        t_forward_odeint = t_forward_odeint - common_timestamps[rnd_stop - 1]

        pred_z = odeint(latent_node_fwd, z0, t_forward_odeint).permute(1, 0, 2)
        pred_x = decoder_fwd(pred_z)

        noise_std_forward = torch.zeros(pred_x.size()).to(device) + noise_std
        noise_logvar_forward = 2.0 * torch.log(noise_std_forward).to(device)

        logpx_forward = log_normal_prob_dist(
            pred_true_AB_correct, pred_x, noise_logvar_forward
        ).sum(-1).sum(-1)

        pz0_mean = pz0_logvar = torch.zeros(z0.size()).to(device)
        analytic_kl_forward = kullback_leibler_divergence_normaldist(
            qz0_mean_forward, qz0_logvar_forward, pz0_mean, pz0_logvar
        ).sum(-1)

        loss = torch.mean(
            (-logpx_forward) + analytic_kl_forward,
            dim=0,
        )
        #+ loss_enc_dec   - Can be added to loss for training encoder decoder
        loss.backward()
        optimizer.step()
        loss_meter.update(loss.item())
        print("Iter: {}, running avg elbo: {:.4f}".format(itr, -loss_meter.avg))

        if args.visualize and itr % args.vis_freq == 0:
            print("vizing")
            with torch.no_grad():

                h_A = rnn_encoder_A.init_hidden_state().to(device)
                h_B = rnn_encoder_B.init_hidden_state().to(device)

                training_times_A = reversed(range(int(np.trunc(true_y_dt_A.size(1)))))
                training_times_B = reversed(range(int(np.trunc(true_y_dt_B.size(1)))))
                # Run through sensor A
                for t in training_times_A:
                    obs_A = true_y_dt_A[:, t, :]
                    out_A, h_A = rnn_encoder_A.forward(obs_A, h_A)
                # picking out mean and log variance for A sensor
                qz0_mean_A, qz0_logvar_A = (
                    out_A[:, :latent_dim_sensors],
                    out_A[:, latent_dim_sensors:],
                )

                # Run through sensor B
                for t in training_times_B:
                    obs_B = true_y_dt_B[:, t, :]
                    out_B, h_B = rnn_encoder_B.forward(obs_B, h_B)
                # picking out mean and log variance for B sensor
                qz0_mean_B, qz0_logvar_B = (
                    out_B[:, :latent_dim_sensors],
                    out_B[:, latent_dim_sensors:],
                )

                qz0_mean_comb = torch.cat((qz0_mean_A, qz0_mean_B), dim=1)
                qz0_logvar_comb = torch.cat((qz0_logvar_A, qz0_logvar_B), dim=1)
                epsilon = torch.randn(qz0_mean_comb.size()).to(device)
                z0 = qz0_mean_comb

                # Sending z0 to odeint to integrate 'sample' timesteps forward
                # Permute just permutes the tensor.
                # First traject for vis.
                z0 = z0[0]
                pred_z = odeint(latent_node, z0, time_cont_A)
                pred_x = decoder(pred_z)

                # Predicting into the future with rnn forward!
                h_forward = rnn_encoder_forw.init_hidden_state().to(device)
                pred_times_forward = range(pred_true_AB.size()[1])
                pred_true_AB_dt_correct = pred_true_AB[:, args.samples : args.samples * 2 - 1, :]
                for t in pred_times_forward:
                    obs_forward = pred_true_AB[:, t, :]
                    out_forward, h_forward = rnn_encoder_forw.forward(obs_forward, h_forward)
                # picking out mean and log variance
                qz0_mean_forward, qz0_logvar_forward = (
                    out_forward[:, :latent_dim_total],
                    out_forward[:, latent_dim_total:],
                )

                epsilon = torch.randn(qz0_mean_forward.size()).to(device)
                #z0 = epsilon * torch.exp(0.5 * qz0_logvar_forward) + qz0_mean_forward
                z0 = qz0_mean_forward
                z0 = z0[0]
                # Creating time vector of timepoints we want to evaluate in the untrained half
                # Picks out the times from the first timepoint after we stopped the RNN forward
                # and then subtracts the time at which we stopped.
                t_forward_pred = torch.linspace(0,args.end_time,steps=100)

                pred_z = odeint(latent_node_fwd, z0, t_forward_pred)
                pred_x_forward = decoder_fwd(pred_z)

                plt.figure()
                plt.plot(
                    time_cont_A.cpu().detach().numpy(),
                    pred_x.cpu().detach().numpy(),
                    label="Predictions Back-RNN",
                )
                plt.plot(
                    time_cont_A.cpu().detach().numpy(),
                    true_y_cont_A.cpu().detach().numpy(),
                    label="Ground truth A",
                )
                plt.plot(
                    time_cont_B.cpu().detach().numpy(),
                    true_y_cont_B.cpu().detach().numpy(),
                    label="Ground truth B",
                )
                plt.plot((t_forward_pred+args.end_time).cpu().detach().numpy(), pred_x_forward.cpu().detach().numpy(), '-', label="RNN forward prediction")
                plt.axvspan(0, 20, facecolor="b", alpha=0.5)
                plt.legend()
                plt.draw()
                plt.savefig(
                    directory_name+"/ground_truth_with_pred_forward" + str(itr) + ".png",
                    dpi=500,
                )
                plt.close()


#!/usr/bin/env python
import sys
import pandas
import numpy as np

import matplotlib.pyplot as plt

import sklearn.cluster as cluster
import sklearn.decomposition as decomposition

MANUALLY_SELECTED_ATT = [
    'Entropy_Domain',
    'argPathRatio',
    'ArgUrlRatio',
    'argDomanRatio',
    'pathurlRatio',
    'CharacterContinuityRate',
    'NumberRate_FileName',
    'domainUrlRatio',
    'NumberRate_URL',
    'pathDomainRatio',
    'NumberRate_AfterPath',
    'avgpathtokenlen',
    'URL_Type_obf_Type'
]

#More info in https://scikit-learn.org/stable/auto_examples/cluster/plot_kmeans_digits.html
def kmeans_cluster_graph(csv_file, class_column_name, n_clusters=None, selected_att=None):
    """
    Creates a kmeans cluster graph for the a dataset, using Principal component analysis (PCA) to reduce dimensions

    :param str csv_file: path to the csv file
    :param str class_column_name: the name of the csv column with the class labels
    :param int n_clusters: Number of clusters to use. Defaults to number of unique class labels of the class column
    :param list[str] selected_att: The attributes (column names) from the CSV file to
    use
    """

    # Reads CSV file with data
    dataset = pandas.read_csv(csv_file)

    # Splits CSV file by chosen atributes
    if selected_att:
        df = dataset[ selected_att ]
    else:
        df = dataset

    #Remove valuens with 'NaN'
    df.dropna(inplace=True)

    #Counts unique class labels
    if not n_clusters:
        n_clusters = len(df[class_column_name].unique())
    print(f"Using {n_clusters} clusters")

    #PDA to project the data in two dimensions for the graph
    df_data = df.drop(columns=[class_column_name])
    reduced_data = decomposition.PCA(n_components=2).fit_transform(df_data)

    #Kmeans
    kmeans = cluster.KMeans(init="k-means++", n_clusters=n_clusters)
    kmeans.fit(reduced_data)

    # Step size of the mesh. Decrease to increase the quality of the VQ.
    h = 0.02  # point in the mesh [x_min, x_max]x[y_min, y_max].

    # Plot the decision boundary. For that, we will assign a color to each
    x_min, x_max = reduced_data[:, 0].min() - 1, reduced_data[:, 0].max() + 1
    y_min, y_max = reduced_data[:, 1].min() - 1, reduced_data[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))

    # Obtain labels for each point in mesh. Use last trained model.
    Z = kmeans.predict(np.c_[xx.ravel(), yy.ravel()])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    plt.figure(1)
    plt.clf()
    plt.imshow(
        Z,
        interpolation="nearest",
        extent=(xx.min(), xx.max(), yy.min(), yy.max()),
        cmap=plt.cm.Paired,
        aspect="auto",
        origin="lower",
    )

    plt.plot(reduced_data[:, 0], reduced_data[:, 1], "k.", markersize=2)
    # Plot the centroids as a white X
    centroids = kmeans.cluster_centers_
    plt.scatter(
        centroids[:, 0],
        centroids[:, 1],
        marker="x",
        s=169,
        linewidths=3,
        color="w",
        zorder=10,
    )
    plt.title(
        "K-means clustering on the dataset (PCA-reduced data)\n"
        "Centroids are marked with white cross"
    )
    plt.xlim(x_min, x_max)
    plt.ylim(y_min, y_max)
    plt.xticks(())
    plt.yticks(())
    plt.show()

if __name__ == '__main__':

    if len(sys.argv) < 3:
        exit("Usage: kmeans_cluster_graph.py <csv_file> <class_column_name> (opt) <n_clusters>")

    kmeans_cluster_graph(
        csv_file=sys.argv[1],
        class_column_name=sys.argv[2],
    )


import pathlib

from collections import defaultdict

root_folder = "/home/nvaessen/data/transfer/paper_n4/aam/ablation/"
experiment_name = "speaker_wav2vec2_aam"
tag_prefix = "ablation"
test_run = False
test_sets_to_use = [
    # "voxceleb2",
    "voxceleb2_test_everyone",
    # "voxceleb2_test_hard"
]

postfix_map = {
    "voxceleb2": "o",
    "voxceleb2_test_hard": "h",
    "voxceleb2_test_everyone": "e",
}

path_dict = defaultdict(set)
for ckpt in sorted(pathlib.Path(root_folder).glob("*.ckpt")):
    first_underscore = ckpt.stem.find("_")
    first_dot = ckpt.stem.find(".")

    ablation_name = ckpt.stem[first_underscore+1:first_dot]

    path_dict[ablation_name].add(ckpt.absolute())

assert len(path_dict) == 10
for key, v in path_dict.items():
    assert len(v) == 3

for test_set in test_sets_to_use:
    for ablation_name, ckpt_set in path_dict.items():
        for ckpt in ckpt_set:
            command_template = (
                f"python run.py -m +experiment={experiment_name} "
                f"data/module={test_set} "
                f"fit_model=False "
                f"network.stat_pooling_type=first+cls "
                f"tag={ablation_name}_eval_{postfix_map[test_set]} "
                f"load_network_from_checkpoint={ckpt} "
                f"network.explicit_num_speakers=5994 "
                f"hydra/launcher=slurm "
            )

            print(f"{command_template} & ")

            if test_run:
                exit()


import json
import pickle
import re

from bs4 import BeautifulSoup, SoupStrainer
from six.moves import urllib_parse
from resources.lib.ui import control, database
from resources.lib.ui.jscrypto import jscrypto
from resources.lib.ui.BrowserBase import BrowserBase
from resources.lib.indexers.malsync import MALSYNC

class sources(BrowserBase):
    _BASE_URL = 'https://aniwatch.to/'
    keyurl = 'https://raw.githubusercontent.com/Claudemirovsky/keys/e1/key'
    keyhints = [[106, 7], [109, 6], [137, 7], [147, 7]]

    def get_sources(self, anilist_id, episode, get_backup):
        show = database.get_show(anilist_id)
        kodi_meta = pickle.loads(show.get('kodi_meta'))
        title = kodi_meta.get('name')
        title = self._clean_title(title)
        keyword = title

        all_results = []
        srcs = ['sub', 'dub']
        if control.getSetting('general.source') == 'Sub':
            srcs.remove('dub')
        elif control.getSetting('general.source') == 'Dub':
            srcs.remove('sub')

        items = MALSYNC().get_slugs(anilist_id=anilist_id, site='Zoro')
        if not items:
            if kodi_meta.get('start_date'):
                year = kodi_meta.get('start_date').split('-')[0]
                keyword += ' {0}'.format(year)

            headers = {'Referer': self._BASE_URL}
            params = {'keyword': keyword}
            res = database.get(
                self._get_request,
                8,
                self._BASE_URL + 'search',
                data=params,
                headers=headers
            )

            mlink = SoupStrainer('div', {'class': 'flw-item'})
            mdiv = BeautifulSoup(res, "html.parser", parse_only=mlink)
            sdivs = mdiv.find_all('h3')
            sitems = []
            for sdiv in sdivs:
                try:
                    slug = sdiv.find('a').get('href').split('?')[0]
                    stitle = sdiv.find('a').get('data-jname')
                    sitems.append({'title': stitle, 'slug': slug})
                except AttributeError:
                    pass

            if sitems:
                if title[-1].isdigit():
                    items = [x.get('slug') for x in sitems if title.lower() in x.get('title').lower()]
                else:
                    items = [x.get('slug') for x in sitems if (title.lower() + '  ') in (x.get('title').lower() + '  ')]
                if not items and ':' in title:
                    title = title.split(':')[0]
                    items = [x.get('slug') for x in sitems if (title.lower() + '  ') in (x.get('title').lower() + '  ')]

        if items:
            slug = items[0]
            all_results = self._process_aw(slug, title=title, episode=episode, langs=srcs)

        return all_results

    def _process_aw(self, slug, title, episode, langs):
        sources = []
        headers = {'Referer': self._BASE_URL}
        r = database.get(
            self._get_request,
            8,
            self._BASE_URL + 'ajax/v2/episode/list/' + slug.split('-')[-1],
            headers=headers,
            XHR=True
        )
        res = json.loads(r).get('html')
        elink = SoupStrainer('div', {'class': re.compile('^ss-list')})
        ediv = BeautifulSoup(res, "html.parser", parse_only=elink)
        items = ediv.find_all('a')
        e_id = [x.get('data-id') for x in items if x.get('data-number') == episode]
        if e_id:
            params = {'episodeId': e_id[0]}
            r = database.get(
                self._get_request,
                8,
                self._BASE_URL + 'ajax/v2/episode/servers',
                data=params,
                headers=headers,
                XHR=True
            )
            eres = json.loads(r).get('html')
            for lang in langs:
                elink = SoupStrainer('div', {'class': re.compile('servers-{0}$'.format(lang))})
                sdiv = BeautifulSoup(eres, "html.parser", parse_only=elink)
                srcs = sdiv.find_all('div', {'class': 'item'})
                for src in srcs:
                    edata_id = src.get('data-id')
                    edata_name = src.text.strip().lower()
                    eserver_id = src.get('data-server-id')
                    if edata_name in ['megacloud', 'vidstreaming', 'streamtape']:
                        params = {'id': edata_id}
                        r = self._get_request(
                            self._BASE_URL + 'ajax/v2/episode/sources',
                            data=params,
                            headers=headers,
                            XHR=True
                        )
                        slink = json.loads(r).get('link')
                        if edata_name == 'streamtape':
                            source = {
                                'release_title': '{0} - Ep {1}'.format(title, episode),
                                'hash': slink,
                                'type': 'embed',
                                'quality': 'EQ',
                                'debrid_provider': '',
                                'provider': 'aniwatch',
                                'size': 'NA',
                                'info': ['DUB' if lang == 'dub' else 'SUB', edata_name],
                                'lang': 2 if lang == 'dub' else 0,
                            }
                            sources.append(source)
                        else:
                            headers = {'Referer': slink}
                            sl = urllib_parse.urlparse(slink)
                            spath = sl.path.split('/')
                            spath.insert(2, 'ajax')
                            sid = spath.pop(-1)
                            eurl = '{}://{}{}/getSources'.format(sl.scheme, sl.netloc, '/'.join(spath))
                            params = {'id': sid}
                            res = self._get_request(
                                eurl,
                                data=params,
                                headers=headers,
                                XHR=True
                            )
                            res = json.loads(res)
                            subs = res.get('tracks')
                            if subs:
                                subs = [{'url': x.get('file'), 'lang': x.get('label')} for x in subs if x.get('kind') == 'captions']
                            if res.get('encrypted'):
                                slink = self._process_link(res.get('sources'))
                            else:
                                slink = res.get('sources')[0].get('file')
                            if not slink:
                                continue
                            res = self._get_request(slink, headers=headers)
                            quals = re.findall(r'#EXT.+?RESOLUTION=\d+x(\d+).+\n(?!#)(.+)', res)

                            for qual, qlink in quals:
                                qual = int(qual)
                                if qual < 577:
                                    quality = 'NA'
                                elif qual < 721:
                                    quality = '720p'
                                elif qual < 1081:
                                    quality = '1080p'
                                else:
                                    quality = '4K'

                                source = {
                                    'release_title': '{0} - Ep {1}'.format(title, episode),
                                    'hash': urllib_parse.urljoin(slink, qlink) + '|User-Agent=iPad',
                                    'type': 'direct',
                                    'quality': quality,
                                    'debrid_provider': '',
                                    'provider': 'aniwatch',
                                    'size': 'NA',
                                    'info': ['DUB' if lang == 'dub' else 'SUB', 'Server ' + eserver_id],
                                    'lang': 2 if lang == 'dub' else 0,
                                    'subs': subs
                                }
                                sources.append(source)
        return sources

    def _process_link(self, sources):
        r = database.get(
            self._get_request,
            0.2,
            self.keyurl
        )
        keyhints = json.loads(r) or self.keyhints
        try:
            key = ''
            orig_src = sources
            y = 0
            for m, p in keyhints:
                f = m + y
                x = f + p
                key += orig_src[f:x]
                sources = sources.replace(orig_src[f:x], '')
                y += p
            sources = json.loads(jscrypto.decode(sources, key))
            return sources[0].get('file')
        except:
            database.remove(
                self._get_request,
                self.keyurl
            )
            control.log('decryption key not working')
            return ''


from collections import deque

n, m = map(int, input().split())
arr = [list(map(int, str(input()))) for _ in range(n)]

visited = [[[0] * 2 for _ in range(m)] for _ in range(n)]  # 특정 위치와 벽을 부쉈을 경우, 안부쉈을경우를 나타내기위해 3차원 리스트로 초기화

# visited의 [x]는 행 [y]는 열 [z][0]은 벽을 안부순경우, [z][1] 부순경우 이다

def bfs():
    queue = deque([(0, 0, 0)])
    visited[0][0][0] = 1
    dxy = [(1, 0), (-1, 0), (0, 1), (0, -1)]
    while queue:
        x, y, z = queue.popleft()
        if x == n - 1 and y == m - 1:  # 목표 지점에 도달할 경우 리턴, 일단 먼저 도착하는게 최단경로이므로 이래도 됨
            return visited[x][y][z]
        for dx, dy in dxy:
            nx = x + dx
            ny = y + dy
            if nx < 0 or nx >= n or ny < 0 or ny >= m:
                continue
            if arr[nx][ny] == 0 and visited[nx][ny][z] == 0:  # 만약 벽이 없고 방문 하지 않았다면
                visited[nx][ny][z] = visited[x][y][z] + 1  # 거리를 기존 위치에서 1 늘려서 저장
                queue.append((nx, ny, z))  # 다음 위치를 큐에 넣음
            elif arr[nx][ny] == 1 and z == 0:  # 만약 벽이고 지금까지 벽을 부순적이 없다면
                visited[nx][ny][z + 1] = visited[x][y][z] + 1  # 벽을 부순 경우의 값을 현재 위치 +1로 저장
                queue.append((nx, ny, z + 1))  # 벽을 부순 경우를 큐에 넣음
    return -1

print(bfs())

백준/Gold/2206. 벽 부수고 이동하기/벽 부수고 이동하기.py

벽 부수고 이동하기.py


from Chern_insulator import *
import matplotlib.pyplot as plt
import time
import argparse
import pickle
import numpy as np
from mpi4py.futures import MPIPoolExecutor
from copy import copy
import itertools

def run(p):
    Lx,Ly,m,prob=p
    params=Params(m=m,Lx=Lx,Ly=Ly)
    # prob=list(np.random.permutation([1,0]*(Lx//4*Ly)))
    params.measure_all_Born([np.arange(Lx//4,Lx//2),np.arange(Ly)],prob=prob)
    LN=params.log_neg([np.arange(Lx//4),np.arange(Ly)],[np.arange(Lx//4)+Lx//2,np.arange(Ly)])
    fn=params.fermion_number([np.arange(Lx//4,Lx//2),np.arange(Ly)])
    total_prob=params.snap_prob([np.arange(Lx//4,Lx//2),np.arange(Ly)],occ=np.array(prob)).prod()
    return LN,fn,total_prob

if __name__=="__main__":
    st=time.time()
    parser=argparse.ArgumentParser()
    parser.add_argument('--kmax',default=2,type=int)
    parser.add_argument('--Lx',default=32,type=int)
    parser.add_argument('--Ly',default=4,type=int)
    parser.add_argument('--m',default=1,type=float)

    args=parser.parse_args()
    executor=MPIPoolExecutor()
    kmax=args.kmax
    Lx=args.Lx
    Ly=args.Ly
    m=args.m
    prob=([1,0]*(Lx//4*Ly))
    prob1_list=[[prob,]]
    for k in range(1,kmax):
        prob_k_list=[]
        nchoosek=list(itertools.combinations(np.arange(Lx*Ly//4),k))
        for i1 in (nchoosek):
            for i2 in (nchoosek):
                prob_k=copy(prob)
                prob_k=np.array(prob_k)
                prob_k[2*np.array(list(i1))]=0
                prob_k[2*np.array(list(i2))+1]=1
                prob_k_list.append(list(prob_k))
        prob1_list.append(prob_k_list)

    LN_map=[]
    fn_map=[]
    total_prob_map=[]
    for k,prob_list in enumerate(prob1_list):
        st=time.time()
        inputs=[(Lx,Ly,m,prob) for prob in prob_list]
        LN_list=np.zeros(len(prob_list))
        fn_list=np.zeros(len(prob_list))
        total_prob_list=np.zeros(len(prob_list))
        pool=executor.map(run, inputs)
        # pool=map(run, inputs)
        for es_i,result in enumerate(pool):
            LN,fn,total_prob=result
            LN_list[es_i]=LN
            fn_list[es_i]=fn
            total_prob_list[es_i]=total_prob
        LN_map.append(LN_list)
        fn_map.append(fn_list)
        total_prob_map.append(total_prob_list)

        with open('CI_kmax{:d}_Lx{:d}_Ly{:d}_m{:.1f}.pickle'.format(kmax,Lx,Ly,m),'wb') as f:
            pickle.dump([LN_map,fn_map,total_prob_map],f)

        print('k={:d}, len={:d}, {:.1f}'.format(k,len(prob_list),time.time()-st))


import torchvision
import torch
import csv

from os.path import join
from collections import OrderedDict
from itertools import islice

import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from io import BytesIO

def get_image_pairs():
    test_dir = './CW_data/test/'
    training_dir = './CW_data/training/'
    answers = OrderedDict()
    with open('test-training_map.txt', mode='r') as csv_file:
        csv_reader = csv.DictReader(csv_file)
        for row in csv_reader:
            if 'f' in row['training']:
                continue
            answers[ join(test_dir, row['test']) ] = join(training_dir, row['training'])

    return answers

def do_img(filename):
    image = Image.open(filename)

    image_tensor = torchvision.transforms.functional.to_tensor(image)

    # pass a list of (potentially different sized) tensors
    # to the model, in 0-1 range. The model will take care of
    # batching them together and normalizing
    output = do_img.model([image_tensor])
    # output is a list of dict, containing the postprocessed predictions

    plt.subplots(figsize=(50,100))
    tmp = np.array(image)
    mask = torch.zeros(tmp.shape)
    for keypoints in output[0]["keypoints"]:
        for keypoint in keypoints:
            x,y,_ = keypoint.data
            # print(int(y), int(x))
            for dx in range(-5,6):
                # mask[int(y), int(x),1] = 1
                mask[int(y+dx), int(x+dx),1] = 1
                mask[int(y-dx), int(x+dx),1] = 1
        break
    mask = mask  + tmp / 255  *0.6
    # plt.subplot(2,1,1)
    # plt.imshow(tmp)
    # plt.subplot(2,1,1)
    plt.imshow(mask)
    plt.show()

def main():

    do_img.model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True).eval()
    test_training_mapping = get_image_pairs()

    num_of_pairs = 22
    sliced = islice(test_training_mapping.items(), num_of_pairs)
    trimmed_mapping = OrderedDict(sliced)

    for test,training in trimmed_mapping.items():
        print (test, training)
        do_img(test)
        do_img(training)

if __name__ == "__main__":
    main()


# -*- coding: gb18030 -*-
#
# $Id: ItemResearch.py, fangpengjun Exp $

"""
implement item Research window class

"""
from guis import *
import ShareTexts
from LabelGather import labelGather
from guis.common.Window import Window
from guis.common.PyGUI import PyGUI
from guis.controls.BaseObjectItem import BaseObjectItem as BOItem
from guis.controls.ProgressBar import HFProgressBar
from guis.controls.Button import Button
from guis.controls.StaticText import StaticText
from guis.tooluis.CSRichText import CSRichText
from guis.controls.Icon import Icon
from ModelRender import GodRender
from guis.controls.Icon import Icon
from guis.tooluis.richtext_plugins.PL_Font import PL_Font
from guis.tooluis.richtext_plugins.PL_NewLine import PL_NewLine
from guis.tooluis.richtext_plugins.PL_Space import PL_Space
from guis.tooluis.richtext_plugins.PL_Image import PL_Image
from ItemsFactory import ObjectItem as ItemInfo
from guis.controls.SelectorGroup import SelectorGroup
from TongBeastData import TongBeastData
tongBeastData = TongBeastData.instance()
from NPCModelLoader import NPCModelLoader
modelLoader = NPCModelLoader.instance()
from config.client.msgboxtexts import Datas as mbmsgs
import items
import GUIFacade
import csconst
import csdefine
import random
import Language
from config.client import tongBeastDsp

TONG_ACTIONVAL_LIMIT = {
	0 : ( ShareTexts.NUM_0, 0 ),
	1 : ( ShareTexts.NUM_1, 100 ),
	2 : ( ShareTexts.NUM_2, 200 ),
	3 : ( ShareTexts.NUM_3, 300 ),
	4 : ( ShareTexts.NUM_4, 400 ),
	5 : ( ShareTexts.NUM_5, 500 )
}

BEAS_INDEX_MAP = { 1: "10111139",
		2: "10111140",
		3: "10111141",
		4: "10111142"
		}

class ShenShouBeckon( Window ):

	def __init__( self ):
		wnd = GUI.load( "guis/general/tongabout/shenshoubeckon/window.gui" )
		uiFixer.firstLoadFix( wnd )
		Window.__init__( self, wnd )
		self.posZSegment = ZSegs.L4
		self.activable_ = True
		self.escHide_ 	 = True
		self.selIndex = 0 #当前选择的神兽索引
		self.beckonIndex = 0 #当前召唤出的神兽
		self.__trapID = 0
		self.shenshouLevel = 0 #召唤出的神兽等级
		self.__turnModelCBID = 0
		self.__defaultColor = 255, 255, 255, 255
		self.__triggers = {}
		self.__registerTriggers()
		self.__initialize( wnd )

	def __initialize( self, wnd ):
		self.__pyBeckonBtn = Button( wnd.btnBeckon )
		self.__pyBeckonBtn.setStatesMapping( UIState.MODE_R4C1 )
		self.__pyBeckonBtn.onLClick.bind( self.__beckonShenShou )
		labelGather.setPyBgLabel( self.__pyBeckonBtn, "TongAbout:ShenShouBeckon", "btnBeckon" )

		self.__pyStAction = StaticText( wnd.actionValue )
		self.__pyStAction.text = ""

		self.__pyStAcadeLevel = StaticText( wnd.stAcadLevel )#帮会神兽殿等级
		self.__pyStAcadeLevel.text = ""

		self.__pyRtUseMoney = CSRichText( wnd.rtUseMoney )
		self.__pyRtUseMoney.align = "L"
		self.__pyRtUseMoney.text = ""

		self.__pyStSelGod = StaticText( wnd.selGodAnimal ) #当前召唤出的神兽
		self.__pyStSelGod.text = ""

		self.__pyBtnLeft = Button( wnd.btnLeft )
		self.__pyBtnLeft.setStatesMapping( UIState.MODE_R2C2 )
		self.__pyBtnLeft.onLMouseDown.bind( self.__turnLeft )

		self.__pyBtnRight = Button( wnd.btnRight )
		self.__pyBtnRight.setStatesMapping( UIState.MODE_R2C2 )
		self.__pyBtnRight.onLMouseDown.bind( self.__turnRight )

		self.__pyGodRender = GodRender( wnd.modelRender )

		self.__pyRtResReq = CSRichText( wnd.rtReshReq )
		self.__pyRtResReq.align = "L"
		self.__pyRtResReq.lineFlat = "T"
		self.__pyRtResReq.text = ""

		self.__pyGodBeckon = Icon( wnd.godBeckon.icon )
		self.__pyGodBeckon.icon = ""
		self.__godsGroup = SelectorGroup()
		self.beastDatas = tongBeastData.getDatas()
		self.__pyGodIcons = {}
		self.__pyStAttrs = {}
		self.__pyAttrBar = {}
		for name, item in wnd.children:
			if name.startswith( "godAnimal_" ): #神兽图标
				index = int( name.split( "_" )[1])
				pyGodIcon = GodIcon( item, index )
				self.__pyGodIcons[index] = pyGodIcon
				self.__godsGroup.addSelector( pyGodIcon )
				beaClassName = BEAS_INDEX_MAP[index]
				if self.beastDatas.has_key( beaClassName ): #设置神兽头像
					modelNums = self.beastDatas[beaClassName][1]["modelNums"]
					if len( modelNums ) <= 0:return
					modelNumber = modelNums[random.randint(0, len( modelNums ) - 1 )]
					pyGodIcon.godTexture = modelLoader.getHeadTexture( modelNumber )
					pyGodIcon.modelNumber = modelNumber
			if name.startswith( "attr_" ): #神兽战斗属性
				attrStr = name[5:]
				pyStattr = AttrItem( item )
				pyStattr.title = labelGather.getText( "TongAbout:ShenShouBeckon", attrStr )
				pyStattr.text = ""
				self.__pyStAttrs[attrStr] = pyStattr
			if name.endswith( "_bar" ): #神兽血条和蓝条
				attrbar = name.split( "_" )[0]
				pyAttrBar = HFProgressBar( item.bar )
				pyAttrBar.clipMode = "RIGHT"
				pyAttrBar.value = 0.0
				self.__pyAttrBar[attrbar] = pyAttrBar

		labelGather.setLabel( wnd.lbTitle, "TongAbout:ShenShouBeckon", "lbTitle" )
		labelGather.setLabel( wnd.useMoneyText, "TongAbout:BuildReSearch", "useMoneyText" )
		labelGather.setLabel( wnd.actionText, "TongAbout:BuildReSearch", "actionText" )
		labelGather.setLabel( wnd.rendPanel.stTitle, "TongAbout:BuildReSearch", "resList" )
		labelGather.setLabel( wnd.reqPanel.bgTitle.stTitle, "TongAbout:ShenShouBeckon", "beckonReq" )
		labelGather.setLabel( wnd.statusPanel.bgTitle.stTitle, "TongAbout:ShenShouBeckon", "tongStatus" )
	# --------------------------------------------------------
	def __registerTriggers( self ) :
		self.__triggers["EVT_ON_TOGGLE_TONG_SHENSHOU_BRECKON"] = self.__toggleShows #弹出窗口
		self.__triggers["EVT_ON_TOGGLE_TONG_SET_SHENSHOU_INFO"] = self.__onCurShenshouChange #当前选择神兽改变通知
		self.__triggers["EVT_ON_TOGGLE_TONG_MONEY_CHANGE"] = self.__onTongMoneyChange #帮会资金改变
		self.__triggers["EVT_ON_TONG_SHENSHOU_SLECTED"] = self.__onSelectedShenshou #选择某一个神兽
		self.__triggers["EVT_ON_TONG_SHENSHOU_BRECKON_SUCCEED"] = self.__onBreckonSucc #神兽召唤成功后的回调
		self.__triggers["EVT_ON_TOGGLE_TONG_BUILD_INFO"] = self.__onReciveBuildInfo	#获取建筑信息
		for key in self.__triggers.iterkeys() :
			ECenter.registerEvent( key, self )

	def __deregisterTriggers( self ) :
		for key in self.__triggers.iterkeys() :
			ECenter.unregisterEvent( key, self )

	# -------------------------------------------------------
	def __toggleShows( self, shenshouLevel, shenshouType ):
		player = BigWorld.player()
		distance = csconst.COMMUNICATE_DISTANCE
		if hasattr( GUIFacade.getGossipTarget(), "getRoleAndNpcSpeakDistance" ):
			distance = GUIFacade.getGossipTarget().getRoleAndNpcSpeakDistance()
		self.__trapID = player.addTrapExt( distance, self.__onEntitiesTrapThrough )#打开窗口后为玩家添加对话陷阱
		self.__pyStAcadeLevel.text = labelGather.getText( "TongAbout:ShenShouBeckon", "ssLevel" )%shenshouLevel
		beaClassName = ""
		if shenshouType > 0: #没有召唤出来的神兽
			beaClassName = BEAS_INDEX_MAP[shenshouType]
		self.beckonIndex = shenshouType
		self.shenshouLevel = shenshouLevel
		if self.beastDatas.has_key( beaClassName ): #设置召唤出的神兽头像
			modelNums = self.beastDatas[beaClassName][1]["modelNums"]
			modelNum = ""
			if len( modelNums ) > 0:
				modelNum = modelNums[random.randint(0, len( modelNums ) - 1 )]
			self.__pyGodBeckon.icon = modelLoader.getHeadTexture( modelNum )
			self.__pyGodRender.resetModel( beaClassName )
			for attr, pyStattr in self.__pyStAttrs.iteritems():
				pyStattr.text = self.beastDatas[beaClassName][shenshouLevel][attr]
			for attrbar, pyAttrBar in self.__pyAttrBar.items():
				attrVal = self.beastDatas[beaClassName][shenshouLevel][attrbar]
				if attrVal <= 0:
					pyAttrBar.value = 0.0
				else:
					pyAttrBar.value = float( attrVal/attrVal )
			self.__pyStSelGod.text = self.beastDatas[beaClassName][shenshouLevel]["uname"]
		self.show()

	def __delTrap( self ) :
		player = BigWorld.player()
		if self.__trapID :
			player.delTrap( self.__trapID )									#删除玩家的对话陷阱
			self.__trapID = 0

	def __onEntitiesTrapThrough( self, entitiesInTrap ):
		gossiptarget = GUIFacade.getGossipTarget()							#获取当前对话NPC
		if gossiptarget and gossiptarget not in entitiesInTrap:				#如果NPC离开玩家对话陷阱
			self.__delTrap()
			self.hide()														#隐藏当前交易窗口

	def __onCurShenshouChange( self, shenshouLevel, shenshouType ):
		"""
		当前神兽改变通知
		"""
		self.__pyStAcadeLevel.text = labelGather.getText( "TongAbout:ShenShouBeckon", "ssLevel" )%shenshouLevel
		beaClassName = ""
		if shenshouType > 0: #没有召唤出来的神兽
			beaClassName = BEAS_INDEX_MAP[shenshouType]
		self.beckonIndex = shenshouType
		self.shenshouLevel = shenshouLevel
		if self.beastDatas.has_key( beaClassName ): #设置召唤出的神兽头像
			modelNums = self.beastDatas[beaClassName][1]["modelNums"]
			modelNum = ""
			if len( modelNums ) > 0:
				modelNum = modelNums[random.randint(0, len( modelNums ) - 1 )]
			self.__pyGodBeckon.icon = modelLoader.getHeadTexture( modelNum )
			self.__pyGodRender.resetModel( beaClassName )
			for attr, pyStattr in self.__pyStAttrs.iteritems():
				pyStattr.text = self.beastDatas[beaClassName][shenshouLevel][attr]
			for attrbar, pyAttrBar in self.__pyAttrBar.items():
				attrVal = self.beastDatas[beaClassName][shenshouLevel][attrbar]
				if attrVal <= 0:
					pyAttrBar.value = 0.0
				else:
					pyAttrBar.value = float( attrVal/attrVal )
			self.__pyStSelGod.text = self.beastDatas[beaClassName][shenshouLevel]["uname"]

	def __onSelectedShenshou( self, selIndex ):
		self.selIndex = selIndex
		for index, pyGodIcon in self.__pyGodIcons.items():
			pyGodIcon.selected = index == selIndex
		if self.__pyGodIcons.has_key( selIndex ):
			selPyIcon = self.__pyGodIcons[selIndex]
			modelNumber = selPyIcon.modelNumber
			self.__pyGodRender.resetModel( BEAS_INDEX_MAP[selIndex] )
		if BEAS_INDEX_MAP.has_key( selIndex ):
			className = BEAS_INDEX_MAP[selIndex]
			for attr, pyStattr in self.__pyStAttrs.iteritems():
				pyStattr.text = self.beastDatas[className][self.shenshouLevel][attr]
			for attrbar, pyAttrBar in self.__pyAttrBar.items():
				attrVal = self.beastDatas[className][self.shenshouLevel][attrbar]
				if attrVal <= 0:
					pyAttrBar.value = 0.0
				else:
					pyAttrBar.value = float( attrVal/attrVal )
		if BigWorld.player().tong_getCanUseMoney() >= 500000:
			reqmoney_color = ( 230, 227, 185 )
		else:
			reqmoney_color = ( 255, 0, 0 )
		priceText = utils.currencyToViewText( 500000 )
		req_Expense = labelGather.getText( "TongAbout:ItemReSearch", "consumMoney" )%( PL_Space.getSource( 4 ), priceText )
		req_Expense = PL_Font.getSource( req_Expense, fc = ( 230, 227, 185, 255 ) )
		req_Expense += PL_Font.getSource( fc = self.__defaultColor )
		self.__pyRtResReq.text = req_Expense
		self.__pyBeckonBtn.enable = selIndex != self.beckonIndex

	def __onBreckonSucc( self, index ):
		pass

	def __onTongMoneyChange( self, money ):
		"""
		帮会资金发生变化
		"""
		if not BigWorld.player().inWorld or not self.visible:
			return
		self.__showTongCanUseMoney()

	def __onReciveBuildInfo( self, buildData ):
		if buildData["type"] == csdefine.TONG_BUILDING_TYPE_JK: #金库信息
			self.__showTongCanUseMoney()

	def __showTongCanUseMoney( self ):
		"""
		显示帮会可用资金
		"""
		money = BigWorld.player().tong_getCanUseMoney()
		self.__pyRtUseMoney.text = utils.currencyToViewText( money )

	def __turnLeft( self ):
		BigWorld.cancelCallback( self.__turnModelCBID )
		self.__turnModel( False )
		LastKeyUpEvent.attach( self.__onLastKeyUpEvent )
		return True

	def __turnRight( self ):
		BigWorld.cancelCallback( self.__turnModelCBID )
		self.__turnModel( True )
		LastKeyUpEvent.attach( self.__onLastKeyUpEvent )
		return True

	def __turnModel( self, isRTurn ) :
		"""
		turning model on the mirror
		"""
		self.__pyGodRender.yaw += ( isRTurn and -0.1 or 0.1 )
		if BigWorld.isKeyDown( KEY_LEFTMOUSE ) :
			self.__turnModelCBID = BigWorld.callback( 0.1, Functor( self.__turnModel, isRTurn ) )

	def __onLastKeyUpEvent( self, key, mods ) :
		if key != KEY_LEFTMOUSE : return
		BigWorld.cancelCallback( self.__turnModelCBID )
		LastKeyUpEvent.detach( self.__onLastKeyUpEvent )

	def __beckonShenShou( self ):
		player = BigWorld.player()
		if self.selIndex == 0: return #第一个神兽为出战
		if self.selIndex == self.beckonIndex:
			# "该神兽已经存在!"
			showAutoHideMessage( 3.0, 0x09a1, mbmsgs[0x0c22] )
			return
		elif not player.tong_checkDutyRights( player.tong_grade, csdefine.TONG_RIGHT_PET_SELECT ):
			# "你没有选择神兽的权限!"
			showAutoHideMessage( 3.0, 0x09a2, mbmsgs[0x0c22] )
			return
		elif player.tong_getCanUseMoney() < 500000:
			# "帮会资金不足,选择或者更换神兽需要%d金帮会资金。"
			showAutoHideMessage( 3.0, mbmsgs[0x09a3] % 50, mbmsgs[0x0c22] )
			return

		player.tong_shenShouSelect( self.selIndex )

	# -----------------------------------------------------------
	# public
	# -----------------------------------------------------------
	def onEvent( self, macroName, *args ) :
		self.__triggers[macroName]( *args )

	def show( self ):
		self.__showTongCanUseMoney()
		self.__onSelectedShenshou( 0 )
		self.__pyGodRender.enableDrawModel()
		Window.show( self )

	def hide( self ):
		Window.hide( self )
		GUIFacade.cancelTurnCB( GUIFacade.getGossipTarget() )

	def onLeaveWorld( self ) :
		self.selIndex = 0
		self.beckonIndex = 0
		self.__pyGodRender.clearModel()
		self.__pyGodRender.disableDrawModel()
		for pyGodIcon in self.__pyGodIcons.itervalues():
			pyGodIcon.selected = False
		self.hide()

# -----------------------------------------------
from guis.common.PyGUI import PyGUI
from guis.controls.SelectableButton import SelectableButton

class GodIcon( SelectableButton ):
	def __init__( self, godIcon, index ):
		SelectableButton.__init__( self, godIcon )
		self.__pyIconBg = PyGUI( godIcon.iconBg )
		self.__pyGodIcon = Icon( godIcon.icon )
		self.__pyGodIcon.focus = True
		self.__pyGodIcon.crossFocus = True
		self.__pyGodIcon.onLClick.bind( self.__onGodSelected )
		self.__pyGodIcon.onMouseEnter.bind( self.__onGodEnter )
		self.__pyGodIcon.onMouseLeave.bind( self.__onGodLeave )
		if hasattr( godIcon,"cover" ):
			self.__pyCover = PyGUI( godIcon.cover )
			self.__pyCover.visible = False
		self.__selected = False
		self.index = index
		self.__modelNumber = ""

	def __onGodSelected( self, pyGodIcon ):
		if pyGodIcon is None:return
		ECenter.fireEvent( "EVT_ON_TONG_SHENSHOU_SLECTED", self.index )
		toolbox.infoTip.hide()

	def __onGodEnter( self, pyIcon ): #显示神兽信息
		beaDsp = self.getTongBeaDsp( self.index )
		toolbox.infoTip.showItemTips( self, beaDsp )

	def __onGodLeave( self, pyIcon ):
		toolbox.infoTip.hide()

	def getTongBeaDsp( self, selIndex ): #读取帮会建筑描述
		return tongBeastDsp.Datas[ selIndex ]
	# ----------------------------------------------------
	def __select( self ):
		if self.__pyCover:
			self.__pyCover.visible = True

	def __deselect( self ):
		if self.__pyCover:
			self.__pyCover.visible = False

	def _getSelected( self ):
		return self.__selected

	def _setSelected( self, selected ):
#		self.__pyReItem.selected = selected
		if selected:
			self.__select()
		else:
			self.__deselect()
		self.__selected = selected

	def _getGodTexture( self ):
		return self.__pyGodIcon.icon

	def _setGodTexure( self, texture ):
		self.__pyGodIcon.icon = texture

	def _getModelNumber( self ):
		return self.__modelNumber

	def _setModelNumber( self, modelNumber ):
		self.__modelNumber = modelNumber

	selected = property( _getSelected, _setSelected )
	godTexture = property( _getGodTexture, _setGodTexure )
	modelNumber = property( _getModelNumber, _setModelNumber )

class AttrItem( PyGUI ):
	def __init__( self, item ):
		PyGUI.__init__( self, item )
		self.__pyAttrText = StaticText( item.attrText )
		self.__pyStAttr = StaticText( item.attrValue )
		self.__pyStAttr.color = 255, 255, 255
		self.__pyStAttr.text = ""

	def _getText( self ):
		return self.__pyStAttr.text

	def _setText( self, text ):
		self.__pyStAttr.text = text

	def _getTitle( self ):
		return self.__pyAttrText.text

	def _setTitle( self, title ):
		self.__pyAttrText.text = title

	text = property( _getText, _setText )
	title = property( _getTitle, _setTitle )


import django.db.models.deletion
from django.db import migrations, models

class Migration(migrations.Migration):

    dependencies = [
        ('cms', '__first__'),
    ]

    operations = [
        migrations.CreateModel(
            name='GoogleMap',
            fields=[
                ('cmsplugin_ptr', models.OneToOneField(serialize=False, parent_link=True, auto_created=True, to='cms.CMSPlugin', primary_key=True, on_delete=django.db.models.deletion.CASCADE)),
                ('title', models.CharField(verbose_name='map title', blank=True, null=True, max_length=100)),
                ('address', models.CharField(verbose_name='address', max_length=150)),
                ('zipcode', models.CharField(verbose_name='zip code', max_length=30)),
                ('city', models.CharField(verbose_name='city', max_length=100)),
                ('content', models.CharField(help_text='Displayed under address in the bubble.', blank=True, max_length=255, verbose_name='additional content')),
                ('zoom', models.PositiveSmallIntegerField(verbose_name='zoom level', default=13, choices=[(0, '0'), (1, '1'), (2, '2'), (3, '3'), (4, '4'), (5, '5'), (6, '6'), (7, '7'), (8, '8'), (9, '9'), (10, '10'), (11, '11'), (12, '12'), (13, '13'), (14, '14'), (15, '15'), (16, '16'), (17, '17'), (18, '18'), (19, '19'), (20, '20'), (21, '21')])),
                ('lat', models.DecimalField(help_text='Use latitude & longitude to fine tune the map position.', blank=True, max_digits=10, verbose_name='latitude', null=True, decimal_places=6)),
                ('lng', models.DecimalField(max_digits=10, verbose_name='longitude', blank=True, null=True, decimal_places=6)),
                ('route_planer_title', models.CharField(verbose_name='route planer title', blank=True, null=True, max_length=150, default='Calculate your fastest way to here')),
                ('route_planer', models.BooleanField(verbose_name='route planer', default=False)),
                ('width', models.CharField(help_text='Plugin width (in pixels or percent).', default='100%', max_length=6, verbose_name='width')),
                ('height', models.CharField(help_text='Plugin height (in pixels).', default='400px', max_length=6, verbose_name='height')),
                ('info_window', models.BooleanField(help_text='Show textbox over marker', default=True, verbose_name='info window')),
                ('scrollwheel', models.BooleanField(help_text='Enable scrollwheel zooming on the map', default=True, verbose_name='scrollwheel')),
                ('double_click_zoom', models.BooleanField(verbose_name='double click zoom', default=True)),
                ('draggable', models.BooleanField(verbose_name='draggable', default=True)),
                ('keyboard_shortcuts', models.BooleanField(verbose_name='keyboard shortcuts', default=True)),
                ('pan_control', models.BooleanField(verbose_name='Pan control', default=True)),
                ('zoom_control', models.BooleanField(verbose_name='zoom control', default=True)),
                ('street_view_control', models.BooleanField(verbose_name='Street View control', default=True)),
            ],
            options={
                'abstract': False,
            },
            bases=('cms.cmsplugin',),
        ),
    ]


import cv2
import mediapipe as mp
import time

class FaceDetector:
	def __init__(self, min_detection_con=0.5):
		self.min_detection_con = min_detection_con
		self.mpFace = mp.solutions.face_detection
		self.face = self.mpFace.FaceDetection(min_detection_con)
		self.mpDraw = mp.solutions.drawing_utils

	def find_faces(self, img, draw=True):
		imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
		results = self.face.process(imgRGB)
		bboxs = []
		self.detections = results.detections
		if self.detections:
			for id, detection in enumerate(self.detections):
				bboxC = detection.location_data.relative_bounding_box
				h, w, c = img.shape
				bbox = (int(bboxC.xmin * w), int(bboxC.ymin * h),
					int(bboxC.width * w), int(bboxC.height * h))
				bboxs.append([id, bbox, detection.score])
				if draw:
					img = self.fancy_draw(img, bbox)
					cv2.putText(img, f'{int(detection.score[0]*100)}%',
						(bbox[0],bbox[1]-20),
						cv2.FONT_HERSHEY_PLAIN,
						2, (255, 0, 255), 2)
		return img, bboxs

	def fancy_draw(self, img, bbox, length=30, thick=5, rect_thick=1):
		x, y, w, h = bbox
		x1, y1 = x + w, y + h
		# Top left
		cv2.line(img, (x,y), (x+length,y), (255,0,255), thick)
		cv2.line(img, (x,y), (x,y+length), (255,0,255), thick)
		# Top right
		cv2.line(img, (x1,y), (x1-length,y), (255,0,255), thick)
		cv2.line(img, (x1,y), (x1,y+length), (255,0,255), thick)
		# Bottom left
		cv2.line(img, (x,y1), (x+length,y1), (255,0,255), thick)
		cv2.line(img, (x,y1), (x,y1-length), (255,0,255), thick)
		# Bottom right
		cv2.line(img, (x1,y1), (x1-length,y1), (255,0,255), thick)
		cv2.line(img, (x1,y1), (x1,y1-length), (255,0,255), thick)
		# Rectangle
		cv2.rectangle(img, bbox, (255,0,255), rect_thick)
		return img

def main():
	cap = cv2.VideoCapture(0) # 1 ?
	detector = FaceDetector()
	pTime = 0
	cTime = 0
	while 1:
		success, img = cap.read()
		img, bboxs = detector.find_faces(img)
		print(bboxs)
		cTime = time.time()
		fps = 1/(cTime - pTime)
		pTime = cTime
		cv2.putText(img, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3, (0, 255, 0), 3)
		cv2.imshow("Image", img)
		cv2.waitKey(1)

if __name__ == '__main__':
	main()