text:
#Scrape Trustee details
#Diarmuid McDonnell
#18/10/18
#This file scrapes trustee information from the Charities Regulator website.

################################# Import packages #################################
from urllib.request import urlopen as uReq
from bs4 import BeautifulSoup as soup
from lxml.html import fromstring
from itertools import cycle
from datetime import datetime
from downloaddate_function import downloaddate
from time import sleep
from dropbox.files import WriteMode
from selenium import webdriver
import requests
import json
import dropbox
import pandas as pd
import csv
import os
import io
import numpy as np

# Define a function for generating proxies
def get_proxies():
    url = 'https://free-proxy-list.net/'
    response = requests.get(url)
    parser = fromstring(response.text)
    proxies = set()
    for i in parser.xpath('//tbody/tr')[:10]:
        if i.xpath('.//td[7][contains(text(),"yes")]'):
            #Grabbing IP and corresponding PORT
            proxy = ":".join([i.xpath('.//td[1]/text()')[0], i.xpath('.//td[2]/text()')[0]])
            proxies.add(proxy)
    return proxies

# Fetch Dropbox authentication
dbtokenpath = 'C:/Users/mcdonndz-local/Desktop/admin/db_token.txt'
#dbtokenpath_pi = '/home/pi/admin/dp_token.txt'
dbtokenfile = open(dbtokenpath, "r")
dbapitoken = dbtokenfile.read()
print(dbapitoken)

dbx = dropbox.Dropbox(dbapitoken) # Create an object for accessing Dropbox API

# Get the current date
ddate = downloaddate()

# Define paths
localdatapath = 'C:/Users/mcdonndz-local/Desktop/data/temp/'
remotedatapath = '/trustee_data/data_raw/' # Dropbox path
remotelogpath = '/trustee_data/logs/' # Dropbox path
projpath = './' # Location of syntax

inputfile = localdatapath + 'ire_charityregister_' + ddate + '.xlsx'
outputfile = localdatapath + 'ire_trustee_data_' + ddate + '.json'
logoutputfile = localdatapath + 'ire_trustee_log_' + ddate + '.csv'

# Delete output file if already exists
try:
    os.remove(outputfile)
except OSError:
    pass

# Download latest copy of IRE charity register from the Regulator's website
if not os.path.isfile(inputfile):
	import ire_download
	print('Finished executing ire_download.py')
	print('                                             ')
	print('---------------------------------------------')
	print('                                             ')
	sleep(10)
else:
	print('Already downloaded charity register extract')

# Create a panda's dataframe from the CSV
with open(inputfile, 'rb') as f:
	df = pd.read_excel(f, sheet_name='Public Register')
	print(df.dtypes)

df.reset_index(inplace=True)
df.set_index(['Registered Charity Number'], inplace=True)
regid_list = df.index.values.tolist()
#print(regid_list)

# Define variable names for the output files
logvarnames = ['timestamp', 'regid', 'entid', 'success', 'attempts', 'status code', 'execution time']

with open(logoutputfile, 'a') as f:
    writer = csv.writer(f, logvarnames)
    writer.writerow(logvarnames)

# Define a counter to track how many rows of the input file the script processes
counter = 1

# Scrape proxies
proxies = get_proxies()
print(proxies)
proxy_pool = cycle(proxies)

print(' ') # Whitespace used to make the output window more readable
print('>>> Run started') # Header of the output, with the start time.
print('\r')

trustee_list = [] # Create a list for storing the results of the scrape

# Loop through list of charity numbers and scrape info from webpages
for regid in regid_list[0:10]:

	starttime = datetime.now() # Track how long it takes to scrape data for each charity

	# Define counter to track number of times the script tries to use a valid proxy
	proxytry = 1
	while proxytry < 1000:
		try:
			proxy = next(proxy_pool) # Grab a proxy from the pool
			proxytry = 2000
			charurl = 'https://www.charitiesregulator.ie/umbraco/api/search/getcharitybyregid?regid=' + str(regid)
			headers = {'http': proxy, 'https': proxy, 'User-Agent': 'Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/68.0.3440.106 Safari/537.36'} # Spoof the user-agent of the request to return the content seen by Chrome.

			# Define counter to track number of webpage request attempts; try three times
			attempt = 1

			while attempt < 4:

				rorg = requests.get(charurl, headers=headers) # Grab the page using the URL and headers.
				#print(proxy) # Checks if script is cycling through proxies

				if rorg.status_code==200: # Only proceed if the webpage can be requested successfully
					chardata = rorg.json()
					entid = chardata['data'][0]['entityId']

					# Request page using entid
					trusturl = 'https://www.charitiesregulator.ie/umbraco/api/search/getentitydetailsbyid?id=' + str(entid)
					rtrust = requests.get(trusturl, headers=headers) # Grab the page using the URL and headers.

					charinfo = rtrust.json()
					trustees = charinfo['data']['trusteeDetails']
					print(trustees)

					#fye = charinfo['data']['annualReturnBarcodes'] # Get financial year end date the trustee information pertains to
					#fye = fye[0:10]
					#print(fye)
					'''
						Might have to get this information from the Annual Returns sheet in the Charity Register.
					'''
					if len(trustees) == 0:
						chardetails = {'name': '.', "position": ".", "startDate": ".",'regid': str(regid), 'entityid': str(entid)}
						trustees.append(chardetails)
					else:
						for el in trustees: # Add charity details to the list of trustee data
							el['regid'] = str(regid)
							el['entityid'] = str(entid)

					trustee_list.append(trustees) # Append trustee data to trustee_list

					success = 1
					runtime = datetime.now() - starttime
					attempt = 5
				else:
					print('\r')
					print(rorg.status_code, '| Could not resolve address of webpage')
					print('Will try to request webpage a couple more times')
					attempt +=1
					success = 0

				# Write the results of the scrape search to the log file
				runtime = datetime.now() - starttime
				with open(logoutputfile, 'a', newline='') as f:
					writer = csv.writer(f)
					writer.writerow([datetime.today().strftime('%Y%m%d %H:%M'), str(regid), str(entid), success, attempt, rorg.status_code, runtime])

		except Exception as e:
			print(e) # Could not cycle through the list of proxies
			sleep(1)
			# Scrape proxies again
			proxies = get_proxies()
			print(proxies)
			proxy_pool = cycle(proxies)
			proxytry +=1

# Export the results to a json file
with open(outputfile, 'w') as f:
    json.dump(trustee_list, f)

# Upload files to Dropbox
infile = outputfile
outfile = remotedatapath + 'ire_trustee_data_' + ddate + '.json' # Dataset
loginfile = logoutputfile
logoutfile = remotelogpath + 'ire_trustee_log_' + ddate + '.csv' # Log file

with open(infile, 'rb') as f:
	#print(f.read())
	dbx.files_upload(f.read(), outfile, dropbox.files.WriteMode.overwrite, mute=True)

with open(loginfile, 'rb') as f:
	dbx.files_upload(f.read(), logoutfile, dropbox.files.WriteMode.overwrite, mute=True)

print('\r')
print('>>> Finished')

max_stars_repo_name: DiarmuidM/charity-dissolution

max_issues_repo_name: DiarmuidM/charity-dissolution

max_forks_repo_name: DiarmuidM/charity-dissolution


text:
###############################################################################
### PIECHART SHOWING DISTRIBUTION OF FLAGS
###############################################################################

### Description of the KPI functions:
# - 'flag_piechart' creates a pie chart showing flag distribution for a
# single parameter (sensor).

#------------------------------------------------------------------------------
### Import packages
import matplotlib.pyplot as plt
import os
import numpy as np

#------------------------------------------------------------------------------
### Declare constants etc.

# Figure sizes
PIE_FIG_SIZE = 2.5

# Specify the color representing each QC flag (and missing value)
COLOR_DICT = {'2':'#85C0F9','3':'#A95AA1','4':'#F5793A', 'nan':'grey'}

#------------------------------------------------------------------------------
### Functions

# Create a piechart showing the distribution of flags for a given sensor
def flag_piechart(parameter, vocab, df, output_dir):

	# Store the parameter QC flags in a list (remove '.0')
	col_header_name = vocab[parameter]['col_header_name']
	flag_list = list(df[col_header_name + ' QC Flag'])
	flag_list_cleaned = [str(item).split('.')[0] for item in flag_list]

	# Get the unique flags and their frequencies (as arrays)
	unique_flags, freq = np.unique(flag_list_cleaned, return_counts=True)

	# Create color list according to the color dictionary
	color_list = []
	for flag in list(unique_flags):
		if flag == 'nan':
			color_list.append('grey')
		else:
			color_list.append(COLOR_DICT[flag])

	# Create pie chart
	plt.subplots(figsize=(PIE_FIG_SIZE,PIE_FIG_SIZE))
	patches, text, autotexts = plt.pie(freq, colors=color_list, startangle=90,
		autopct='%1.0f%%')
	for autotext in autotexts:
		autotext.set_color('white')
	plt.legend(patches, list(unique_flags), loc="best")
	plt.axis('equal')
	plt.tight_layout()

	# Save figure to file and close figure
	filename = parameter + '_flag_piechart.png'
	filepath = os.path.join(output_dir, filename)
	plt.savefig(filepath, bbox_inches='tight')
	plt.close()

ext: py

max_stars_repo_name: BjerknesClimateDataCentre/KPI

max_issues_repo_name: BjerknesClimateDataCentre/KPI

reason: import numpy


text:
from pylab import plot, show, legend
import numpy
data = numpy.loadtxt("error_est.gp")
x = data[:, 0]
y = data[:, 1]
plot(x, y, label="max FTR error")
data = numpy.loadtxt("error_exact.gp")
x = data[:, 0]
y = data[:, 1]
plot(x, y, label="exact error")
legend()
show()


text:
import sys, cv2, bisect
import numpy as np

#===============================================================================
class Graph():
    def __init__(self, img=None):
        self.seed = None
        self.node = {}
        self.img = img
        self.visited = set()

    def set_img(self, img):
        self.img = img

    def set_arch(self, node_A, node_B):
        if node_A.has_out_arch():
            node = node_A.get_out_arch()
            node.remove_in_arch(node_A)
        node_A.set_out_arch(node_B)
        node_B.set_in_arch(node_A)
        self.node[node_A.pixel] = node_A
        self.node[node_B.pixel] = node_B
        #self._mark_node(node_B.pixel)

    def has_arch(self, pixel_A, pixel_B):
        node_B = self.node[pixel_B]
        if self.node[pixel_A].has_out_arch(node_B):
            return True
        return False

    def add_node(self, node):
        self.node[node.pixel] = node

    def get_node(self, pixel):
        return self.node[pixel]

    def has_node(self, pixel):
        if pixel in self.node:
            return True
        return False

    def add_visited(self, pixel):
        self.visited.add(pixel)

    def is_visited(self, pixel):
        if pixel in self.visited:
            return True
        return False

    def has_8_neighbor(self, pixel):
        adj_list = [(-1,-1), (-1,0), (-1,1), (0,1), (1,1), (1,0), (1,-1), (0,-1)]
        for i in adj_list:
            p = (pixel[0] + i[0], pixel[1] + i[1])
            if p in self.node:
                return self.node[p]
        return False

    def has_4_neighbor(self, pixel):
        adj_list = [(-1,0), (0,1), (1,0), (0,-1)]
        for i in adj_list:
            p = (pixel[0] + i[0], pixel[1] + i[1])
            if p in self.node:
                return self.node[p]
        return False

    def _mark_node(self, pixel):
        self.img[pixel] = (0,100,0)

#===============================================================================
class Node():
    def __init__(self, pixel, cost=sys.maxsize):
        self.in_arch = set()
        self.out_arch = None
        self.pixel = pixel
        self.cost = cost

    def has_in_arch(self):
        if self.in_arch:
            return True
        return False

    def has_out_arch(self):
        if self.out_arch:
            return True
        return False

    def get_in_arch(self):
        if len(self.in_arch) > 0:
            return list(self.in_arch)[0]

    def get_out_arch(self):
        return self.out_arch

    def get_in_degree(self):
        return len(self.in_arch)

    def set_in_arch(self, node):
        self.in_arch.add(node)

    def set_out_arch(self, node):
        self.out_arch = node

    def remove_in_arch(self, node):
        self.in_arch.discard(node)

    def y(self):
        return self.pixel[0]

    def x(self):
        return self.pixel[1]

    def __lt__(self, other):
        self.cost < other.cost

#===============================================================================

class PriorityQueue:
    def __init__(self):
        self.queue = []
        self.nodes = {}

    def put(self, item, priority):
        if item in self.nodes:
            pos = bisect.bisect_right(self.queue, [self.nodes[item], item])
            del self.queue[pos-1]
        bisect.insort_right(self.queue, [priority, item])
        self.nodes[item] = priority

    def pop(self):
        if self.queue:
            item = self.queue.pop(0)[1]
            #print("queue:", len(self.queue))
            #print("nodes:", len(self.nodes))
            if item in self.nodes:
                del self.nodes[item]
            return item
        raise KeyError('pop from an empty priority queue')

    def empty(self):
        self.queue = []
        self.nodes = {}

#===============================================================================


text:
# %%
import os
import pandas as pd
import numpy as np
import datetime
import xlrd
from vininfo import Vin
from scripts import origensegunvin

# %% CARGA DE DATOS

# PAIS - ORIGINALES
costa_rica = pd.read_csv(r"F:\Bases de datos\Costa Rica\csv\costa_rica.csv")

# %%
costa_rica["ORIGEN"] = None

# %%
costa_rica.head()

# %%
costa_rica = origensegunvin(costa_rica, "IDENTIFICACION DEL VEHICULO")

# %% CAMBIAR POSICION DE CILINDROS - ARREGLAR
# PROBLEMA: CILINDROS TIENE QUE SER NUMEROS Y POSICION TIENE QUE SER LETRA
cambiarcilindroyposicion = costa_rica.loc[costa_rica["CILINDROS"].astype("str").str.isupper() == True]

# %%
costa_rica = costa_rica.drop(costa_rica[costa_rica["CILINDROS"].str.isalpha() == True].index)

# %%
cambiarcilindroyposicion.rename(columns={
                'CILINDROS': 'cilindros',
                'Posición Cilindros': 'CILINDROS',
                'cilindros': 'CILINDROS',
                },
                inplace=True)

# %%
costa_rica = pd.concat([costa_rica, cambiarcilindroyposicion], join='outer')
costa_rica.drop(columns=['cilindros'], inplace=True)

# %% EQUIVALENCIAS
costa_rica["MERCADO"] = "COSTA RICA"
costa_rica["CANTIDAD"] = 1
costa_rica["TRANSMISION"].replace('M', 'MANUAL', inplace=True)
costa_rica["TRANSMISION"].replace('A', 'AUTOMATICO', inplace=True)
costa_rica["DISPOSICION CILINDROS"] =  costa_rica["Posición Cilindros"].astype("str") + costa_rica["CILINDROS"].astype("str")

costa_rica.rename(columns={
                'MODELO': 'MODELOAREVISAR',
                'CODIGO MARCA': 'MARCA',
                'ESTILO': 'MODELO',
                'Número Motor': 'NUMERO MOTOR',
                'AÑO MODELO': 'AÑO',
                'CILINDRAJE': 'CILINDRADA',
                'TRACCION': 'TRACCION'
                },
                inplace=True)

# %%
costa_rica["CILINDRADA"].replace('nan', None, inplace=True)
costa_rica["CILINDRADA"].replace('NaN', None, inplace=True)

# %%
costa_rica["CILINDRADA"] = costa_rica["CILINDRADA"].astype("float").map('{:.0f}'.format)

# %%
costa_rica["AÑO"] = costa_rica["AÑO"].astype(int)
costa_rica["AÑO"] = costa_rica["AÑO"].astype(str)

# %% AGREGAR MARCAS
marcas = pd.read_excel(r'D:\Basededatos\Origen\Costa Rica\Codigo Marca Costa Rica.xlsx', engine='openpyxl' )

# %%
costa_rica = pd.merge(
    costa_rica,
    marcas[['MARCA', 'Nombre de marca']],
    on='MARCA'
    )

# %%
costa_rica.rename(columns={
                'MARCA': 'CODIGO DE MARCA',
                'Nombre de marca': 'MARCA',
                },
                inplace=True)

# %%

# %% GENERAR ORIGEN SEGUN VIN
costa_ricaconvin["ORIGEN"] = costa_ricaconvin["IDENTIFICACION DEL VEHICULO"].map(lambda x: Vin(x).country)
costa_ricaconvin['ORIGEN'].replace('China (Mainland)', 'China', inplace=True)
costa_ricaconvin['ORIGEN'].replace('Taiwan, China', 'China', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Cote d'Ivoire", 'Costa de Marfil', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Germany/West Germany", 'Alemania', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Korea (South)", 'Corea del Sur', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Saudi Arabia", 'Arabia Saudita', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"United Kingdom", 'Reino Unido', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Italy", 'Italia', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Greece", 'Grecia', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Belgium", 'Belgica', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"Luxembourg", 'Luxemburgo', inplace=True)
costa_ricaconvin['ORIGEN'].replace(r"United States", 'Estados Unidos', inplace=True)
costa_ricasinvin = pd.concat([costa_rica, costa_ricaconvin]).drop_duplicates(keep=False)
costa_rica = pd.concat([costa_ricaconvin, costa_ricasinvin])

# %%
columnasutiles = [
                  "IDENTIFICACION DEL VEHICULO",
                  "MERCADO",
                  "MARCA",
                  "MODELO",
                  "DISPOSICION CILINDROS",
                  "CILINDRADA",
                  "CILINDROS",
                  "TRACCION",
                  "TRANSMISION",
                  "AÑO",
                  "CANTIDAD",
                  "NUMERO MOTOR",
                  "ORIGEN",
                  "PARTIDA_AR"]

costa_rica = costa_rica[columnasutiles]

# %%
costa_rica.rename(columns={"MODELO": "MODELO/VERSION", "IDENTIFICACION DEL VEHICULO": "NUMERO CHASIS / VIN"}, inplace=True)
new = costa_rica["MODELO/VERSION"].str.split(" ", n = 1, expand = True)
costa_rica["MODELO"] = new[0]

# %%
costa_rica["SEGMENTO.1"] = None
listamarca = r'''(HINO|ISUZU|MITSUBISHI|MACK|VOLVO|INTERNATIONAL|FREIGHLINER)'''
condicion1 = costa_rica["PARTIDA_AR"].str.contains("8704", na=False)
condicion2 = costa_rica["MARCA"].str.contains(listamarca)
costa_rica[condicion].head()
costa_rica.loc[condicion1 & condicion2, "SEGMENTO.1"] = "CAMION"

# %%
costa_rica.head()

# %%
costa_rica.to_csv(r'D:\Basededatos\Limpioparaunir\costa_rica.csv', index=False)

# %%
costa_rica[condicion1 & condicion2]

# %%
listamarca = r'''(HINO|ISUZU|MITSUBISHI|MACK|VOLVO|INTERNATIONAL|FREIGHLINER)'''
condicion1 = costa_rica["PARTIDA_AR"].str.contains("8704", na=False)
condicion2 = costa_rica["MARCA"].str.contains(listamarca)

# %%

max_stars_repo_name: Demonliquid/cars-python-cleaning

max_issues_repo_name: Demonliquid/cars-python-cleaning

max_forks_repo_name: Demonliquid/cars-python-cleaning


text:
import glob
import os
import random
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
import matplotlib.pyplot as plt
import cv2
import torchio as tio
from seg_model_utils.torchio_transforms import *

class BraTS2021(Dataset):
    """
    Dataset for BraTS2021 challenge - Includes tasks 1 and 2
    """

    def __init__(self, mode,
                 npy_fns_list,
                 label_list=[],
                 augmentations=None,
                 tio_augmentations=None,
                 volume_normalize=True,
                 segmentation_classes=False,
                 max_out_size=(256,256,96),
                 label_smoothing=0.4
                ):
        """
        :param mode: 'train','val','test'
        :param npy_fns_list: list of numpy array paths
        :param label_list: list of binary label integers
        :param augmentations: 3D augmentations
        :param tio_augmentations: TorchIO 3D augmentations
        :param volume_normalize: z-score normalize each channel in sample
        :param segmentation_classes: use two tumor segmentation classes for No-MGMT and MGMT
        :param max_out_size: max out dimensions (x,y,z), if data is larger, it is cropped to this
        :param label_smoothing: applied to label values only on train mode
        """
        self.mode = mode
        self.augmentations = augmentations
        self.tio_augmentations = tio_augmentations
        self.volume_normalize = volume_normalize
        self.segmentation_classes = segmentation_classes
        self.max_out_size = max_out_size
        self.label_smoothing = label_smoothing

        self.seg_fn_list = []
        self.fn_list = npy_fns_list
        self.label_list = label_list

        if self.mode != "test":
            fn_list_with_seg = []
            label_list_with_seg = []
            for fn, lbl in zip(npy_fns_list, label_list):
                fn_seg = str(fn).replace('.npy', '_seg.npy')
                if os.path.exists(fn_seg) and os.path.exists(fn):
                    fn_list_with_seg.append(fn)
                    self.seg_fn_list.append(fn_seg)
                    label_list_with_seg.append(int(lbl))

            # update fn list and exclude ones that don't have seg map
            self.fn_list = fn_list_with_seg
            self.label_list = label_list_with_seg
            assert len(self.fn_list) == len(self.seg_fn_list) == len(self.label_list)

        ########################
        #### Sanity checks #####

        assert len(self.fn_list) > 0

        # Check one sample
        one_sample = np.load(self.fn_list[0])
        assert len(one_sample.shape) == 4
        d,x,y,c = one_sample.shape
        self.full_vol_dim = (x,y,d)
        self.channels = c

        # Check one seg map
        if len(self.seg_fn_list) > 0:
            one_seg_sample = np.load(self.seg_fn_list[0])
            assert len(one_seg_sample.shape) == 3
            seg_d,seg_x,seg_y = one_seg_sample.shape
            assert seg_d == d
            assert seg_x == x
            assert seg_y == y

        #### Sanity checks #####
        ########################

        # shuffle samples
        if self.mode == "train":
            all_lists = list(zip(self.fn_list, self.seg_fn_list, self.label_list))
            random.shuffle(all_lists)
            self.fn_list, self.seg_fn_list, self.label_list = zip(*all_lists)

    def __len__(self):
        return len(self.fn_list)

    def __getitem__(self, index):
        sample = np.load(self.fn_list[index]).astype(np.float32)
        # fillna
        sample = np.nan_to_num(sample, nan=0)
        out_dict = {
            'BraTSID' : int(os.path.basename(self.fn_list[index]).split('.')[0]),
            'image' : sample
        }

        # Calculate stats from full image
        sample_mean = np.mean(sample, axis=tuple([0,1,2]))
        sample_std = np.std(sample, axis=tuple([0,1,2])) + 1e-6

        # Load labels
        if self.mode != "test":
            seg = np.load(self.seg_fn_list[index])
            # set seg to binary values
            seg = (seg > 0).astype(np.float32)
            lbl = self.label_list[index]
            if self.mode == 'train':
                if lbl > 0.5:
                    lbl = lbl - self.label_smoothing
                else:
                    lbl = self.label_smoothing
            out_dict['segmentation'] = seg
            out_dict['label'] = lbl

        # augment
        if self.augmentations:
            out_dict = self.augmentations(out_dict)

        # TorchIO augs
        if self.tio_augmentations:
            subject = tio.Subject(
                rgb_image = arr_2_tio_image(out_dict['image']),
                segmentation = arr_2_tio_seg_image(out_dict['segmentation'])
            )
            subject = self.tio_augmentations(subject)
            out_dict['image'] = subject.rgb_image.numpy().swapaxes(0,3)
            out_dict['segmentation'] = subject.segmentation.numpy().swapaxes(0,3)[:,:,:,0]

        # Process segmentation map
        if 'segmentation' in out_dict:
            if self.segmentation_classes:
                # split segmentation to two
                empty_seg = np.zeros_like(out_dict['segmentation'])
                if lbl == 0:
                    out_dict['segmentation'] = np.stack([out_dict['segmentation'], empty_seg], axis=3)
                else:
                    out_dict['segmentation'] = np.stack([empty_seg, out_dict['segmentation']], axis=3)
            else:
                # Add channel dimension to segmentation
                out_dict['segmentation'] = np.expand_dims(out_dict['segmentation'], axis=3)

            # swap channel axis in image with depth axis - torch wants channel first
            out_dict['segmentation'] = np.swapaxes(out_dict['segmentation'], 0, 3)

        # z-score norm each channel - done after augmentations
        if self.volume_normalize:
            sample = out_dict['image']
            sample = (sample - sample_mean) / sample_std
            out_dict['image'] = sample
            out_dict['mean'] = sample_mean
            out_dict['std'] = sample_std
        else:
            out_dict['mean'] = np.array([0 for _ in range(sample.shape[3])])
            out_dict['std'] = np.array([1. for _ in range(sample.shape[3])])

        # swap channel axis in image with depth axis - torch wants channel first
        out_dict['image'] = np.swapaxes(out_dict['image'], 0, 3)

        # check that size doesn't exceed max size
        im_shape = list(out_dict['image'].shape[1:])
        for i, max_sz in enumerate(self.max_out_size):
            im_shape[i] = min(im_shape[i], max_sz)
        out_dict['image'] = out_dict['image'][:,:im_shape[0],:im_shape[1],:im_shape[2]]
        if 'segmentation' in out_dict:
            out_dict['segmentation'] = out_dict['segmentation'][:,:im_shape[0],:im_shape[1],:im_shape[2]]

        # convert values to torch and arrays
        for key, val in out_dict.items():
            if key == 'label':
                out_dict[key] = np.array(val)
            elif key == 'BraTSID':
                out_dict[key] = np.array(val)
            else:
                out_dict[key] = torch.FloatTensor(val.copy())

        return out_dict


text:
# -*- coding: utf-8 -*-
r"""
This example is a complete study of a crystal analyzer. The generator
``plot_generator()`` is rather complex and therefore needs some explanation.
The main loop changes the source type. After a flat-energy source has been
ray-traced (after ``yield``), the widths of *z* and energy distributions are
saved. Then a single line source is ray-traced and provides the width of *z*
distribution. From these 3 numbers we calculate energy resolution and, as a
check, ray-trace a third source with 7 energy lines with a spacing equal to the
previously calculated energy resolution. The source sizes, axis limits, number
of iterations etc. were determined experimentally and are given by lists in the
upper part of the script. The outputs are the plots and a text file with the
resulted energy resolutions."""
__author__ = "Konstantin Klementiev"
__date__ = "08 Mar 2016"
import os, sys; sys.path.append(os.path.join('..', '..', '..'))  # analysis:ignore
import numpy as np
#import matplotlib as mpl

import xrt.backends.raycing as raycing
import xrt.backends.raycing.sources as rs
#import xrt.backends.raycing.apertures as ra
import xrt.backends.raycing.oes as roe
import xrt.backends.raycing.run as rr
import xrt.backends.raycing.materials as rm
import xrt.backends.raycing.screens as rsc

import xrt.plotter as xrtp
import xrt.runner as xrtr

showIn3D = False

crystalMaterial = 'Si'
if crystalMaterial == 'Si':
    d111 = 3.1354161
elif crystalMaterial == 'Ge':
    d111 = 3.2662725
else:
    raise

crystal = rm.CrystalDiamond((4, 4, 4), d111/4, elements=crystalMaterial)
#numiter = 640
numiter = 40

Rm = 1e9  # meridional radius, mm
Rs = 1000  # tmp sagittal radius, mm
#Rs = 250  # tmp sagittal radius, mm
dphi = 0

beamV = 0.1/2.35  # vertical beam size
beamH = 0.1/2.35  # horizontal beam size

yAxesLim = 20

dxCrystal = 100.
dyCrystal = 100.
#dyCrystal = 50.
#dxCrystal = 300.
#dyCrystal = 70.

isDiced = False
yAxisLim = 32  # Mythen length = 64 mm
yAxisLine = 0.4
if isDiced:
    xAxisLim = 12
    facetKWargs = {'dxFacet': 1-0.05, 'dyFacet': 50.,
                   'dxGap': 0.05, 'dyGap': 0.}
    Toroid = roe.DicedJohannToroid
    analyzerName = crystalMaterial + 'vonHamosDiced'
else:
    xAxisLim = 5
    facetKWargs = {}
    Toroid = roe.JohannToroid
    analyzerName = crystalMaterial + 'vonHamos'

thetaDegree = 40

if isDiced:
    dphi = np.arcsin((facetKWargs['dxFacet'] + facetKWargs['dxGap']) / Rs)

if thetaDegree == 40:
    if Rs > 800:
        eAxisFlat = 7.5e-3  # @ 40 deg, R=1000
    else:
        eAxisFlat = 3e-2  # @ 40 deg
elif thetaDegree == 60:
    if Rs > 800:
        eAxisFlat = 6.8e-3  # @ 60 deg, R=1000
    else:
        eAxisFlat = 2.6e-2  # @ 60 deg
elif thetaDegree == 80:
    if Rs > 800:
        eAxisFlat = 2.6e-3  # @ 80 deg, R=1000
    else:
        eAxisFlat = 9.0e-3  # @ 80 deg
else:
    raise

def build_beamline(nrays=2e5):
    beamLine = raycing.BeamLine(azimuth=0, height=0)
    rs.GeometricSource(
        beamLine, 'GeometricSource', nrays=nrays, dx=beamH, dy=0,
        dz=beamV, distxprime='flat', distzprime='flat', polarization=None)
    beamLine.analyzer = Toroid(
        beamLine, analyzerName, surface=('',),
        limPhysX=(-dxCrystal/2, dxCrystal/2),
        limPhysY=(-dyCrystal/2*3, dyCrystal/2*3),
        Rm=Rm, Rs=Rs, shape='rect',
        **facetKWargs)
    beamLine.detector = rsc.Screen(beamLine, 'Detector', z=(0, 0, 1))
#    beamLine.s1h = ra.RectangularAperture(
#        beamLine, 'horizontal. slit', 0, Rs-10.,
#        ('left', 'right'), [-0.1, 0.1])
    return beamLine

def run_process(beamLine):
    beamSource = beamLine.sources[0].shine()
#    beamLine.s1h.propagate(beamSource)
    beamAnalyzerGlobal, beamAnalyzerLocal = \
        beamLine.analyzer.reflect(beamSource)
    beamDetector = beamLine.detector.expose(beamAnalyzerGlobal)
    outDict = {'beamSource': beamSource,
               'beamAnalyzerGlobal': beamAnalyzerGlobal,
               'beamAnalyzerLocal': beamAnalyzerLocal,
               'beamDetector': beamDetector
               }
    if showIn3D:
        beamLine.prepare_flow()
    return outDict
rr.run_process = run_process

def align_spectrometer_Rs(beamLine, theta, Rs):
    sinTheta = np.sin(theta)
    cosTheta = np.cos(theta)
    sin2Theta = np.sin(2 * theta)
    p = Rs / sinTheta
    yDet = p * 2 * cosTheta**2
    zDet = p * sin2Theta

    beamLine.analyzer.center = 0, p, 0
    beamLine.analyzer.pitch = theta
    beamLine.detector.center = 0, yDet, zDet
    beamLine.detector.set_orientation(z=(0, cosTheta, sinTheta))

    beamLine.sources[0].dxprime = 1.1 * dxCrystal / p
    beamLine.sources[0].dzprime = dyCrystal * np.sin(theta) / p
    print('theta={0}deg, Rs={1}mm: p={2}mm'.format(np.degrees(theta), Rs, p))

#def align_spectrometer_p(beamLine, theta, p):
#    sinTheta = np.sin(theta)
#    cosTheta = np.cos(theta)
#    sin2Theta = np.sin(2 * theta)
#    Rs = p * sinTheta
#    yDet = p * 2 * cosTheta**2
#    zDet = p * sin2Theta
#
#    beamLine.analyzer.center = 0, p, 0
#    beamLine.analyzer.Rs = Rs
#    beamLine.analyzer.pitch = theta
#    beamLine.detector.center = 0, yDet, zDet
#    beamLine.detector.set_orientation(z=(0, cosTheta, sinTheta))
#
#    beamLine.sources[0].dxprime = 1.1 * dxCrystal / p
#    beamLine.sources[0].dzprime = dyCrystal * np.sin(theta) / p
#    print('theta={0}deg, p={1}mm: Rs={2}mm'.format(np.degrees(theta), p, Rs))

def stripe_number(beam):
    phi = np.arcsin(beam.x / Rs)
    return np.round(phi / dphi)

def define_plots(beamLine):
    fwhmFormatStrE = '%.2f'
    plots = []
    plotsAnalyzer = []
    plotsDetector = []
    plotsE = []

    limits = [-dxCrystal/2 - 5, dxCrystal/2 + 5]
    plotAnE = xrtp.XYCPlot(
        'beamAnalyzerLocal', (1,),
        xaxis=xrtp.XYCAxis(r'$x$', 'mm', bins=400, ppb=1, limits=limits),
        yaxis=xrtp.XYCAxis(r'$y$', 'mm', bins=400, ppb=1, limits=limits),
        caxis=xrtp.XYCAxis('energy', 'eV', fwhmFormatStr='%.2f',
                           bins=200, ppb=2),
        title='xtal_E', oe=beamLine.analyzer, raycingParam=1000)
    plotAnE.caxis.fwhmFormatStr = fwhmFormatStrE
    plotAnE.textPanel = plotAnE.fig.text(
        0.88, 0.85, '', transform=plotAnE.fig.transFigure, size=14, color='r',
        ha='center')
    plotsAnalyzer.append(plotAnE)
    plotsE.append(plotAnE)

    if isDiced:
        plotAnS = xrtp.XYCPlot(
            'beamAnalyzerLocal', (1,),
            xaxis=xrtp.XYCAxis(r'$x$', 'mm', bins=400, ppb=1, limits=limits),
            yaxis=xrtp.XYCAxis(r'$y$', 'mm', bins=400, ppb=1, limits=limits),
            caxis=xrtp.XYCAxis('stripe number', '', data=stripe_number),
            title='xtal_stripes', oe=beamLine.analyzer, raycingParam=1000)
        plotAnS.textPanel = plotAnS.fig.text(
            0.88, 0.85, '', transform=plotAnS.fig.transFigure, size=14,
            color='r', ha='center')
        plotsAnalyzer.append(plotAnS)

    if isDiced:
        xMax = (facetKWargs['dxFacet']+1) / 2
        limits = [-xMax, xMax]
    else:
        limits = [-3, 3]
    plot = xrtp.XYCPlot(
        'beamAnalyzerLocal', (1,),
        xaxis=xrtp.XYCAxis(r'$x$', 'mm', limits=limits),
        yaxis=xrtp.XYCAxis(r'$y$', 'mm', limits=limits),
        caxis=xrtp.XYCAxis('energy', 'eV', fwhmFormatStr='%.2f'),
        title='xtal_E_zoom', oe=beamLine.analyzer, raycingParam=1000)
    plot.caxis.fwhmFormatStr = fwhmFormatStrE
    plot.textPanel = plot.fig.text(
        0.88, 0.8, '', transform=plot.fig.transFigure, size=14, color='r',
        ha='center')
    plotsAnalyzer.append(plot)
    plotsE.append(plot)

    if isDiced:
        plot = xrtp.XYCPlot(
            'beamAnalyzerLocal', (1,),
            xaxis=xrtp.XYCAxis(r'$x$', 'mm', limits=limits),
            yaxis=xrtp.XYCAxis(r'$y$', 'mm', limits=limits),
            caxis=xrtp.XYCAxis('stripe number', '', data=stripe_number),
            title='xtal_stripes_zoom', oe=beamLine.analyzer,
            raycingParam=1000)
        plot.textPanel = plot.fig.text(
            0.88, 0.8, '', transform=plot.fig.transFigure, size=14, color='r',
            ha='center')
        plotsAnalyzer.append(plot)

    plotDetE = xrtp.XYCPlot(
        'beamDetector', (1,), aspect='auto',
        xaxis=xrtp.XYCAxis(r'$x$', 'mm', fwhmFormatStr='%.3f'),
        yaxis=xrtp.XYCAxis(r'$z$', 'mm', fwhmFormatStr='%.3f'),
        title='det_E')
    plotDetE.xaxis.limits = 'symmetric'
    plotDetE.yaxis.limits = 'symmetric'
#    plotDetE.xaxis.limits = -xAxisLim, xAxisLim
#    plotDetE.yaxis.limits = -yAxisLim, yAxisLim
    plotDetE.caxis.fwhmFormatStr = fwhmFormatStrE
    plotDetE.textPanel = plotDetE.fig.text(
        0.88, 0.8, '', transform=plotDetE.fig.transFigure, size=14, color='r',
        ha='center')
    plotsDetector.append(plotDetE)
    plotsE.append(plotDetE)

    if isDiced:
        plotDetS = xrtp.XYCPlot(
            'beamDetector', (1,), aspect='auto',
            xaxis=xrtp.XYCAxis(r'$x$', 'mm', fwhmFormatStr='%.3f'),
            yaxis=xrtp.XYCAxis(r'$z$', 'mm', fwhmFormatStr='%.3f'),
            caxis=xrtp.XYCAxis('stripe number', '', data=stripe_number),
            beamC='beamAnalyzerLocal', title='det_stripes')
        plotDetS.xaxis.limits = -xAxisLim, xAxisLim
        plotDetS.yaxis.limits = -yAxisLim, yAxisLim
        plotDetS.textPanel = plotDetS.fig.text(
            0.88, 0.8, '', transform=plotDetS.fig.transFigure, size=14,
            color='r', ha='center')
        plotsDetector.append(plotDetS)

    for plot in plotsAnalyzer:
        plots.append(plot)
    for plot in plotsDetector:
        plots.append(plot)
    return plots, plotsAnalyzer, plotsDetector, plotsE, plotDetE

def plot_generator(beamLine, plots=[], plotsAnalyzer=[], plotsDetector=[],
                   plotsE=[], plotDetE=[]):
    hklSeparator = ',' if np.any(np.array(crystal.hkl) > 10) else ''
    crystalLabel = '{0}{1[0]}{2}{1[1]}{2}{1[2]}'.format(
        crystalMaterial, crystal.hkl, hklSeparator)
    beamLine.analyzer.surface = crystalLabel,
    for plot in plotsAnalyzer:
        plot.draw_footprint_area()
    beamLine.analyzer.material = crystal
    theta = np.radians(thetaDegree)
    align_spectrometer_Rs(beamLine, theta, Rs)

    sinTheta = np.sin(theta)
    E0raw = rm.ch / (2 * crystal.d * sinTheta)
    dTheta = crystal.get_dtheta_symmetric_Bragg(E0raw)
    E0 = rm.ch / (2 * crystal.d * np.sin(theta + dTheta))
    offsetE = round(E0, 3)

    dELine = 0
    dzLine = 0
    for isource in np.arange(3):
#    for isource in [-1, ]:
        xrtr.set_repeats(numiter)
        if isource == 0 or isource == -1:  # flat or norm
#            xrtr.set_repeats(0)
            eAxisMin = E0 * (1 - eAxisFlat)
            eAxisMax = E0 * (1 + eAxisFlat)
            dELine = E0 * eAxisFlat/3.  # for showIn3D
            for plot in plotsE:
                if plot is None:
                    continue
                plot.caxis.offset = int(round(offsetE, -2))
                plot.caxis.limits = [eAxisMin, eAxisMax]
            tt = r'{0}{1}$\theta = {2:.0f}^\circ$'.format(
                crystalLabel, '\n', thetaDegree)
            for plot in plots:
                try:
                    plot.textPanel.set_text(tt)
                except AttributeError:
                    pass
            if isource == -1:
                beamLine.sources[0].distE = 'normal'
                beamLine.sources[0].energies = E0, eAxisFlat/2.355
                sourcename = 'norm'
            else:
                beamLine.sources[0].distE = 'flat'
                beamLine.sources[0].energies = eAxisMin, eAxisMax
                sourcename = 'flat'
        elif isource == 1:  # line
#            xrtr.set_repeats(0)
            beamLine.sources[0].distE = 'lines'
            beamLine.sources[0].energies = E0,
            sourcename = 'line'
            for plot in plotsDetector:
                plot.yaxis.limits = [-yAxisLine, yAxisLine]
        else:
#            xrtr.set_repeats(2560*16L)
            for plot in plotsDetector:
                plot.xaxis.limits = -6, 6
            tt = (r'{0}{1}$\theta = {2:.0f}^\circ${1}$' +
                  '\delta E = ${3:.3f} eV').format(
                crystalLabel, '\n', thetaDegree, dELine)
            for plot in plots:
                try:
                    plot.textPanel.set_text(tt)
                except AttributeError:
                    pass
            beamLine.sources[0].distE = 'lines'
            sourcename = '7lin'
            for plot in plotsDetector:
                plot.yaxis.limits = [-dzLine*7, dzLine*7]
            dEStep = dELine
            beamLine.sources[0].energies = \
                [E0 + dEStep * i for i in range(-3, 4)]
            eAxisMin = E0 - dEStep * 4
            eAxisMax = E0 + dEStep * 4
            for plot in plotsE:
                if plot is None:
                    continue
                plot.caxis.limits = [eAxisMin, eAxisMax]

        for plot in plots:
            filename = '{0}{1}-{2}-{3}'.format(
                analyzerName, thetaDegree, plot.title, sourcename)
            plot.saveName = filename + '.png'
#            plot.persistentName = filename + '.pickle'
        if showIn3D:
            beamLine.glowFrameName = \
                '{0}{1}-{2}-{3}.jpg'.format(analyzerName, thetaDegree, isource,
                                            sourcename)
        yield

        if not showIn3D:
            if isource == 0:
                dzFlat = plotDetE.dy
                dEFlat = plotDetE.dE
            elif isource == 1:
                dzLine = plotDetE.dy
                try:
                    dELine = dzLine * dEFlat / dzFlat
                except:
                    print('dzFlat={0}'.format(dzFlat))
                    dELine = 0

def main():
    beamLine = build_beamline()
    if showIn3D:
        beamLine.glow(scale=4, centerAt=analyzerName,
                      generator=plot_generator, generatorArgs=[beamLine])
        return
    plots, plotsAnalyzer, plotsDetector, plotsE, plotDetE = \
        define_plots(beamLine)
    args = [beamLine, plots, plotsAnalyzer, plotsDetector, plotsE, plotDetE]
    xrtr.run_ray_tracing(
        plots, generator=plot_generator, generatorArgs=args,
        beamLine=beamLine, processes='half')

#this is necessary to use multiprocessing in Windows, otherwise the new Python
#contexts cannot be initialized:
if __name__ == '__main__':
    main()

max_stars_repo_name: adinatan/xrt

max_issues_repo_name: adinatan/xrt

max_forks_repo_name: adinatan/xrt


text:
import os
import sys
import json
import gzip

import pandas as pd
import numpy as np
from nameparser import HumanName
import unicodedata

# determine if we are loading from a jupyter notebook (to make pretty progress bars)
if 'ipykernel' in sys.modules:
    from tqdm.notebook import tqdm
else:
    from tqdm import tqdm

from pyscisci.datasource.readwrite import load_preprocessed_data, load_int, load_float
from pyscisci.database import BibDataBase

class MAG(BibDataBase):
    """
    Base class for Microsoft Academic Graph interface.

    This is an extension of 'BibDataBase' with processing functions specific to the MAG.
    See 'BibDataBase' in database.py for details of non-MAG specific functions.

    The MAG comes structured into three folders: mag, advanced, nlp.
    Explain downloading etc.

    """

    def __init__(self, path2database = '', keep_in_memory = False, global_filter = None, show_progress=True):

        self._default_init(path2database, keep_in_memory, global_filter, show_progress)

        self.PublicationIdType = int
        self.AffiliationIdType = int
        self.AuthorIdType = int

    def preprocess(self, dflist = None, show_progress=True):
        """
        Bulk preprocess the MAG raw data.

        Parameters
        ----------
        dflist: list, default None
            The list of DataFrames to preprocess.  If None, all MAG DataFrames are preprocessed.

        show_progress: bool, default True
            Show progress with processing of the data.

        """
        if dflist is None:
            dflist = ['affiliation', 'author', 'publication', 'reference', 'publicationauthoraffiliation', 'fields']

        if 'affiliation' in dflist:
            self.parse_affiliations(preprocess = True, show_progress=show_progress)

        if 'author' in dflist:
            self.parse_authors(preprocess = True, show_progress=show_progress)

        if 'publication' in dflist:
            self.parse_publications(preprocess = True, show_progress=show_progress)

        if 'reference' in dflist:
            self.parse_references(preprocess = True, show_progress=show_progress)

        if 'publicationauthoraffiliation' in dflist:
            self.parse_publicationauthoraffiliation(preprocess = True, show_progress=show_progress)

        if 'fields' in dflist:
            self.parse_fields(preprocess=True, show_progress=show_progress)

    def download_from_source(self):
        #TODO: Error Raising when database isnt defined
        raise NotImplementedError("ToDo")

    def parse_affiliations(self, preprocess=True, show_progress=True):
        """
        Parse the MAG Affilation raw data.

        Parameters
        ----------
        preprocess: bool, default True
            Save the processed data in new DataFrames.

        show_progress: bool, default True
            Show progress with processing of the data.

        Returns
        ----------
        DataFrame
            Affiliation DataFrame.
        """

        affil_int_columns = [0, 7, 8]
        affil_str_columns = [3, 4, 5, 6]
        affil_float_columns = [9, 10]

        affil_column_names = ['AffiliationId', 'NumberPublications', 'NumberCitations', 'FullName', 'GridId', 'OfficialPage', 'WikiPage', 'Latitude', 'Longitude']

        file_name = os.path.join(self.path2database, 'mag', 'Affiliations.txt')

        affiliation_info = []
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, desc='Affiliations', leave=True, disable=not show_progress) as pbar:
            with open(file_name, 'r') as infile:
                for line in infile:
                    sline = line.replace('\n', '').split('\t')
                    affline = [load_int(sline[i]) for i in affil_int_columns]
                    affline += [sline[i] for i in affil_str_columns]
                    affline += [load_float(sline[i]) for i in affil_float_columns]
                    affiliation_info.append(affline)

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

        aff_df = pd.DataFrame(affiliation_info, columns = affil_column_names)

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'affiliation')):
                os.mkdir(os.path.join(self.path2database, 'affiliation'))
            aff_df.to_hdf(os.path.join(self.path2database, 'affiliation', 'affiliation0.hdf'), key = 'affiliation', mode = 'w')

        return aff_df

    def parse_authors(self, preprocess = False, process_name = True, num_file_lines = 5*10**6, show_progress=True):
        """
        Parse the MAG Author raw data.

        Parameters
        ----------
        preprocess: bool, default True
            Save the processed data in new DataFrames.

        process_name: bool, default True
            If True, then when processing the raw file, the package `NameParser <https://nameparser.readthedocs.io/en/latest/>`_
            will be used to split author FullNames.

        num_file_lines: int, default 5*10**6
            The processed data will be saved into smaller DataFrames, each with `num_file_lines` rows.

        show_progress: bool, default True
            Show progress with processing of the data.

        Returns
        ----------
        DataFrame
            Author DataFrame.
        """

        author_int_columns = [0, 4, 5, 6]

        author_column_names = ['AuthorId', 'LastKnownAffiliationId', 'NumberPublications', 'NumberCitations', 'FullName']
        if process_name:
            author_column_names += ['LastName', 'FirstName', 'MiddleName']

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'author')):
                os.mkdir(os.path.join(self.path2database, 'author'))

        file_name = os.path.join(self.path2database, 'mag', 'Authors.txt')

        iauthor = 0
        ifile = 0
        authorinfo = []
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, desc='Authors', disable=not show_progress, leave=True) as pbar:
            with open(file_name, 'r') as infile:

                for line in infile:

                    # split the line and keep only the relevant columns
                    sline = line.split('\t')
                    adata = [load_int(sline[ip]) for ip in author_int_columns] + [sline[2]]

                    # process the first, middle, and last names for the author
                    if process_name:
                        hname = HumanName(unicodedata.normalize('NFD', sline[2]))
                        adata += [hname.last, hname.first, hname.middle]

                    authorinfo.append(adata)
                    iauthor += 1

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

                    # time to save
                    if preprocess and iauthor % num_file_lines == 0:
                        pd.DataFrame(authorinfo, columns = author_column_names).to_hdf(
                            os.path.join(self.path2database, 'author', 'author{}.hdf'.format(ifile)),
                                                                                key = 'author', mode = 'w')

                        ifile += 1
                        authorinfo = []

                author_df = pd.DataFrame(authorinfo, columns = author_column_names)
                if preprocess:
                    author_df.to_hdf(os.path.join(self.path2database, 'author', 'author{}.hdf'.format(ifile)),
                                                                                key = 'author', mode = 'w')

        return author_df

    def parse_publications(self, preprocess = True, num_file_lines=5*10**6, preprocess_dicts = True, show_progress=True):
        """
        Parse the MAG Publication and Journal raw data.

        Parameters
        ----------
        preprocess: bool, default True
            Save the processed data in new DataFrames.

        num_file_lines: int, default 5*10**6
            The processed data will be saved into smaller DataFrames, each with `num_file_lines` rows.

        preprocess_dicts: bool, default True
            Save the processed Year and DocType data as dictionaries.

        show_progress: bool, default True
            Show progress with processing of the data.

        Returns
        ----------
        DataFrame
            Publication DataFrame.
        """

        # first do the journal information
        journal_str_col = [2, 4, 5, 6]
        journal_column_names = ['JournalId', 'FullName', 'Issn', 'Publisher', 'Webpage']

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'journal')):
                os.mkdir(os.path.join(self.path2database, 'journal'))

        file_name = os.path.join(self.path2database, 'mag', 'Journals.txt')

        journal_info = []
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, desc='Journals', leave=True, disable=not show_progress) as pbar:
            with open(os.path.join(self.path2database, 'mag', 'Journals.txt'), 'r') as infile:
                for line in infile:
                    # split the line and keep only the relevant columns
                    sline = line.replace('\n', '').split('\t')
                    jline = [load_int(sline[0])] + [sline[i] for i in journal_str_col]
                    journal_info.append(jline)

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

        journal_df = pd.DataFrame(journal_info, columns = journal_column_names)
        if preprocess:
            journal_df.to_hdf(os.path.join(self.path2database, 'journal', 'journal.hdf'), key = 'journal', mode = 'w')

        #now lets do the publication information

        doctype = {'Journal': 'j', 'Book':'b', '':'', 'BookChapter':'bc', 'Conference':'c', 'Dataset':'d', 'Patent':'p', 'Repository':'r'}

        pub_int_columns = [0, 7, 10, 21]
        pub_str_columns = [2, 4, 8, 13, 14]
        pub_column_names = ['PublicationId', 'Year', 'JournalId', 'FamilyId',  'Doi', 'Title', 'Date', 'Volume', 'Issue', 'DocType']

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'publication')):
                os.mkdir(os.path.join(self.path2database, 'publication'))

        file_name = os.path.join(self.path2database, 'mag', 'Papers.txt')

        ipub = 0
        ifile = 0
        pubinfo = []

        pub2year = {}
        pub2doctype = {}
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, desc='Publications', leave=True, disable=not show_progress) as pbar:
            with open(file_name, 'r') as infile:
                for line in infile:
                    # split the line and keep only the relevant columns
                    sline = line.replace('\n', '').split('\t')
                    pline = [load_int(sline[ip]) for ip in pub_int_columns] + [sline[ip] for ip in pub_str_columns] + [doctype[sline[3]]]
                    pub2year[pline[0]] = pline[1]
                    if doctype[sline[3]] != '':
                        pub2doctype[pline[0]] = doctype[sline[3]]

                    pubinfo.append(pline)
                    ipub += 1

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

                    if preprocess and ipub % num_file_lines == 0:
                            pd.DataFrame(pubinfo, columns = pub_column_names).to_hdf(
                                os.path.join(self.path2database, 'publication', 'publication{}.hdf'.format(ifile)),
                                                                                key = 'publication', mode = 'w')

                            ifile += 1
                            pubinfo = []

            pub_df = pd.DataFrame(pubinfo, columns = pub_column_names)
            if preprocess:
                pub_df.to_hdf(os.path.join(self.path2database, 'publication', 'publication{}.hdf'.format(ifile)),
                                                                                key = 'publication', mode = 'w')

                if preprocess_dicts:
                    with gzip.open(os.path.join(self.path2database, 'pub2year.json.gz'), 'w') as outfile:
                        outfile.write(json.dumps(pub2year).encode('utf8'))

                    with gzip.open(os.path.join(self.path2database, 'pub2doctype.json.gz'), 'w') as outfile:
                        outfile.write(json.dumps(pub2doctype).encode('utf8'))

        return pub_df

    def parse_references(self, preprocess = False, num_file_lines=10**7, show_progress=True):
        """
        Parse the MAG References raw data.

        Parameters
        ----------
        preprocess: bool, default True
            Save the processed data in new DataFrames.

        num_file_lines: int, default 10**7
            The processed data will be saved into smaller DataFrames, each with `num_file_lines` rows.

        show_progress: bool, default True
            Show progress with processing of the data.

        Returns
        ----------
        DataFrame
            Pub2Ref DataFrame.
        """
        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'pub2ref')):
                os.mkdir(os.path.join(self.path2database, 'pub2ref'))

        file_name = os.path.join(self.path2database, 'mag', 'PaperReferences.txt')

        iref = 0
        ifile = 0
        pub2ref_info = []
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, unit_divisor=1024, desc='References', leave=True, disable=not show_progress) as pbar:
            with open(file_name, 'r') as infile:
                for line in infile:
                    # split the line and keep only the relevant columns
                    sline = line.replace('\n', '').split('\t')
                    pub2ref_info.append([load_int(sline[ip]) for ip in range(2)])
                    iref += 1

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

                    if preprocess and iref % num_file_lines == 0:
                        pd.DataFrame(pub2ref_info, columns = ['CitingPublicationId', 'CitedPublicationId']).to_hdf(
                            os.path.join(self.path2database, 'pub2ref', 'pub2ref{}.hdf'.format(ifile)),
                                                                                key = 'pub2ref', mode = 'w')

                        ifile += 1
                        pub2ref_info = []

        pub2ref_df = pd.DataFrame(pub2ref_info, columns = ['CitingPublicationId', 'CitedPublicationId'])

        if preprocess:
            pub2ref_df.to_hdf(os.path.join(self.path2database, 'pub2ref', 'pub2ref{}.hdf'.format(ifile)),
                                                                            key = 'pub2ref', mode = 'w')

        return pub2ref_df

    def parse_publicationauthoraffiliation(self, preprocess = False, num_file_lines=5*10**6, show_progress=True):
        """
        Parse the MAG PublicationAuthorAffiliation raw data.

        Parameters
        ----------
        preprocess: bool, default True
            Save the processed data in new DataFrames.

        num_file_lines: int, default 10**7
            The processed data will be saved into smaller DataFrames, each with `num_file_lines` rows.

        show_progress: bool, default True
            Show progress with processing of the data.

        Returns
        ----------
        DataFrame
            PublicationAuthorAffiliation DataFrame.
        """
        pubauthaff_int_columns = [0, 1, 2, 3]
        pubauthaff_str_columns = [4, 5]
        pub_column_names = ['PublicationId', 'AuthorId', 'AffiliationId', 'AuthorSequence',  'OrigAuthorName', 'OrigAffiliationName']

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'publicationauthoraffiliation')):
                os.mkdir(os.path.join(self.path2database, 'publicationauthoraffiliation'))

        file_name = os.path.join(self.path2database, 'mag', 'PaperAuthorAffiliations.txt')

        iref = 0
        ifile = 0
        pubauthaff_info = []
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, desc='PaperAuthorAffiliations', leave=True, disable=not show_progress) as pbar:
            with open(file_name, 'r') as infile:
                for line in infile:
                    sline = line.replace('\n', '').split('\t')
                    pubauthaff_info.append([load_int(sline[ip]) for ip in pubauthaff_int_columns] + [sline[ip] if len(sline) > ip else '' for ip in pubauthaff_str_columns ])
                    iref += 1

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

                    if preprocess and iref % num_file_lines == 0:
                        pd.DataFrame(pubauthaff_info, columns = pub_column_names).to_hdf(
                            os.path.join(self.path2database, 'publicationauthoraffiliation', 'publicationauthoraffiliation{}.hdf'.format(ifile)),
                                                                                key = 'publicationauthoraffiliation', mode = 'w')

                        ifile += 1
                        pubauthaff_info = []

        paa_df = pd.DataFrame(pubauthaff_info, columns = pub_column_names)
        if preprocess:
            paa_df.to_hdf(os.path.join(self.path2database, 'publicationauthoraffiliation', 'publicationauthoraffiliation{}.hdf'.format(ifile)),
                                                                        key = 'publicationauthoraffiliation', mode = 'w')
        return paa_df

    def parse_fields(self, preprocess = False, num_file_lines=10**7, show_progress=True):
        """
        Parse the MAG Paper Field raw data.

        Parameters
        ----------
        preprocess: bool, default True
            Save the processed data in new DataFrames.

        num_file_lines: int, default 10**7
            The processed data will be saved into smaller DataFrames, each with `num_file_lines` rows.

        show_progress: bool, default True
            Show progress with processing of the data.

        Returns
        ----------
        DataFrame
            Pub2Field DataFrame.
        """
        field2get = [0, 5, 6]
        fieldnames = ['FieldId', 'FieldLevel', 'NumberPublications', 'FieldName']

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'fieldinfo')):
                os.mkdir(os.path.join(self.path2database, 'fieldinfo'))

        fieldinfo = []
        with open(os.path.join(self.path2database, 'advanced', 'FieldsOfStudy.txt'), 'r') as infile:

            for line in infile:
                sline = line.split('\t')
                fielddata = [load_int(sline[ip]) for ip in field2get] + [sline[2]]
                fieldinfo.append(fielddata)

        field_df = pd.DataFrame(fieldinfo, columns = fieldnames)
        if preprocess:
            field_df.to_hdf(os.path.join(self.path2database, 'fieldinfo', 'fieldinfo0.hdf'), key = 'field', mode = 'w')

        # and now do pub2field
        paperfields = [0, 1]
        paperfieldnames = ['PublicationId', 'FieldId']

        if preprocess:
            if not os.path.exists(os.path.join(self.path2database, 'pub2field')):
                os.mkdir(os.path.join(self.path2database, 'pub2field'))

        file_name = os.path.join(self.path2database, 'advanced', 'PaperFieldsOfStudy.txt')

        ipaper = 0
        ifile = 0
        fieldinfo = []
        with tqdm(total=os.path.getsize(file_name), unit='iB', unit_scale=True, desc='PaperFieldsOfStudy', leave=True, disable=not show_progress) as pbar:
            with open(file_name, 'r') as infile:

                for line in infile:
                    sline = line.split('\t')
                    fielddata = [int(sline[ip]) for ip in paperfields]
                    fieldinfo.append(fielddata)
                    ipaper += 1

                    # update progress bar
                    pbar.update(sys.getsizeof(line))

                    if preprocess and ipaper % num_file_lines == 0:
                        pd.DataFrame(fieldinfo, columns = paperfieldnames).to_hdf(
                            os.path.join(self.path2database, 'pub2field', 'pub2field' + str(ifile) + '.hdf'),
                                                                                key = 'pub2field', mode = 'w')

                        ifile += 1
                        fieldinfo = []

        pub2field_df = pd.DataFrame(fieldinfo, columns = paperfieldnames)
        if preprocess:
            pub2field_df.to_hdf(os.path.join(self.path2database, 'pub2field', 'pub2field' + str(ifile) + '.hdf'),
                                                                                key = 'pub2field', mode = 'w')
        return pub2field_df

max_stars_repo_name: kishorevasan/pyscisci

max_issues_repo_name: kishorevasan/pyscisci

max_forks_repo_name: kishorevasan/pyscisci


text:
import numpy as np
import pandas as pd
import statsmodels.api as sm

import asym_io

def pred_fold(X, coefs):
    return coefs[0] + coefs[1] * X

def linear_fit(X, Y):
    X = sm.add_constant(X)
    mod = sm.OLS(Y, X)
    return mod.fit()

def linear_fit_with_ci(X, Y, ci=0.05):
    res = linear_fit(X, Y)
    fit_ci = np.array(res.conf_int(ci))
    lo = fit_ci[[0,1],[0,1]]
    hi = fit_ci[[0,1],[1,0]]
    return res.params, lo, hi

def get_folding_kinetics(pfdb, ci=0, X='L', Y='log_kf'):
    Xf = np.log10(pfdb[X])
    Yf = pfdb[Y]
    if ci != 0:
        return linear_fit_with_ci(Xf, Yf)
    else:
        return linear_fit(Xf, Yf).params

def get_folding_translation_rates(df, which='best', acpro=False, reduce_pfdb=True, only2s=False):
    if acpro:
        pfdb = asym_io.load_acpro()
        reduce_pfdb = False
    else:
        pfdb = asym_io.load_pfdb()
        if only2s:
            pfdb = pfdb.loc[:88]
    if reduce_pfdb:
        pfdb = pfdb.loc[pfdb.use]
    if which == 'best':
        coef = get_folding_kinetics(pfdb)
    else:
        idx = {'lo':1, 'hi':2}[which]
        coef = get_folding_kinetics(pfdb, ci=0.05)[idx]

    df['ln_kf'] = df.AA_PDB.apply(lambda x: pred_fold(np.log10(x), coef))
    df['T_TRANS'] = np.log10(df.AA_PDB / df.k_trans)
    df['REL_RATE'] = - df['ln_kf'] - df['T_TRANS']
    return df

def kdb_results(pdb, dom, kdb):
    pdb = pdb.copy()
    dom = dom.copy()
    coef = linear_fit(np.log10(kdb['Protein Length']), kdb['ln kf']).params
    pdb['ln_kf'] = pred_fold(np.log10(pdb.AA_PDB), coef)
    pdb = get_rel_rate(pdb)

    dom['ln_kf'] = pred_fold(np.log10(dom.AA_PDB), coef)
    dom = get_rel_rate(dom)

    boot_R_and_save(pdb, kdb='_kdb')
    boot_enrich_and_save(pdb, dom, kdb='_kdb')

max_stars_repo_name: jomimc/FoldAsymCode

max_issues_repo_name: jomimc/FoldAsymCode

max_forks_repo_name: jomimc/FoldAsymCode

reason: import numpy,import statsmodels


text:
## PLOT TEMPORAL AUTOCORRELATIONS

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
exec(open('python/ecco2/colormap.py').read())

#tau1 = np.load('python/gyres/theta_lowres_acftau.npy')
tau1 = np.load('python/gyres/patterns/randpattern_ar5_20-120_high_acftau.npy')
tau2 = np.load('python/gyres/theta_highres_acftau.npy')
(time,lat,lon) = np.load('python/gyres/temp_lowres_dim.npy')

#correct lonna,latna for pcolor
dlon = lon[-1]-lon[-2]
dlat = lat[-1]-lat[-2]
lon = np.hstack((lon,lon[-1]+dlon))-dlon/2.
lat = np.hstack((lat,lat[-1]+dlat))-dlat/2.
xx,yy = np.meshgrid(lon,lat)

## PLOTTING

plt.figure(1,figsize=(10,6))
m = Basemap(llcrnrlon=-2,llcrnrlat=15,urcrnrlon=40,urcrnrlat=58,projection='mill',\
rsphere=6371200.,resolution='l',area_thresh=5000,fix_aspect=False)
m.drawparallels(np.arange(15.,65.,5.),labels=[1,0,0,0],fontsize=10)
m.drawmeridians(np.arange(0.,45.,5.),labels=[0,0,0,1],fontsize=10)
lons, lats = m(xx, yy)

c = m.pcolormesh(lons,lats,tau1)
cc = plt.colorbar(c)
cc.set_label('days')
plt.clim((0,10))
plt.title('T_low: Decorrelation time scale')

plt.figure(2,figsize=(10,6))
m.drawparallels(np.arange(15.,65.,5.),labels=[1,0,0,0],fontsize=10)
m.drawmeridians(np.arange(0.,45.,5.),labels=[0,0,0,1],fontsize=10)
c = m.pcolormesh(lons,lats,tau2)
cc = plt.colorbar(c)
cc.set_label('days')
plt.clim((0,10))
plt.title('T_high: Decorrelation time scale')

plt.show()

max_stars_repo_name: milankl/misc

max_issues_repo_name: milankl/misc

max_forks_repo_name: milankl/misc


text:
from distutils.core import setup, Extension
from Cython.Build import cythonize
import numpy as np

ext = Extension(name='dummyc',
                sources=['dummyc.pyx', 'dummy.c'])

setup(
    ext_modules=cythonize(ext),
    include_dirs=[np.get_include()]
)


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

import re
import numpy as np

data = ''
claims = []

pattern = '#([0-9]*) @ ([0-9]*),([0-9]*): ([0-9]*)x([0-9]*)'

with open('input.txt') as f:
    data = f.read()

claims = [(int(c[0]), int(c[1]), int(c[2]), int(c[3]), int(c[4])) for c in re.findall(pattern, data)]

sq_size = 2000
square = np.zeros((sq_size, sq_size), dtype='int')

for c in claims:
    for r_i in range(c[2], c[2]+c[4]):
        for c_i in range(c[1], c[1]+c[3]):
            square[r_i, c_i] += 1

count_overlap = 0
for r_i in range(sq_size):
    for c_i in range(sq_size):
        if square[r_i, c_i] > 1:
            count_overlap += 1

print(count_overlap)

for c in claims:
    overlaps = False
    for r_i in range(c[2], c[2]+c[4]):
        for c_i in range(c[1], c[1]+c[3]):
            if square[r_i, c_i] > 1:
                overlaps = True
                break
        if overlaps:
            break
    if not(overlaps):
        print c[0]


text:
import matplotlib.pyplot as plt
import numpy as np

from utils import seconds_to_hour

class EDONumericalResolution:
  def __init__(self, edo, ci, full_time_range, time_step):
    """
    @edo: fonction qui retourne [y' = CI, y'' = EDO] sous forme de matrice
    @ci : liste des conditions initiales
    @full_time_range : ndarray contenant [time_start -> time_end] pour chaque time_step
    @time_step : intervalle de temps
    """
    # function that return [y' = CI, y'' = EDO] as a matrix
    self.edo = edo
    self.ci = ci
    self.full_time_range = full_time_range
    self.time_step = time_step

    # keep track of data
    #self._time_set = np.zeros([len(full_time_range)])
    self.y_set = np.zeros([len(full_time_range), len(ci)])

  def resolve(self, **kwargs):
    raise NotImplemented

  def graph(self, title, edo_legends, x_label, y_label):
    full_time_range_in_hours = np.fromiter(map(lambda t: seconds_to_hour(t), self.full_time_range), dtype=np.float)
    plt.plot(full_time_range_in_hours, self.y_set)
    plt.title(title)
    # https://stackoverflow.com/questions/44632571/pyplot-legend-only-displaying-one-letter?noredirect=1&lq=1
    plt.legend(edo_legends, loc="upper right")
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.yscale('log')
    plt.show()

  def graph_sigma_b(self, x_label):
    plt.plot(self.full_time_range, np.array(self.sigma_b_set))
    plt.title("Barres de contrôle")
    # https://stackoverflow.com/questions/44632571/pyplot-legend-only-displaying-one-letter?noredirect=1&lq=1
    plt.legend(["Section efficace des neutrons avec les barres de contrôle"], loc="upper right")
    plt.xlabel(x_label)
    plt.ylabel("Sigma B")
    plt.show()

max_stars_repo_name: Mathieu-R/tchernobyl-simulation

max_issues_repo_name: Mathieu-R/tchernobyl-simulation

max_forks_repo_name: Mathieu-R/tchernobyl-simulation


text:
import os
import sys
import cv2
import time
import click
import numpy as np
import matplotlib.pyplot as plt

# comment out below line to enable tensorflow logging outputs
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# deep sort imports
from deep_sort import preprocessing, nn_matching
from deep_sort.detection import Detection
from deep_sort.tracker import Tracker
from tools import generate_detections as gdet

# YOLO detectors
from detectors import *

@click.command()
@click.option('-f', '--framework', default='tf', type=str, help='Inference framework: {tf, tflite, trt, opencv, openvino}')
@click.option('-m', '--model_path', default='./data/tf/yolov4-416', type=str, help='Path to detection model')
@click.option('-n', '--yolo_names', default='./data/classes/coco.names', type=str, help='Path to YOLO class names file')
@click.option('-s', '--size', default=416, type=int, help='Model input size')
@click.option('-t', '--tiny', default=False, type=bool, help='If YOLO tiny architecture')
@click.option('-t', '--model_type', default='yolov4', type=str, help='yolov3 or yolov4')
@click.option('-v', '--video_path', default='./data/video/test.mp4', type=str, help='Path to input video')
@click.option('-o', '--output', default=None, type=str, help='Path to output, inferenced video')
@click.option('--output_format', default='XVID', type=str, help='Codec used in VideoWriter when saving video to file')
@click.option('--iou', default=0.45, type=float, help='IoU threshold')
@click.option('--score_threshold', default=0.5, type=float, help='Confidence score threshold')
@click.option('--opencv_dnn_target', default='CPU', type=str, help='Precision of OpenCV DNN model')
@click.option('--device', default='MYRIAD', type=str, help='OpenVINO inference device, available: {MYRIAD, CPU, GPU}')
@click.option('--dont_show', default=True, type=bool, help='Do not show video output')
@click.option('--info', default=False, type=bool, help='Show detailed info of tracked objects')
@click.option('--count', default=False, type=bool, help='Count objects being tracked on screen')
def main(framework, model_path, yolo_names, size, tiny, model_type,
         video_path, output, output_format, iou, score_threshold,
         opencv_dnn_target, device, dont_show, info, count):
    # Definition of the parameters
    max_cosine_distance = 0.4
    nn_budget = None
    nms_max_overlap = 1.0

    # initialize deep sort
    model_filename = '/content/DeepSORT-YOLOv4-TensorRT-OpenVINO/data/deep_sort_model/mars-small128.pb'
    encoder = gdet.create_box_encoder(model_filename, batch_size=1)
    # calculate cosine distance metric
    metric = nn_matching.NearestNeighborDistanceMetric("cosine", max_cosine_distance, nn_budget)
    # initialize tracker
    tracker = Tracker(metric)

    # read in all class names from config
    with open(yolo_names, 'r') as f:
        class_names = f.read().split('\n')

    # by default allow all classes in .names file
    allowed_classes = class_names

    # custom allowed classes (uncomment line below to customize tracker for only people)
    # allowed_classes = ['person']

    detector_config = {
        'model_path': model_path,
        'iou_threshold': iou,
        'score_threshold': score_threshold
    }

    # load tf saved model
    if framework == 'tf':
        detector_config['input_size'] = size

        detector = TfYOLO(**detector_config)
    # load tflite model if flag is set
    elif framework == 'tflite':
        detector_config['input_size'] = size
        detector_config['model_type'] = model_type
        detector_config['tiny'] = tiny

        detector = TfliteYOLO(**detector_config)
    # load tensorrt engine if flag is set
    elif framework == 'trt':
        detector = TrtYOLO(**detector_config)
    # load openvino xml model file
    elif framework == 'openvino':
        detector_config['device'] = device

        detector = OpenvinoYOLO(**detector_config)
    # load darknet weights with opencv if flag is set
    elif framework == 'opencv':
        detector_config['cfg_file'] = model_path.replace('weights', 'cfg')
        detector_config['input_size'] = size
        detector_config['opencv_dnn_target'] = opencv_dnn_target

        detector = OpenCVYOLO(**detector_config)

    # begin video capture
    try:
        vid = cv2.VideoCapture(int(video_path))
    except:
        vid = cv2.VideoCapture(video_path)

    # get video ready to save locally if flag is set
    if output:
        # by default VideoCapture returns float instead of int
        width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(vid.get(cv2.CAP_PROP_FPS))
        codec = cv2.VideoWriter_fourcc(*output_format)
        out = cv2.VideoWriter(output, codec, fps, (width, height))

    # detector warm up
    print("{} detector warm up".format(framework))
    for _ in range(5):
        _, _, _, _, _ = detector.detect(np.zeros((2160,3840,3), dtype=np.uint8))

    frame_num = 0
    all_time = 0
    all_detection_time = 0
    all_time_tracker = 0

    # while video is running
    while True:
        return_value, frame = vid.read()
        if return_value:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        else:
            print('Video has ended or failed, try a different video format!')
            break
        frame_num +=1
        print('Frame #: ', frame_num)

        # start frame inference timer
        start_time = time.time()

        # get detections from detector
        bboxes, scores, classes, num_objects, detection_time = detector.detect(frame)
        all_detection_time += detection_time

        # loop through objects and use class index to get class name, allow only classes in allowed_classes list
        names = []
        deleted_indx = []
        for class_indx in classes.astype(int):
            class_name = class_names[class_indx]
            if class_name not in allowed_classes:
                deleted_indx.append(i)
            else:
                names.append(class_name)

        # delete detections that are not in allowed_classes
        bboxes = np.delete(bboxes, deleted_indx, axis=0)
        scores = np.delete(scores, deleted_indx, axis=0)
        names = np.array(names)

        # encode yolo detections and feed to tracker
        features = encoder(frame, bboxes)
        detections = [Detection(bbox, score, class_name, feature) for bbox, score, class_name, feature in zip(bboxes, scores, names, features)]

        #initialize color map
        cmap = plt.get_cmap('tab20b')
        colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]

        # run non-maxima supression
        boxs = np.array([d.tlwh for d in detections])
        scores = np.array([d.confidence for d in detections])
        classes = np.array([d.class_name for d in detections])
        indices = preprocessing.non_max_suppression(boxs, classes, nms_max_overlap, scores)
        detections = [detections[i] for i in indices]

        # call the tracker
        tracker_start = time.time()
        tracker.predict()
        tracker.update(detections)
        all_time_tracker += time.time() - tracker_start

        # update tracks
        for track in tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            bbox = track.to_tlbr()
            class_name = track.get_class()

            # draw bbox on screen
            color = colors[int(track.track_id) % len(colors)]
            color = [i * 255 for i in color]
            cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 2)
            cv2.rectangle(frame, (int(bbox[0]), int(bbox[1]-30)), (int(bbox[0])+(len(class_name)+len(str(track.track_id)))*17, int(bbox[1])), color, -1)
            cv2.putText(frame, class_name + "-" + str(track.track_id),(int(bbox[0]), int(bbox[1]-10)),0, 0.75, (255,255,255),2)

            # if enable info flag then print details about each track
            if info:
                print("Tracker ID: {}, Class: {},  BBox Coords (xmin, ymin, xmax, ymax): {}".format(str(track.track_id), class_name, (int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]))))

        # calculate frames per second of running detections
        loop_time = time.time() - start_time
        all_time += loop_time
        fps = 1.0 / loop_time
        print("FPS: %.2f" % fps)

        if count:
            object_count = len(names)
            cv2.putText(frame, "Objects being tracked: {}".format(object_count), (5, 35), cv2.FONT_HERSHEY_COMPLEX_SMALL, 2, (0, 255, 0), 2)
            print("Objects being tracked: {}".format(object_count))

        result = np.asarray(frame)
        result = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

        if not dont_show:
            cv2.imshow("Output Video", result)

        # if output flag is set, save video file
        if output:
            out.write(result)
        if cv2.waitKey(1) & 0xFF == ord('q'): break
    cv2.destroyAllWindows()

    print("Inference total time: %.4f" % all_time)
    print("Frames: ", frame_num)
    mean_loop_time = all_time / frame_num
    mean_detection_time = all_detection_time / frame_num
    mean_tracker_time = all_time_tracker / frame_num

    mean_fps = frame_num / all_time
    mean_detection_fps = frame_num / all_detection_time
    mean_tracker_fps = frame_num / all_time_tracker

    print("Mean inference time: %.4f" % mean_loop_time)
    print("Mean FPS: %.4f" % mean_fps)
    print("Mean detection time: %.4f" % mean_detection_time)
    print("Mean detection FPS: %.4f" % mean_detection_fps)
    print("Mean tracker time: %.4f" % mean_tracker_time)
    print("Mean tracker FPS: %.4f" % mean_tracker_fps)

if __name__ == '__main__':
    main()