# -*- coding: utf-8 -*-
from utils.ops import unpool, variable_summaries, get_scope_variable, scope, kl_div
import os
import config
import tensorflow as tf
from itertools import combinations, permutations
from tensorflow.python.saved_model import builder as saved_model_builder
import numpy as np
from network import Network

#############################################
#     Adaptive Front and Back End Model     #
#############################################

class Adapt(Network):
	def __init__(self, *args, **kwargs):
		##
		## Model Configuration
		##
		super(Adapt, self).__init__(*args, **kwargs)

		if kwargs is not None:
			self.N = kwargs['filters']
			self.max_pool_value = kwargs['max_pool']
			self.l = kwargs['regularization']
			self.beta = kwargs['beta']
			self.p = kwargs['sparsity']
			self.window = kwargs['window_size']
			self.pretraining = kwargs['pretraining']
			self.overlap_coef = kwargs['overlap_coef']
			self.overlap_value = kwargs['overlap_value']
			self.loss = kwargs['loss']
			self.separation = kwargs['separation']
			self.with_max_pool = kwargs['with_max_pool']
			self.with_average_pool = kwargs['with_average_pool']
			self.hop_size = kwargs['hop_size']
			self.non_negativity = kwargs['non_negativity']

		with tf.get_default_graph().as_default():

			with tf.name_scope('preprocessing'):
				if self.pretraining:
					self.x = tf.concat([self.x_mix,  tf.reshape(self.x_non_mix, [self.B*self.S, self.L])], axis=0)
					self.B_tot = tf.shape(self.x)[0]
				else:

					self.x =tf.concat([self.x_mix,  tf.reshape(self.x_non_mix, [self.B*self.S, self.L])], axis=0)
					self.B_tot = self.B*(self.S+1)

			if self.pretraining:
				self.front
				self.separator
				self.back
				self.cost_model = self.cost
				self.finish_construction()
				self.optimize
			else:
				# self.sepNet = kwargs['separator']
				self.front

	def create_centroids_saver(self):
		self.centroids_saver = tf.train.Saver([self.sepNet.speaker_vectors], max_to_keep=10000000)

	def savedModel(self):
		path = os.path.join(config.log_dir,self.folder ,self.runID, 'SavedModel')
		builder = saved_model_builder.SavedModelBuilder(path)

		# Build signatures
		input_tensor_info = tf.saved_model.utils.build_tensor_info(self.x_mix)
		output_tensor_info = tf.saved_model.utils.build_tensor_info(self.back)

		signature = tf.saved_model.signature_def_utils.build_signature_def(
			inputs={
				'mixed_audio':
				input_tensor_info
			},
			outputs={
				'unmixed_audio':
				output_tensor_info
			}
			)

		builder.add_meta_graph_and_variables(
			self.sess, ['validating'],
			signature_def_map={'separate_audio':signature}
			)

		builder.save()
		print 'Successfully exported model to %s' % path

	##
	## Front End creating STFT like data
	##
	@scope
	def front(self):
		# Front-End

		# x : [ Btot , 1, L , 1]
		# Equivalent to B_tot batches of image of height = 1, width = L and 1 channel -> for Conv1D with Conv2D
		input_front =  tf.reshape(self.x, [self.B_tot, 1, self.L, 1])

		# Filter [filter_height, filter_width, input_channels, output_channels] = [1, W, 1, N]
		self.window_filter = get_scope_variable('window', 'w', shape=[self.window], initializer=tf.contrib.layers.xavier_initializer_conv2d())
		self.bases = get_scope_variable('bases', 'bases', shape=[self.window, self.N], initializer=tf.contrib.layers.xavier_initializer_conv2d())
		self.conv_filter = tf.reshape(tf.abs(tf.expand_dims(self.window_filter,1))*self.bases , [1, self.window, 1, self.N])
		# self.conv_filter = get_scope_variable('filters_front','filters_front', shape=[1, self.window, 1, self.N])
		variable_summaries(self.conv_filter)

		# 1 Dimensional convolution along T axis with a window length = self.window
		# And N = 256 filters -> Create a [Btot, 1, T, N]
		self.T = tf.shape(input_front)[2]

		if self.with_max_pool:
			self.X = tf.nn.conv2d(input_front, self.conv_filter, strides=[1, 1, 1, 1], padding="SAME", name='Conv_STFT')
			self.y, self.argmax = tf.nn.max_pool_with_argmax(self.X, [1, 1, self.max_pool_value, 1],
										strides=[1, 1, self.hop_size, 1], padding="VALID", name='output')
		elif self.with_average_pool:
			self.X = tf.nn.conv2d(input_front, self.conv_filter, strides=[1, 1, 1, 1], padding="SAME", name='Conv_STFT')
			self.y = tf.layers.average_pooling2d(self.X, (1, self.max_pool_value), strides=(1, self.max_pool_value), name='output')
		else:
			self.y = tf.nn.conv2d(input_front, self.conv_filter, strides=[1, 1, self.hop_size, 1], padding="SAME", name='Conv_STFT')

		# Reshape to Btot batches of T x N images with 1 channel
		# [Btot, 1, T_pool, N] -> [Btot, T-pool, N, 1]
		self.y = tf.transpose(self.y,[0, 2, 3, 1], name='output')

		tf.summary.image('front/output', self.y, max_outputs=3)
		y_shape = tf.shape(self.y)
		y = tf.reshape(self.y, [self.B_tot, y_shape[1]*y_shape[2]])
		self.p_hat = tf.reduce_sum(tf.abs(y), 0)
		self.sparse_constraint = tf.reduce_sum(kl_div(self.p, self.p_hat))

		return self.y

	@scope
	def separator(self):
		# shape = [B_tot, T_, N, 1], shape = [B(1+S), T , N, 1], [B(1+S), T_ , N, 1]
		separator_in = self.front
		# Compute the overlapping rate:
		nb = 2
		comb = list(combinations(range(self.S), nb))
		len_comb = len(comb)
		combs = tf.reshape(tf.constant(comb), [1, len_comb, 2, 1])
		combs = tf.tile(combs, [self.B, 1, 1, 1])
		batch_range = tf.tile(tf.reshape(tf.range(self.B, dtype=tf.int32), shape=[self.B, 1, 1, 1]), [1, len_comb, nb, 1])
		comb_range = tf.tile(tf.reshape(tf.range(len_comb, dtype=tf.int32), shape=[1, len_comb, 1, 1]), [self.B, 1, nb, 1])
		indicies = tf.concat([batch_range, comb_range, combs], axis=3)
		self.T_max_pooled = tf.shape(separator_in)[1]
		input = tf.reshape(separator_in, [self.B_tot, self.T_max_pooled, self.N])
		input_non_mix = tf.reshape(input[self.B:, : , :], [self.B, self.S, self.T_max_pooled, self.N]) # B*S others non mix
		comb_non_mix = tf.gather_nd(tf.tile(tf.reshape(input_non_mix, [self.B, 1, self.S, self.T_max_pooled*self.N]), [1, len_comb, 1, 1]), indicies) #

		# Combination of non mixed representations : [B, len(comb), nb, T*N]
		comb_non_mix = tf.abs(comb_non_mix)
		measure = 1.0 - tf.abs(comb_non_mix[:,:,0,:] - comb_non_mix[:,:,1,:]) / (tf.reduce_max(comb_non_mix, axis=2) + 1e-8)
		overlapping = tf.reduce_mean(measure, -1) # Mean over the bins
		overlapping = tf.reduce_mean(overlapping, -1) # Mean over combinations
		self.overlapping = tf.reduce_mean(overlapping, -1) # Mean over batces
		self.overlapping_constraint = self.overlapping

		input_mix = tf.reshape(input[:self.B, : , :], [self.B, 1, self.T_max_pooled, self.N]) # B first batches correspond to mix input
		input_non_mix = tf.reshape(input[self.B:, : , :], [self.B, self.S, self.T_max_pooled, self.N]) # B*S others non mix

		input_mix = tf.tile(input_mix, [1, self.S, 1, 1])

		filters = tf.divide(input_non_mix, input_mix)
		output = tf.reshape(input_mix * filters, [self.B*self.S, self.T_max_pooled, self.N, 1])
		output = tf.transpose(output, [0, 3, 1, 2])

		self.filters = filters

		if self.pretraining:
			self.T_max_pooled = tf.shape(separator_in)[1]

			input = tf.reshape(separator_in, [self.B_tot, self.T_max_pooled, self.N])

			if self.separation == 'mask':
				filters = tf.divide(input_non_mix, input_mix)
				output = tf.reshape(input_mix * filters, [self.B*self.S, self.T_max_pooled, self.N, 1])
				output = tf.transpose(output, [0, 3, 1, 2])

				self.filters = filters

			elif self.separation == 'perfect':
				# From [a, b, c ,d] -> [a+b+c+d, a+b+c+d, a+b+c+d, a+b+c+d]
				tiled_sum = tf.tile(tf.reduce_sum(input_non_mix, 1, keep_dims=True), [1, self.S, 1, 1])
				# From [a+b+c+d, a+b+c+d, a+b+c+d, a+b+c+d] ->  [b+c+d, a+c+d, a+b+d, a+b+c]
				X_add = tiled_sum - input_non_mix
				output = input_mix - X_add

				# with tf.control_dependencies([tf.assert_equal(output,input_non_mix)]):
				output = tf.reshape(output, [self.B*self.S, self.T_max_pooled, self.N, 1])

			return output
		else:
			output = self.sepNet.output
			self.prediction = self.sepNet.prediction
			self.true_masks = self.sepNet.true_masks
			self.X_non_mix = self.sepNet.X_non_mix
			self.X = self.sepNet.X
			return output

	@scope
	def back(self):

		input_tensor = self.separator

		if self.with_max_pool:

			argmax_mix = self.argmax[:self.B]

			repeats = [self.S, 1, 1 ,1]
			shape = tf.shape(argmax_mix)
			argmax_mix = tf.expand_dims(argmax_mix, 1)
			argmax_mix = tf.tile(argmax_mix, [1, self.S, 1, 1 ,1])
			argmax_mix = tf.reshape(argmax_mix, shape*repeats)

			output_shape = [self.B*self.S, 1, self.T, self.N]

			output = unpool(input_tensor, argmax_mix, ksize=[1, 1, self.max_pool_value, 1], output_shape=output_shape, scope='unpool')
			output = tf.reshape(output, [self.B*self.S, 1, self.T, self.N])

		elif self.with_average_pool:
			input_tensor = tf.reshape(input_tensor, [self.B*self.S, 1, self.T_max_pooled, self.N])
			output = tf.keras.layers.UpSampling2D((1, self.max_pool_value))(input_tensor)
			output = tf.reshape(output, [self.B*self.S, 1, self.T, self.N])
		else:
			output = tf.reshape(input_tensor, [self.B*self.S, 1, self.T_max_pooled, self.N])

		self.window_filter_2 = get_scope_variable('window', 'value', shape=[self.window], initializer=tf.contrib.layers.xavier_initializer_conv2d())
		self.bases_2 = get_scope_variable('bases', 'value', shape=[self.window, self.N], initializer=tf.contrib.layers.xavier_initializer_conv2d())
		self.conv_filter_2 = tf.reshape(tf.abs(tf.expand_dims(self.window_filter_2,1))*self.bases_2 , [1, self.window, 1, self.N], name='filters')# self.conv_filter_2 = tf.Variable(self.conv_filter.initialized_value(), name="filters_back")

		if not (self.with_max_pool or self.with_average_pool):
			strides = [1, 1, self.hop_size, 1]
		else:
			strides = [1, 1, 1, 1]

		output = tf.nn.conv2d_transpose(output , filter=self.conv_filter_2,
									 output_shape=[self.B*self.S, 1, self.L, 1],
									 strides=strides, padding='SAME')

		output = tf.reshape(output, [self.B, self.S, self.L], name='back_output')
		self.output = output

		sdr_improvement, sdr = self.sdr_improvement(self.x_non_mix, output)
		self.sdr_imp = sdr_improvement

		return output

	@scope
	def enhance(self):

		input_enhance = self.output # Take the reconstructed audio as input
		# shape [B, S, L]

		x_tile  = tf.tile(tf.expand_dims(self.x_mix, 1), [1, self.S, 1]) # [B, L] -> [B, S, L]

		z = tf.concat([tf.expand_dims(x_tile, 2), tf.expand_dims(input_enhance, 2)], axis=2)

		z = tf.reshape(z, [self.B*self.S, 2, self.L, 1]) # B*S images of [2, L] with 1 channel

		z = tf.layers.conv2d(z, 32, [1, 8], activation=tf.nn.relu, padding='same') # [2,L, 256]
		print z
		z = tf.layers.conv2d(z, 32, [2, 8], activation=tf.nn.relu, padding='same') # [2, L, 256]
		print z
		z = tf.layers.conv2d(z, 1, [1, 8], activation=tf.nn.relu, padding='same') # [2, L, 256]
		print z
		# z = tf.layers.conv2d(z, 1, [1, 1024], activation=tf.nn.relu) # [2, L, 1]
		z = tf.layers.average_pooling2d(z, [2, 1], strides=[1, 1])
		z = tf.reshape(z, [self.B, self.S, self.L])

		return z

	@scope
	def enhance_cost(self):
		perms = list(permutations(range(self.S))) # ex with 3: [0, 1, 2], [0, 2 ,1], [1, 0, 2], [1, 2, 0], [2, 1, 0], [2, 0, 1]
		length_perm = len(perms)
		perms = tf.reshape(tf.constant(perms), [1, length_perm, self.S, 1])
		perms = tf.tile(perms, [self.B, 1, 1, 1])

		batch_range = tf.tile(tf.reshape(tf.range(self.B, dtype=tf.int32), shape=[self.B, 1, 1, 1]), [1, length_perm, self.S, 1])
		perm_range = tf.tile(tf.reshape(tf.range(length_perm, dtype=tf.int32), shape=[1, length_perm, 1, 1]), [self.B, 1, self.S, 1])
		indicies = tf.concat([batch_range, perm_range, perms], axis=3)

		# [B, P, S, L]
		permuted_back = tf.gather_nd(tf.tile(tf.reshape(self.enhance, [self.B, 1, self.S, self.L]), [1, length_perm, 1, 1]), indicies) #

		X_nmr = tf.reshape(self.x_non_mix, [self.B, 1, self.S, self.L])

		l2 = tf.reduce_sum(tf.square(X_nmr - permuted_back), axis=-1) # L2^2 norm
		l2 = tf.reduce_min(l2, axis=1) # Get the minimum over all possible permutations : B S
		l2 = tf.reduce_sum(l2, -1)
		l2 = tf.reduce_mean(l2, -1)

		sdr_improvement, sdr = self.sdr_improvement(X_nmr, self.enhance, True)
		self.sdr_imp = sdr_improvement

		tf.summary.scalar('SDR_improvement', sdr_improvement)
		tf.summary.scalar('cost', l2)

		return l2

	@scope
	def cost(self):
		# Definition of cost for Adapt model
		# Regularisation
		# shape = [B_tot, T, N]
		regularization = self.l * (tf.nn.l2_loss(self.conv_filter_2)+ tf.nn.l2_loss(self.conv_filter))

		# Non-Negativity Loss
		neg = tf.square(tf.where(self.front < 0, self.front, tf.zeros_like(self.front)))
		nn = self.non_negativity*tf.reduce_mean(tf.reduce_sum(neg, [1,2,3]))

		# input_shape = [B, S, L]
		# Doing l2 norm on L axis :
		if self.pretraining:

			l2 = tf.reduce_sum(tf.square(self.x_non_mix - self.back), axis=-1)
			l2 = tf.reduce_sum(l2, -1) # Sum over all the speakers
			l2 = tf.reduce_mean(l2, -1) # Mean over batches

			sdr_improvement, sdr = self.sdr_improvement(self.x_non_mix, self.back)
			self.sdr_imp = sdr_improvement
			sdr = tf.reduce_mean(sdr) # Mean over speakers
			sdr = tf.reduce_mean(sdr) # Mean over batches

			if self.loss == 'l2':
				loss = l2
			elif self.loss == 'sdr':
				loss = sdr
			else:
				loss = l2 + sdr

		else:
			# Compute loss over all possible permutations

			perms = list(permutations(range(self.S))) # ex with 3: [0, 1, 2], [0, 2 ,1], [1, 0, 2], [1, 2, 0], [2, 1, 0], [2, 0, 1]
			length_perm = len(perms)
			perms = tf.reshape(tf.constant(perms), [1, length_perm, self.S, 1])
			perms = tf.tile(perms, [self.B, 1, 1, 1])

			batch_range = tf.tile(tf.reshape(tf.range(self.B, dtype=tf.int32), shape=[self.B, 1, 1, 1]), [1, length_perm, self.S, 1])
			perm_range = tf.tile(tf.reshape(tf.range(length_perm, dtype=tf.int32), shape=[1, length_perm, 1, 1]), [self.B, 1, self.S, 1])
			indicies = tf.concat([batch_range, perm_range, perms], axis=3)

			# [B, P, S, L]
			permuted_back = tf.gather_nd(tf.tile(tf.reshape(self.back, [self.B, 1, self.S, self.L]), [1, length_perm, 1, 1]), indicies) #

			X_nmr = tf.reshape(self.x_non_mix, [self.B, 1, self.S, self.L])

			l2 = tf.reduce_mean(tf.square(X_nmr - permuted_back), axis=-1) # MSE norm
			l2 = tf.reduce_sum(l2, -1)
			l2 = tf.reduce_min(l2,  -1) # Get the minimum over all possible permutations : B
			l2 = tf.reduce_mean(l2, -1)

			sdr_improvement, sdr = self.sdr_improvement(X_nmr, self.back, True)
			self.sdr_imp = sdr_improvement
			sdr = tf.reduce_min(sdr, 1) # Get the minimum over all possible permutations : B S
			sdr = tf.reduce_sum(sdr, -1)
			sdr = tf.reduce_mean(sdr, -1)

			if self.loss == 'l2':
				loss = l2
			elif self.loss == 'sdr':
				loss = sdr
			else:
				loss = 1e-3*l2 + sdr

		# shape = [B]
		# Compute mean over batches
		cost_value = loss
		if self.beta != 0.0:
			cost_value += self.beta * self.sparse_constraint
		if self.l != 0.0:
			cost_value += self.l * regularization
		if self.overlap_coef != 0.0:
			cost_value += self.overlap_coef * self.overlapping_constraint
		if self.non_negativity is not None:
			cost_value += self.non_negativity*nn

		variable_summaries(self.conv_filter_2)

		tf.summary.audio(name= "audio/output/reconstructed", tensor = tf.reshape(self.back, [-1, self.L]), sample_rate = config.fs, max_outputs=2)

		with tf.name_scope('loss_values'):
			tf.summary.scalar('l2_loss', l2)
			tf.summary.scalar('SDR', sdr)
			tf.summary.scalar('SDR_improvement', sdr_improvement)
			tf.summary.scalar('Non negativity loss', self.non_negativity*nn)
			tf.summary.scalar('sparsity', tf.reduce_mean(self.p_hat))
			tf.summary.scalar('sparsity_loss', self.beta * self.sparse_constraint)
			tf.summary.scalar('L2_reg', self.l * regularization)
			tf.summary.scalar('loss', cost_value)
			tf.summary.scalar('overlapping', self.overlapping)
			tf.summary.scalar('overlapping_loss', self.overlap_coef * self.overlapping_constraint)

		return cost_value

	@scope
	def cost_finetuning(self):
		perms = list(permutations(range(self.S))) # ex with 3: [0, 1, 2], [0, 2 ,1], [1, 0, 2], [1, 2, 0], [2, 1, 0], [2, 0, 1]
		length_perm = len(perms)
		perms = tf.reshape(tf.constant(perms), [1, length_perm, self.S, 1])
		perms = tf.tile(perms, [self.B, 1, 1, 1])

		batch_range = tf.tile(tf.reshape(tf.range(self.B, dtype=tf.int32), shape=[self.B, 1, 1, 1]), [1, length_perm, self.S, 1])
		perm_range = tf.tile(tf.reshape(tf.range(length_perm, dtype=tf.int32), shape=[1, length_perm, 1, 1]), [self.B, 1, self.S, 1])
		indicies = tf.concat([batch_range, perm_range, perms], axis=3)

		postproc = tf.tile(tf.reshape(self.back, [self.B, 1, self.S, self.L]), [1, length_perm, 1, 1])

		# [B, P, S, L]
		permuted_back = tf.gather_nd(postproc, indicies) #

		X_nmr = tf.reshape(self.x_non_mix, [self.B, 1, self.S, self.L])

		l2 = 0.5*tf.reduce_sum(tf.square(X_nmr - permuted_back), axis=-1) # L2^2 norm
		l2 = tf.reduce_mean(l2, -1)
		l2 = tf.reduce_min(l2, -1) # Get the minimum over all possible permutations
		l2 = tf.reduce_mean(l2, -1)
		tf.summary.scalar('cost', l2)

		sdr_improvement, sdr = self.sdr_improvement(X_nmr, self.back, True)
		tf.summary.scalar('sdr', sdr_improvement)

		return l2

	def test_prediction(self, X_mix_test, X_non_mix_test, step):
		pred, y = self.sess.run([self.sepNet.prediction, self.sepNet.y_test_export], {self.x_mix: X_mix_test, self.x_non_mix:X_non_mix_test, self.training:True})
		pred = np.reshape(pred, [X_mix_test.shape[0], -1, 40])
		labels = [['r' if b == 1 else 'b' for b in batch]for batch in y]
		np.save(os.path.join(config.log_dir, self.folder ,self.runID, "bins-{}".format(step)), pred)
		np.save(os.path.join(config.log_dir, self.folder ,self.runID, "labels-{}".format(step)), labels)

	def connect_front(self, separator_class):
		self.sepNet = separator_class(True, **self.args)

	def connect_only_front_to_separator(self, separator, freeze_front=True):
		self.connect_front(separator)
		self.sepNet.output = self.sepNet.prediction
		self.cost_model = self.sepNet.cost
		self.back # To save the back values !
		var_list =[v for v in tf.global_variables() if ('back/' in v.name or 'front/' in v.name)]
		self.create_saver(subset=var_list)
		self.restore_model(self.args['model_folder'])
		self.finish_construction()
		self.freeze_all_with('front/')
		self.freeze_all_with('back/')
		self.optimize
		self.tensorboard_init()

	def connect_enhance_to_separator(self, separator):
		self.connect_front(separator)
		self.sepNet.output = self.sepNet.enhance
		self.cost_model = self.sepNet.enhance_cost
		self.back # To save the back values !
		var_list =[v for v in tf.global_variables()
			if ('back/' in v.name or 'front/' in v.name or 'prediction/' in v.name or 'speaker_centroids' in v.name )]
		self.create_saver(subset=var_list)
		self.restore_model(self.args['model_folder'])
		self.finish_construction()
		self.freeze_all_except('enhance/')
		self.optimize
		self.tensorboard_init()

	def restore_front_separator(self, sess, path, separator):
		self.connect_front(separator)
		self.sepNet.output = self.sepNet.prediction
		self.back
		self.restore_model(sess, path)


import cv2
import numpy as np
import dlib
from math import hypot

cap = cv2.VideoCapture(0)

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")

def midpoint(p1 ,p2):
    return int((p1.x + p2.x)/2), int((p1.y + p2.y)/2)

font = cv2.FONT_HERSHEY_PLAIN

def get_blinking_ratio(eye_points, facial_landmarks):
    left_point = (facial_landmarks.part(eye_points[0]).x, facial_landmarks.part(eye_points[0]).y)
    right_point = (facial_landmarks.part(eye_points[3]).x, facial_landmarks.part(eye_points[3]).y)
    center_top = midpoint(facial_landmarks.part(eye_points[1]), facial_landmarks.part(eye_points[2]))
    center_bottom = midpoint(facial_landmarks.part(eye_points[5]), facial_landmarks.part(eye_points[4]))

    hor_line = cv2.line(frame, left_point, right_point, (0, 255, 0), 2)
    ver_line = cv2.line(frame, center_top, center_bottom, (0, 255, 0), 2)

    hor_line_lenght = hypot((left_point[0] - right_point[0]), (left_point[1] - right_point[1]))
    ver_line_lenght = hypot((center_top[0] - center_bottom[0]), (center_top[1] - center_bottom[1]))

    ratio = hor_line_lenght / ver_line_lenght
    return ratio

while True:
    _, frame = cap.read()
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    faces = detector(gray)
    for face in faces:
        #x, y = face.left(), face.top()
        #x1, y1 = face.right(), face.bottom()
        #cv2.rectangle(frame, (x, y), (x1, y1), (0, 255, 0), 2)

        landmarks = predictor(gray, face)

        left_eye_ratio = get_blinking_ratio([36, 37, 38, 39, 40, 41], landmarks)
        right_eye_ratio = get_blinking_ratio([42, 43, 44, 45, 46, 47], landmarks)
        blinking_ratio = (left_eye_ratio + right_eye_ratio) / 2

        if blinking_ratio > 5.7:
            cv2.putText(frame, "BLINKING", (50, 150), font, 7, (255, 0, 0))

    cv2.imshow("Frame", frame)

    key = cv2.waitKey(1)
    if key == 27:
        break

cap.release()
cv2.destroyAllWindows()


# Convert uvfits from gildas into a CASA MS
# For casa 5.1.1

from astropy.io import fits
import astropy.units as u
from astropy.coordinates import SkyCoord
import numpy as np
import os

# Issues with setting vsys and rest freq of the source
# Presumably this is fine and can just be given when imaging
# The dishes are also all set to the same position, but the baselines
# are correct in the actual data.

# Loop through the different spectral resolutions Jonathan made
# str_prefix = ['M33-ARM-', 'M33-ARM05-', 'M33-ARM1-', 'M33-ARM13-', 'M33-ARM2-',
#               'M33-ARMcont-', 'M33-ARM05-merged-']
str_prefix = ['M33-ARM-', 'M33-ARM05-', 'M33-ARM1-', 'M33-ARM13-', 'M33-ARM2-',
              'M33-ARMcont-']

if not os.path.exists('meas_sets'):
    os.mkdir('meas_sets')

for pref in str_prefix:

    print("On {}".format(pref))

    ms_names = []

    for i in range(1, 89):

        print("On {} of 88".format(i))

        str_name = '{}{}'.format(pref, i)

        uvfits = "calib/uvfits/{}.uvfits".format(str_name)
        uncor_vis_name = "meas_sets/{}_uncor.ms".format(str_name)
        vis_name = "meas_sets/{}.ms".format(str_name)

        hdu = fits.open(uvfits, mode='update')
        # UVFITS have two directions: CRVAL5 and 6 are the phase direction,
        # and OBSRA and DEC are the pointing centres. But CASA doesn't
        # recognize the latter keywords and assumes that they are the pointing
        # and phase centres. We'll fix these below using the fixvis routine,
        # but we'll need the pointing centres
        point = SkyCoord(hdu[0].header['OBSRA'] * u.deg,
                         hdu[0].header['OBSDEC'] * u.deg)
        point_string = point.to_string(style='hmsdms')

        # hdu[0].header['CRVAL5'] = hdu[0].header['OBSRA']
        # hdu[0].header['CRVAL6'] = hdu[0].header['OBSDEC']

        # The ref freq is also confusing. It's the frequency in the rest frame
        # So it needs to be altered to be in the observed frame.
        # GILDAS appears to have this conversion built-in everywhere
        restfreq = hdu[0].header['RESTFREQ'] * u.Hz
        rest_cent_freq = hdu[0].header['CRVAL4'] * u.Hz
        source_vel = hdu[0].header['VELO-LSR'] * u.m / u.s

        # Calculate the frequency shift
        del_f = rest_cent_freq - source_vel.to(u.Hz, u.doppler_radio(restfreq))

        # Adjust CRVAL4 to the observed frame frequency
        hdu[0].header['CRVAL4'] -= del_f.value

        # Need to assign the SD visibilities to an ant pair
        # CASA doesn't barf when doing this, but I can't be sure it worked in
        # any useful way. So don't use this for any science products, just as
        # a check against the GILDAS imaging
        if 'merged' in pref:
            raise ValueError("This does not work. Don't import merged files "
                             "into CASA!")
            baselines = hdu[0].data['BASELINE']
            baselines[np.where(baselines == 0)] = baselines[0]
            hdu[0].data['BASELINE'] = baselines

        hdu.flush()
        hdu.close()

        importuvfits(fitsfile=uvfits, vis=uncor_vis_name)

        # Recalculate uv data with the phase at the pointing centre
        fixvis(vis=uncor_vis_name,
               outputvis=vis_name,
               field='M33-ARM',
               phasecenter='J2000 {}'.format(point_string))

        # Remove the uncorrected version
        rmtables(uncor_vis_name)

        # Change the field name
        vishead(vis="meas_sets/{}.ms".format(str_name), mode='put',
                hdkey='field', hdvalue=str_name)

        # If you want to use tclean, you'll need to rename the telescope to
        # IRAMPDB instead of IRAM PDB. The former is the correct beam response
        # name, the latter is the name that CASA recognizes elsewhere.
        vishead(vis='meas_sets/{}.ms'.format(str_name), mode='put',
                hdkey='telescope',
                hdvalue=(np.array(["IRAMPDB"] * 2)))

        # Add in rest freq and vsys
        hdu = fits.open(uvfits)

        tb.open("{}/SOURCE".format(vis_name), nomodify=False)
        # tb.putcol('REST_FREQUENCY', [[hdu[0].header['RESTFREQ']]])
        tb.putcol('SYSVEL', [[hdu[0].header['VELO-LSR']]])
        tb.putcol('NAME', [[str_name]])
        tb.close()

        hdu.close()

        ms_names.append(vis_name)

    # Now concatenate

    concat(vis=ms_names, concatvis='meas_sets/{}.ms'.format(pref[:-1]))

    # Delete the individual MSs
    for ms_name in ms_names:
        rmtables(ms_name)

import numpy,import astropy,from astropy


import argparse
import os
from time import time
from PIL import Image
import align.detect_face as detect_face
import cv2
import numpy as np
import tensorflow as tf
from lib.face_utils import judge_side_face
from lib.utils import Logger, mkdir
from project_root_dir import project_dir
from src.sort import Sort
from load_model.tensorflow_loader import load_tf_model, tf_inference
from utils.anchor_generator import generate_anchors
from utils.anchor_decode import decode_bbox
from utils.nms import single_class_non_max_suppression

logger = Logger()

# 开始人脸对齐

sess, graph = load_tf_model('models/face_mask_detection.pb')
# anchor configuration
feature_map_sizes = [[33, 33], [17, 17], [9, 9], [5, 5], [3, 3]]
anchor_sizes = [[0.04, 0.056], [0.08, 0.11], [0.16, 0.22], [0.32, 0.45], [0.64, 0.72]]
anchor_ratios = [[1, 0.62, 0.42]] * 5

# generate anchors
anchors = generate_anchors(feature_map_sizes, anchor_sizes, anchor_ratios)

# for inference , the batch size is 1, the model output shape is [1, N, 4],
# so we expand dim for anchors to [1, anchor_num, 4]
anchors_exp = np.expand_dims(anchors, axis=0)

id2class = {0: 'Mask', 1: 'NoMask'}

# 人脸对齐方法
def inference(image,
              conf_thresh=0.5,
              iou_thresh=0.4,
              target_shape=(160, 160),
              draw_result=True,
              show_result=True
              ):
    '''
    Main function of detection inference
    :param image: 3D numpy array of image
    :param conf_thresh: the min threshold of classification probabity.
    :param iou_thresh: the IOU threshold of NMS
    :param target_shape: the model input size.
    :param draw_result: whether to daw bounding box to the image.
    :param show_result: whether to display the image.
    :return:
    '''

    output_info = []
    height, width, _ = image.shape
    image_resized = cv2.resize(image, target_shape)
    image_np = image_resized / 255.0  # 归一化到0~1
    image_exp = np.expand_dims(image_np, axis=0)

    # 输出回归框和人脸得分
    y_bboxes_output, y_cls_output = tf_inference(sess, graph, image_exp)

    # remove the batch dimension, for batch is always 1 for inference.
    y_bboxes = decode_bbox(anchors_exp, y_bboxes_output)[0]
    y_cls = y_cls_output[0]
    # To speed up, do single class NMS, not multiple classes NMS.
    bbox_max_scores = np.max(y_cls, axis=1)
    bbox_max_score_classes = np.argmax(y_cls, axis=1)

    # 对单个人脸进行非极大抑制
    keep_idxs = single_class_non_max_suppression(y_bboxes,
                                                 bbox_max_scores,
                                                 conf_thresh=conf_thresh,
                                                 iou_thresh=iou_thresh,
                                                 )

    for idx in keep_idxs:
        conf = float(bbox_max_scores[idx])
        class_id = bbox_max_score_classes[idx]
        bbox = y_bboxes[idx]
        # clip the coordinate, avoid the value exceed the image boundary.
        xmin = max(0, int(bbox[0] * width))
        ymin = max(0, int(bbox[1] * height))
        xmax = min(int(bbox[2] * width), width)
        ymax = min(int(bbox[3] * height), height)

        if draw_result:
            if class_id == 0:
                color = (0, 255, 0)
            else:
                color = (255, 0, 0)
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
            cv2.putText(image, "%s: %.2f" % (id2class[class_id], conf), (xmin + 2, ymin - 2),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, color)
        output_info.append([class_id, conf, xmin, ymin, xmax, ymax])
        print(output_info)
    if show_result:
        Image.fromarray(image).show()

    return output_info

def main():
    global colours, img_size
    args = parse_args()
    videos_dir = args.videos_dir
    output_path = args.output_path
    no_display = args.no_display
    detect_interval = args.detect_interval  # 间隔一帧检测一次
    margin = args.margin  # 脸边距（默认10）
    scale_rate = args.scale_rate  # 检测图像的尺寸设置
    show_rate = args.show_rate  # 展示图像的尺寸设置
    face_score_threshold = args.face_score_threshold  # 人脸判别阈值

    mkdir(output_path)
    # for display
    if not no_display:
        colours = np.random.rand(32, 3)

    # 初始化追踪器
    tracker = Sort()  # create instance of the SORT tracker

    logger.info('Start track and extract......')

    # 影像处理
    for filename in os.listdir(videos_dir):
        logger.info('All files:{}'.format(filename))
    for filename in os.listdir(videos_dir):
        suffix = filename.split('.')[1]
        if suffix != 'mp4' and suffix != 'avi':  # you can specify more video formats if you need
            continue
        video_name = os.path.join(videos_dir, filename)
        directoryname = os.path.join(output_path, filename.split('.')[0])
        logger.info('Video_name:{}'.format(video_name))
        cam = cv2.VideoCapture(video_name)
        c = 0
        while True:
            final_faces = []
            addtional_attribute_list = []
            ret, frame = cam.read()
            if not ret:
                logger.warning("ret false")
                break
            if frame is None:
                logger.warning("frame drop")
                break

            frame = cv2.resize(frame, (0, 0), fx=scale_rate, fy=scale_rate)
            r_g_b_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            # 间隔取帧，默认每帧都取
            # if c % detect_interval == 0:
            #     img_size = np.asarray(frame.shape)[0:2]
            #     faces = inference(r_g_b_frame, show_result=True, target_shape=(260, 260))

    with tf.Graph().as_default():
        with tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True),
                                              log_device_placement=False)) as sess:

            pnet, rnet, onet = detect_face.create_mtcnn(sess, os.path.join(project_dir, "align"))

            minsize = 40  # minimum size of face for mtcnn to detect
            threshold = [0.6, 0.7, 0.7]  # three steps's threshold
            factor = 0.709  # scale factor

            for filename in os.listdir(videos_dir):
                logger.info('All files:{}'.format(filename))
            for filename in os.listdir(videos_dir):
                suffix = filename.split('.')[1]
                if suffix != 'mp4' and suffix != 'avi':  # you can specify more video formats if you need
                    continue
                video_name = os.path.join(videos_dir, filename)
                directoryname = os.path.join(output_path, filename.split('.')[0])
                logger.info('Video_name:{}'.format(video_name))
                cam = cv2.VideoCapture(video_name)
                c = 0
                while True:
                    final_faces = []
                    addtional_attribute_list = []
                    ret, frame = cam.read()
                    if not ret:
                        logger.warning("ret false")
                        break
                    if frame is None:
                        logger.warning("frame drop")
                        break

                    frame = cv2.resize(frame, (0, 0), fx=scale_rate, fy=scale_rate)
                    r_g_b_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    if c % detect_interval == 0:
                        img_size = np.asarray(frame.shape)[0:2]
                        mtcnn_starttime = time()
                        faces, points = detect_face.detect_face(r_g_b_frame, minsize, pnet, rnet, onet, threshold,
                                       factor)
                        logger.info("MTCNN detect face cost time : {} s".format(
                            round(time() - mtcnn_starttime, 3)))  # mtcnn detect ,slow
                        print('testttttt')
                        print(type(faces))
                        face_sums = faces.shape[0]
                        if face_sums > 0:
                            face_list = []
                            for i, item in enumerate(faces):
                                print(item)
                                score = round(faces[i, 4], 6)
                                if score > face_score_threshold:
                                    det = np.squeeze(faces[i, 0:4])

                                    # face rectangle
                                    det[0] = np.maximum(det[0] - margin, 0)
                                    det[1] = np.maximum(det[1] - margin, 0)
                                    det[2] = np.minimum(det[2] + margin, img_size[1])
                                    det[3] = np.minimum(det[3] + margin, img_size[0])
                                    face_list.append(item)

                                    # face cropped
                                    bb = np.array(det, dtype=np.int32)

                                    # use 5 face landmarks  to judge the face is front or side
                                    squeeze_points = np.squeeze(points[:, i])
                                    tolist = squeeze_points.tolist()
                                    facial_landmarks = []
                                    for j in range(5):
                                        item = [tolist[j], tolist[(j + 5)]]
                                        facial_landmarks.append(item)
                                    if args.face_landmarks:
                                        for (x, y) in facial_landmarks:
                                            cv2.circle(frame, (int(x), int(y)), 3, (0, 255, 0), -1)
                                    cropped = frame[bb[1]:bb[3], bb[0]:bb[2], :].copy()

                                    dist_rate, high_ratio_variance, width_rate = judge_side_face(
                                        np.array(facial_landmarks))

                                    # face addtional attribute(index 0:face score; index 1:0 represents front face and 1 for side face )
                                    item_list = [cropped, score, dist_rate, high_ratio_variance, width_rate]
                                    addtional_attribute_list.append(item_list)

                            final_faces = np.array(face_list)

                    trackers = tracker.update(final_faces, img_size, directoryname, addtional_attribute_list, detect_interval)

                    c += 1

                    for d in trackers:
                        if not no_display:
                            d = d.astype(np.int32)
                            cv2.rectangle(frame, (d[0], d[1]), (d[2], d[3]), colours[d[4] % 32, :] * 255, 3)
                            if final_faces != []:
                                cv2.putText(frame, 'ID : %d  DETECT' % (d[4]), (d[0] - 10, d[1] - 10),
                                            cv2.FONT_HERSHEY_SIMPLEX,
                                            0.75,
                                            colours[d[4] % 32, :] * 255, 2)
                                cv2.putText(frame, 'DETECTOR', (5, 45), cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                                            (1, 1, 1), 2)
                            else:
                                cv2.putText(frame, 'ID : %d' % (d[4]), (d[0] - 10, d[1] - 10), cv2.FONT_HERSHEY_SIMPLEX,
                                            0.75,
                                            colours[d[4] % 32, :] * 255, 2)

                    if not no_display:
                        frame = cv2.resize(frame, (0, 0), fx=show_rate, fy=show_rate)
                        cv2.imshow("Frame", frame)
                        if cv2.waitKey(1) & 0xFF == ord('q'):
                            break

def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--videos_dir", type=str,
                        help='Path to the data directory containing aligned your face patches.', default='videos')
    parser.add_argument('--output_path', type=str,
                        help='Path to save face',
                        default='facepics')
    parser.add_argument('--detect_interval',
                        help='how many frames to make a detection',
                        type=int, default=1)
    parser.add_argument('--margin',
                        help='add margin for face',
                        type=int, default=10)
    parser.add_argument('--scale_rate',
                        help='Scale down or enlarge the original video img',
                        type=float, default=0.7)
    parser.add_argument('--show_rate',
                        help='Scale down or enlarge the imgs drawn by opencv',
                        type=float, default=1)
    parser.add_argument('--face_score_threshold',
                        help='The threshold of the extracted faces,range 0<x<=1',
                        type=float, default=0.85)
    parser.add_argument('--face_landmarks',
                        help='Draw five face landmarks on extracted face or not ', action="store_true")
    parser.add_argument('--no_display',
                        help='Display or not', action='store_true')
    args = parser.parse_args()
    return args


'''
This variation is based upon:
https://github.com/rlcode/reinforcement-learning/blob/master/2-cartpole/3-reinforce/cartpole_reinforce.py

We combine angle/throttle into a single 7*3 catagorical output
and use the loss function to define the reward.

---
Methods to create, use, save and load pilots. Pilots
contain the highlevel logic used to determine the angle

models to help direct the vehicles motion.

'''
# STATE_OPENCV:
#           supervised learning based upon OpenCV
#           OpenCv results incl reward @ control rpi & sent via msg
#           cache_state @ RLPi
#           rlfit: predict critic, train actor & critic @ deq @ RLPI
#              - note: OpenCv action not like actor results
#           Computation:
#              - Control: Open CV run once
#              - RLPi:    critic NN run @ enq batch
#                         actor  NN run @ enq batch
#                         actor and critic runs fit in deq batch
#           Proposed Computation:
#                         run actor predict, DUMMY, DUMMY during enqueue
#                         run critic NN during enqueue
#                         parallize actor and critic fits in deq batch
#                         3 Pi for Control, critic @ enq, actor @ enq
#                                  actor fit @ deq, critic fit @ deq
# STATE_TRIAL_NN:
#           RL training randomized trial_run(){actor} @ control pi
#           cache_state @ enq @ RLPI
#           compute_reward: OpenCv @ enq @ RLPI
#           rlfit: predict critic, train actor & critic @ deq @ RLPI
#              - note: randomization of actor actions possible
#              - note: control pi actor will be slightly out-of-sync (OK)
#              - pass both predicted & real values to RLPI
#           Computation:
#              - Control: actor NN run once
#              - RLPi:
#                         opencv run during enq
#                         critic NN run during enq
#                         batch actor and critic runs fit in deq batch
#           Proposed Computation:
#                         5 Pi for Control; critic @ enq; opencv@enq;
#                                  actor fit @ deq; critic fit @ deq
#                         3 Pi for Control; critic @ enq & opencv@enq;
#                            batch actor fit @ deq & critic fit @ deq
import os
import numpy as np
import keras
import donkeycar as dk
from keras.optimizers import Adam
from donkeycar.parts.RLOpenCV import ThrottleBase, HoughBundler, LaneLines
from keras import backend as K
from keras.models import Sequential
from keras.layers import Input, Dense, merge
from keras.layers import Input, Lambda
from keras.layers import Convolution2D, MaxPooling2D, Reshape, BatchNormalization
from keras.layers import Activation, Dropout, Flatten, Dense

def proximal_policy_optimization_loss_continuous(advantage, old_prediction):
      def loss(y_true, y_pred):

        import tensorflow as tf
        # return K.mean(K.square(y_pred - y_true), axis=-1)

        # adv = K.squeeze(advantage, axis = 1)
        # pred = K.squeeze(old_prediction, axis = 1)
        # adv_sum = K.sum(adv, axis=-1)
        # adv_mean = K.mean(adv, axis=-1)
        # pred_sum = K.sum(pred, axis=-1)
        # pred_mean = K.mean(pred, axis=-1)
        # if (pred_mean == PPO_OUT_OF_RANGE):
        # if (K.sum(adv_sum, -pred_sum) == K.sum(adv_mean, -pred_mean) and adv_sum != adv_mean):
        # if (K.equal(adv_sum, pred_sum) and K.equal(adv_mean, pred_mean) and K.not_equal(adv_sum, adv_mean)):
        # out_of_range = tf.constant(PPO_OUT_OF_RANGE, tf.shape(advantage))
        # out_of_range = K.constant(PPO_OUT_OF_RANGE, dtype=old_prediction.dtype, shape=old_prediction.shape)
        # pred_out_of_range = K.equal(old_prediction, out_of_range)
        # pred_out_of_range = K.equal(old_prediction, PPO_OUT_OF_RANGE)

        mean_sq_err = K.mean(K.square(y_pred - y_true), axis=-1)

        try:
          PPO_OUT_OF_RANGE = 1000    # negative of -1000
          checkifzero = K.sum(old_prediction, PPO_OUT_OF_RANGE)
          divbyzero = old_prediction / checkifzero
        except:
          return mean_sq_err

        # pred_out_of_range = K.mean((old_prediction / PPO_OUT_OF_RANGE), axis=-1)
        # pred_out_of_range = K.mean(K.equal(old_prediction , PPO_OUT_OF_RANGE), axis=-1)
        pred_out_of_range = K.mean(old_prediction, axis=-1)

        PPO_NOISE = 1.0
        var = keras.backend.square(PPO_NOISE)
        denom = K.sqrt(2 * np.pi * var)
        prob_num = K.exp(- K.square(y_true - y_pred) / (2 * var))
        old_prob_num = K.exp(- K.square(y_true - old_prediction) / (2 * var))

        prob = prob_num/denom
        old_prob = old_prob_num/denom
        r = prob/(old_prob + 1e-10)

        PPO_LOSS_CLIPPING = 0.2
        PPO_ENTROPY_LOSS = 5 * 1e-3 # Does not converge without entropy penalty
        # return -K.mean(K.minimum(r * advantage, K.clip(r, min_value=1 - PPO_LOSS_CLIPPING, max_value=1 + PPO_LOSS_CLIPPING) * advantage))
        # ppo_loss = -K.mean(K.minimum(r * advantage, K.clip(r, min_value=1 - PPO_LOSS_CLIPPING, max_value=1 + PPO_LOSS_CLIPPING) * advantage)) + PPO_ENTROPY_LOSS * (prob * K.log(prob + 1e-10))
        return -K.mean(K.minimum(r * advantage, K.clip(r, min_value=1 - PPO_LOSS_CLIPPING, max_value=1 + PPO_LOSS_CLIPPING) * advantage)) + PPO_ENTROPY_LOSS * (prob * K.log(prob + 1e-10))

        # out = tf.where(tf.equal(pred_out_of_range, PPO_OUT_OF_RANGE), mean_sq_err,  ppo_loss)
        # out = K.switch(K.equal(-1000, PPO_OUT_OF_RANGE), mean_sq_err,  ppo_loss)
        # out = K.switch(K.equal(pred_out_of_range, PPO_OUT_OF_RANGE), mean_sq_err,  ppo_loss)
        # out = K.switch( pred_out_of_range, K.zeros_like(pred_out_of_range),  K.zeros_like(pred_out_of_range))
        # out = K.switch( K.mean(old_prediction/PPO_OUT_OF_RANGE), mean_sq_err,  ppo_loss)
        # return out
      return loss

class KerasPilot():

    def load(self, model, model_path=None):
      global cfg

      # ARD: TODO: try using save_weights and load_weights instead of the whole
      # model. use the model definition from this file.
      try:
        #load keras model
        if model_path is None:
          model_path = cfg.RL_MODEL_PATH
        print("LOADING %s_%s" % (model_path, model))
        if model == "critic":
          Model = keras.models.load_model(os.path.expanduser(model_path + "_" + model))
        else:
          advantage_in = Input(shape=(1,))
          old_prediction_in = Input(shape=(1,))
          Model = keras.models.load_model(os.path.expanduser(model_path + "_" + model), custom_objects={'loss': proximal_policy_optimization_loss_continuous( advantage=advantage_in, old_prediction=old_prediction_in)})
        print("loaded")
        return Model
      except:
        print("No model to load")
        return None

    def load_weights(self, model, model_path=None):
      global cfg

      # ARD: TODO: try using save_weights and load_weights instead of the whole
      # model. use the model definition from this file.
      try:
        #load keras model
        if model_path is None:
          model_path = cfg.RL_MODEL_PATH
        print("LOADING %s_%s" % (model_path, model))
        if model == "critic":
          self.critic.load_weights(os.path.expanduser(model_path + "_" + model))
        else:
          self.actor.load_weights(os.path.expanduser(model_path + "_" + model))
        print("loaded")
        return True
      except:
        print("No weights to load")
        # raise
        return False

    def set_weights(self, w, model):
        if model == "actor":
          return self.actor.set_weights(w)
        else:
          return self.critic.set_weights(w)

    def get_weights(self, model):
        if model == "actor":
          return self.actor.get_weights()
        else:
          return self.critic.get_weights()

    def clear_session(self, model):
        # otherwise get_weights keeps consuming memory
        if model == "actor":
          return self.actor.backend.clear_session()
        else:
          return self.critic.backend.clear_session()

    def predicted_actions(self, img):
        # print("predicted_actions shape %d" % len(img.shape))
        if len(img.shape) != 4:
          img_arr = img.reshape((1,) + img.shape)
        else:
          img_arr = img
        # PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE, PPO_DUMMY_ZEROS = np.zeros((1, 1)), np.zeros((1, 1)), np.zeros((1, 1)), np.zeros((1, 1))
        # outputs = self.actor.predict([img_arr, PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE, PPO_DUMMY_ZEROS], batch_size=1)
        PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE = np.zeros((1, 1)), np.zeros((1, 1)), np.zeros((1, 1))
        outputs = self.actor.predict([img_arr, PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE], batch_size=1)
        predicted_angle = outputs[0]
        if predicted_angle < -1:
          print("Predicted angle too low: %f" % predicted_angle)
        if predicted_angle > 1:
          print("Predicted angle too high: %f" % predicted_angle)
        # print("predicted actions")
        # print(outputs[0].shape)
        # print(outputs[1].shape)
        predicted_throttle = outputs[1]
        predicted_throttle /= 1000
        predicted_throttle = outputs[1].reshape(1,1)
        predicted_angle = min(max(predicted_angle, -1),1)
        predicted_throttle = max(predicted_throttle, 0)
        # print(predicted_throttle.shape)
        # print("predicted_actions")
        # print(predicted_angle)
        # print(predicted_throttle)
        return predicted_angle, predicted_throttle

    def predicted_reward(self, img):
        img_arr = img.reshape((1,) + img.shape)
        output = self.critic.predict([img_arr], batch_size=1)
        reward = output[0][0]
        return reward

class KerasRLContinuous(KerasPilot):

    def __init__(self, model=None, LaneLine=None, *args, **kwargs):
        global cfg
        self.trial_cnt = 0
        self.ll = LaneLine
        self.critic = None
        self.actor = None
        cfg = dk.load_config(config_path=os.path.expanduser('~/donkeycar/donkeycar/parts/RLConfig.py'))
        super(KerasRLContinuous, self).__init__(*args, **kwargs)
        # following initializes self.actor / self.critic
        self.default_continuous(model)

    ################
    # CONTROL PI
    ################
    # trial_* are run at the control pi

    def trial_run(self, img_arr, trial_gps):
        img_arr = img_arr.reshape((1,) + img_arr.shape)
        # angle,throttle = self.model.predict(img_arr, batch_size=1)
        # PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE, PPO_DUMMY_ZEROS = np.zeros((1, 1)), np.zeros((1, 1)), np.zeros((1, 1)), np.zeros((1, 1))
        # outputs = self.actor.predict([img_arr, PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE, PPO_DUMMY_ZEROS], batch_size=1)
        PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE = np.zeros((1, 1)), np.zeros((1, 1)), np.zeros((1, 1))
        outputs = self.actor.predict([img_arr, PPO_DUMMY_ADVANTAGE, PPO_DUMMY_ANGLE, PPO_DUMMY_THROTTLE], batch_size=1)
        predicted_angle = outputs[0]
        predicted_throttle = outputs[1]
        predicted_throttle /= 1000
        # reward = outputs[2]
        rdm = np.random.random() * 10
        print("t_unb %f a_unb %f " % (predicted_throttle, predicted_angle))
        if (trial_gps):
          rdm = np.random.random() * 10
          # Training is done with throttle = LOW, so higher prob of MED/HIGH
          # random throttle: 0  / 1 /  2
          #                 10% /60%/30%
          if rdm < 1:
            delta_throttle = -1
          elif rdm < 2:
            delta_throttle = 0
          elif rdm < 8:
            delta_throttle = 1
          else:
            delta_throttle = 2
          trial_throttle = predicted_throttle + (delta_throttle * cfg.THROTTLE_INCREMENT)
          rdm = np.random.random() * 5
          if rdm <= 1:
            # fully random
            rdm = np.random.random()
            delta_steering = 2*(rdm - .5)
            trial_angle = delta_steering
            print("random steering: %f" % trial_angle)
          else:
            # bounded random
            # scale = int(np.random.random() * 4) + 1
            scale = 1
            # random delta steering: -2 / -1/  0/ +1/ +2
            #   with distribution:   10% 20% 40% 20% 10%
            # rdm = np.random.random() * 10
            rdm = 6    # no change
            if rdm < 1:
              delta_steering = -2
            elif rdm < 3:
              delta_steering = -1
            elif rdm < 7:
              delta_steering = 0
            elif rdm < 9:
              delta_steering = 1
            else:
              delta_steering = 2
            scale = 4   # scale of randomization
            # maximize entropy
            if predicted_angle + scale * abs(delta_steering) / 15 < -1:
              trial_angle = -1 - scale * abs(delta_steering) / 15
            if predicted_angle + scale * abs(delta_steering) / 15 > 1:
              trial_angle = 1 - scale * abs(delta_steering) / 15
            else:
              # return unbinned values back to the Control Pi
              trial_angle = predicted_angle + scale * delta_steering / 15
            print("tpred %f apred %f ttrl %f atrl %f t_delta %d a_delta %d scale %f" % (predicted_throttle, predicted_angle, trial_throttle, trial_angle, delta_throttle, delta_steering, scale))
        else:
          trial_angle = predicted_angle
          trial_throttle = predicted_throttle
        # angle is passed as unbinned between processes to maintain precision
        # throttle is just low/med/high
        trial_angle = min(max(trial_angle, -1),1)
        trial_throttle = max(trial_throttle, 0)
        return predicted_angle, predicted_throttle, trial_angle, trial_throttle

    ################
    # Reinforcement Learning PI
    ################
    # rl_* are run at the RL pi
    def rl_compute_reward(self, img, actual_throttle, actual_steering):

        simplecl = None
        lines    = None
        ylines   = None
        wlines   = None
        curcline = []
        currline = []
        curlline = []
        if cfg.USE_COLOR:
          line_color_simple, line_color_yellow, line_color_white = self.ll.get_line_color_info()

        if cfg.USE_COLOR and line_color_yellow is not None and line_color_white is not None and line_color_simple is not None:
          simplecl, wlines, ylines, roi = self.ll.process_img_color(img)
          if wlines is not None:
            currline, curlline, dummycurcline = self.ll.lrcsort(wlines)
          if ylines is not None:
            dummycurrline, dummycurlline, curcline = self.ll.lrcsort(ylines)
          if roi is not None:
            steering, throttle = self.ll.lrclines(currline, curlline, curcline, roi)
        else:
          simplecl, lines, roi = self.ll.process_img(img)
          if lines is not None and roi is not None:
            currline, curlline, curcline = self.ll.lrcsort(lines)
            if cfg.USE_COLOR:
              # img used to compute line colors
              steering, throttle = self.ll.lrclines(currline, curlline, curcline, img)
            else:
              # roi only used for debugging
              steering, throttle = self.ll.lrclines(currline, curlline, curcline, roi)

        laneWidth, pixPerFrame, donkeyvpangle, bestll, bestcl, bestrl, width, curpos = self.ll.get_map_data()

        if not cfg.DISABLE_EMERGENCY_STOP and laneWidth >0 and (curpos == 0 or curpos == 6):
          if simplecl is not None:
            curpos = 4
          else:
            # out of bounds
            print("reward = EMERG_STOP")
            # min reward
            return cfg.EMERGENCY_STOP
        '''
        '''
        if simplecl is not None:
          dist_from_center = abs(self.ll.closestX(simplecl) - (width/2))
          print("1 dist_from_center %f" % dist_from_center)
        elif (not cfg.USE_COLOR and lines is not None) or (cfg.USE_COLOR and (ylines is not None or wlines is not None)):
          if bestcl is not None:
            dist_from_center = abs(self.ll.closestX(bestcl) - (width/2))
            print("2 dist_from_center %f" % dist_from_center)
          elif bestll is not None and bestrl is not None:
            dist_from_center = abs((self.ll.closestX(bestll) + self.ll.closestX(bestrl))/2 - (width/2))
            print("3 dist_from_center %f" % dist_from_center)
          elif bestll is not None:
            dist_from_center = abs((self.ll.closestX(bestll) + laneWidth) - (width/2))
            print("4 dist_from_center %f" % dist_from_center)
          elif bestrl is not None:
            dist_from_center = abs((self.ll.closestX(bestrl) - laneWidth) - (width/2))
            print("5 dist_from_center %f" % dist_from_center)
          else:
            dist_from_center = width
            print("6 dist_from_center %f" % dist_from_center)
        else:
          # Major Penalty, resulting in 0 steering, 0 throttle
          # ARD emergency_stop: make sure we have POV before continuing?
          print("reward = EMERG_STOP")
          # min reward
          return cfg.EMERGENCY_STOP
        if dist_from_center < cfg.MIN_DIST_FROM_CENTER:
          dist_from_center = 0
          # print("7 dist_from_center %f" % dist_from_center)
        # steering = min(max((self.ll.closestX(cl) - midpix) / denom,-1),1)
        # reward_center is based upon post-action ROI (n+1)
        # print("reward_center %f laneWidth %f" % (reward_center, laneWidth))
        # reward_throttle is based upon pre-action ROI (n) at ControlPI:
        #   reward_throttle = throttle - minthrottle - battery_adjustment + 1

        # normalize reward_center & reward_throttle betweeen (0,1]
        reward_center = min(max(((3*laneWidth/4 - dist_from_center)/(3*laneWidth/4)),cfg.EMERGENCY_STOP) , 1) * 0.8
        reward_throttle = actual_throttle
        # normalize total reward between (0,1]
        reward = (reward_center) * (1 + (cfg.THROTTLE_BOOST*actual_throttle))
        print("reward %f reward_center %f reward_throttle %f laneWidth %f" % (reward, reward_center, reward_throttle, laneWidth))
        return reward*1000

    # called when batch is ready at RLPI
    # run at RL_pi
    # def rl_fit(self, trial_roi_batch, advantage_batch, angle_batch, predicted_angle_batch, predicted_throttle_batch, trial_angle_batch, trial_throttle_batch, actual_total_reward_batch):
    def rl_fit(self, batch):
      global cfg

      if batch is None:
        return   # batch not ready
      self.trial_cnt += 1
      # trial_roi = trial_roi.reshape((1,) + trial_roi.shape)
      print("actor train_on_batch")
      # print(batch[0])
      # print(batch[0].shape)
      # print(batch[1].shape)
      # print(batch[1])
      # print(batch[2].shape)
      # # print(batch[2])
      # print(batch[3].shape) # (20,1)
      # # print(batch[3])
      # print(batch[4].shape)
      # print(batch[4])
      # print(batch[5].shape)
      # print(batch[5])
      # print(batch[7].shape)
      # print(batch[7])
      # self.actor.train_on_batch([batch[0],batch[1],batch[2],batch[3],batch[7]], [batch[4],batch[5]])
      self.actor.train_on_batch([batch[0],batch[1],batch[2],batch[3]], [batch[4],batch[5]])
      print("critic train_on_batch")
      # print(batch[0].shape)
      # print(batch[0])
      # print(batch[6].shape)
      # print(batch[6])
      self.critic.train_on_batch([batch[0]], [batch[6]])
      if (self.trial_cnt % cfg.SAVE_NN == 0):
        self.actor.save_weights(os.path.expanduser(cfg.RL_MODEL_PATH + "_actor"))
        self.critic.save_weights(os.path.expanduser(cfg.RL_MODEL_PATH + "_critic"))
        # self.actor.save(os.path.expanduser(cfg.RL_MODEL_PATH + "_actor"))
        # self.critic.save(os.path.expanduser(cfg.RL_MODEL_PATH + "_critic"))

##########################################
# PPO
##########################################
# Initial framework taken from https://github.com/jaara/AI-blog/blob/master/CartPole-A3C.py
# This variation derived from https://github.com/OctThe16th/PPO-Keras/blob/master/Main.py

# move to RLConfig
# Only implemented clipping for the surrogate loss, paper said it was best
# PPO_LOSS_CLIPPING = 0.2
# PPO_NOISE = 1.0
# PPO_GAMMA = 0.99
# PPO_DUMMY_ACTION, PPO_DUMMY_VALUE = np.zeros((1, NUM_ACTIONS)), np.zeros((1, 1))
# PPO_ENTROPY_LOSS = 5 * 1e-3 # Does not converge without entropy penalty
# PPO_LR = 1e-4 # Lower lr stabilises training greatly

    '''
    def proximal_policy_optimization_loss_continuous(advantage, old_prediction):
      def loss(y_true, y_pred):
        PPO_NOISE = 1.0
        var = keras.backend.square(PPO_NOISE)
        denom = K.sqrt(2 * np.pi * var)
        prob_num = K.exp(- K.square(y_true - y_pred) / (2 * var))
        old_prob_num = K.exp(- K.square(y_true - old_prediction) / (2 * var))

        prob = prob_num/denom
        old_prob = old_prob_num/denom
        r = prob/(old_prob + 1e-10)

        PPO_LOSS_CLIPPING = 0.2
        PPO_ENTROPY_LOSS = 5 * 1e-3 # Does not converge without entropy penalty
        # return -K.mean(K.minimum(r * advantage, K.clip(r, min_value=1 - PPO_LOSS_CLIPPING, max_value=1 + PPO_LOSS_CLIPPING) * advantage))
        return -K.mean(K.minimum(r * advantage, K.clip(r, min_value=1 - PPO_LOSS_CLIPPING, max_value=1 + PPO_LOSS_CLIPPING) * advantage)) + PPO_ENTROPY_LOSS * (prob * K.log(prob + 1e-10))
      return loss
    '''

    ##########################################
    # MODEL USED BY CONTROLPI, RLPI (and controlpi2 if run)
    ##########################################
    def default_model(self, actor_critic = "actor"):
      from keras.models import Model
      from keras.layers import Input, Lambda
      from keras.layers import multiply, add

      PPO_LR = 1e-4 # Lower lr stabilises training greatly
      img_in = Input(shape=(120, 160, 3), name='img_in')

      if actor_critic == "actor":
        advantage_in = Input(shape=(1,))
        old_angle_prediction = Input(shape=(1,))
        old_throttle_prediction = Input(shape=(1,))
        # dummy_zeros = Input(shape=(1,))
        # continuous output of the angle, throttle and reward
        # lambda function uses the loss-only inputs: requirement for load
        # but lambda function is not trainable...
        # x = Lambda(lambda x: x[0])([img_in, advantage_in, old_angle_prediction, old_throttle_prediction])
        x = img_in
        # y = multiply([advantage_in, old_angle_prediction, old_throttle_prediction, dummy_zeros])
        # Convolution2D class name is an alias for Conv2D
        x = Convolution2D(filters=24, kernel_size=(5, 5), strides=(2, 2), activation='relu')(x)
        x = Convolution2D(filters=32, kernel_size=(5, 5), strides=(2, 2), activation='relu')(x)
        x = Convolution2D(filters=64, kernel_size=(5, 5), strides=(2, 2), activation='relu')(x)
        x = Convolution2D(filters=64, kernel_size=(3, 3), strides=(2, 2), activation='relu')(x)
        x = Convolution2D(filters=64, kernel_size=(3, 3), strides=(1, 1), activation='relu')(x)
        x = Flatten(name='flattened')(x)
        angle_out = Dense(units=1, activation='linear', name='angle_out')(x)
        # angle_out = add([angle_out,y])
        throttle_out = Dense(units=1, activation='linear', name='throttle_out')(x)

        # model = Model(inputs=[l_input], outputs=[angle_out, throttle_out])
        # model = Model(inputs=[img_in, advantage_in, old_angle_prediction, old_throttle_prediction, dummy_zeros], outputs=[angle_out, throttle_out])
        model = Model(inputs=[img_in, advantage_in, old_angle_prediction, old_throttle_prediction], outputs=[angle_out, throttle_out])
        # action, action_matrix, prediction from trial_run
        # reward is a function( angle, throttle)
        model.compile(optimizer=Adam(lr=PPO_LR),
                      loss=[proximal_policy_optimization_loss_continuous( advantage=advantage_in, old_prediction=old_angle_prediction),
                            proximal_policy_optimization_loss_continuous( advantage=advantage_in, old_prediction=old_throttle_prediction)])
      elif actor_critic == "critic":
        y = img_in
        # Convolution2D class name is an alias for Conv2D
        y = Convolution2D(filters=24, kernel_size=(5, 5), strides=(2, 2), activation='relu')(y)
        y = Convolution2D(filters=32, kernel_size=(5, 5), strides=(2, 2), activation='relu')(y)
        y = Convolution2D(filters=64, kernel_size=(5, 5), strides=(2, 2), activation='relu')(y)
        y = Convolution2D(filters=64, kernel_size=(3, 3), strides=(2, 2), activation='relu')(y)
        y = Convolution2D(filters=64, kernel_size=(3, 3), strides=(1, 1), activation='relu')(y)
        y = Flatten(name='flattened')(y)
        total_reward = Dense(units=1, activation='linear', name='total_reward')(y)
        model = Model(inputs=[img_in], outputs=[total_reward])
        model.compile(optimizer=Adam(lr=PPO_LR),
                      loss={'total_reward': 'mean_squared_error'})
      # model.summary()

      return model

    def default_continuous(self, mymodel=None):
      if (mymodel is None):
        # self.critic = self.load("critic")
        # print(self.critic)
        # print(self.actor)
        if self.critic is None:
          self.critic = self.default_model("critic")
          self.load_weights("critic")
        # self.actor = self.load("actor")
        if self.actor is None:
          self.actor = self.default_model("actor")
          self.load_weights("actor")
        print(" ")
        print("ACTOR MODEL SUMMARY:")
        self.actor.summary()
        print(" ")
        print("CRITIC MODEL SUMMARY:")
        self.critic.summary()
      else:
        if mymodel == "actor" or mymodel == "control":
          # self.actor = self.load(mymodel)
          if self.actor is None:
            self.actor = self.default_model(mymodel)
            self.load_weights(mymodel)
          print(" ")
          print("ACTOR MODEL SUMMARY:")
          self.actor.summary()
        else:
          # self.critic = self.load("critic")
          if self.critic is None:
            self.critic = self.default_model("critic")
            self.load_weights("critic")
          print(" ")
          print("CRITIC MODEL SUMMARY:")
          self.critic.summary()


#!/usr/bin/env python
# -*- coding:utf-8 -*-
from __future__ import annotations
from typing import Generator, Dict, Optional
import os
import time
import glob
import subprocess
import simplejson as json
import numpy as np
import cv2
from cv2 import dnn_superres, dnn_superres_DnnSuperResImpl
from PIL import Image, ImageDraw, ImageFont
import matplotlib.cm as cm
from utils.nms import filter_bboxes
from utils.image_preproc import (
    adjust_white_balance,
    smooth_image,
    correct_contrast,
    upsample_image,
    # levelize_histogram,
    # correct_contrast_using_lut,
)

ld_library_path = subprocess.Popen(
    'pyenv which python', stdout=subprocess.PIPE, shell=True
).communicate()[0].decode('utf-8').strip().split(os.sep)[:-2]
ld_library_path.append('lib')
ld_library_path = os.sep.join(ld_library_path)
os.environ['LD_LIBRARY_PATH'] = ld_library_path

class Session(object):
    def __init__(
        self: Session,
        path: str,
        clarify_image: bool = False,
        use_superres: bool = False
    ) -> None:
        self.name = os.path.basename(path)
        self.image_id = os.path.splitext(self.name)[0]
        self.is_clarify_image = clarify_image
        self.use_superres = use_superres
        self.raw_image = cv2.imread(path)
        self.raw_height = self.raw_image.shape[0]
        self.raw_width = self.raw_image.shape[1]
        self.fine_image = self.raw_image.copy()
        self.model_height = None
        self.model_width = None
        self.scale = None
        self.image_height = None
        self.image_width = None
        self.offset_height = None
        self.offset_width = None
        self.pad_image = None
        self.yolo_input = None
        self.yolov5_input = None
        self.crowddet_input = None
        self.elapsed_ms = None
        self.pred_count = None
        self.pred_bboxes = None
        return

    def clarify_image(
        self: Session,
        sr: Optional[dnn_superres_DnnSuperResImpl]
    ) -> None:
        if not self.is_clarify_image:
            return
        image = self.fine_image
        image = adjust_white_balance(image)
        image = correct_contrast(image)
        image = smooth_image(image)
        if self.use_superres:
            assert sr is not None
            image = upsample_image(image, sr)
        # stop using below functions
        # image = levelize_histogram(image)
        # image = correct_contrast_using_lut(image)
        self.fine_image = image
        return

    def padding_image(
        self: Session,
        model_height: int,
        model_width: int
    ) -> None:
        # calc height, width, offset to resize
        self.model_height = model_height
        self.model_width = model_width
        self.scale = min(
            float(self.model_height) / self.raw_height,
            float(self.model_width) / self.raw_width
        )
        self.image_height = int(round(self.raw_height * self.scale))
        self.image_width = int(round(self.raw_width * self.scale))
        self.offset_height = (
            self.model_height - self.image_height
        ) // 2
        self.offset_width = (
            self.model_width - self.image_width
        ) // 2
        # resize image and put it on to the background image
        if self.scale >= 1:
            interpolation = cv2.INTER_CUBIC
        else:
            interpolation = cv2.INTER_AREA
        img = cv2.resize(
            self.fine_image.copy(),
            (self.image_width, self.image_height),
            interpolation=interpolation
        )
        background = np.full(
            (self.model_height, self.model_width, 3),
            128, dtype=np.uint8
        )
        background[
            self.offset_height: self.offset_height + self.image_height,
            self.offset_width: self.offset_width + self.image_width,
        ] = img
        self.pad_image = background
        return

    def rescale_xyxy(self: Session, xyxy: np.ndarray) -> np.ndarray:
        xyxy[:, 0] = np.maximum(
            (xyxy[:, 0] - self.offset_width) / self.scale,
            0
        )
        xyxy[:, 1] = np.maximum(
            (xyxy[:, 1] - self.offset_height) / self.scale,
            0
        )
        xyxy[:, 2] = np.minimum(
            (xyxy[:, 2] - self.offset_width) / self.scale,
            self.raw_width
        )
        xyxy[:, 3] = np.minimum(
            (xyxy[:, 3] - self.offset_height) / self.scale,
            self.raw_height
        )
        return xyxy

    def draw_prediction(self: Session, category_map: Dict) -> Image:
        fontsize = int(round(max(
            self.raw_width, self.raw_height
        ) / 32.0))
        buf = self.raw_image.copy()
        buf = cv2.cvtColor(buf, cv2.COLOR_BGR2RGB)
        buf = Image.fromarray(buf)
        draw = ImageDraw.Draw(buf)
        font = ImageFont.truetype(
            font='TakaoGothic.ttf',
            size=fontsize
        )
        default_color = (0xff, 0xff, 0xff)
        margin = 5
        bboxes = self.pred_bboxes.tolist()
        for bbox in bboxes:
            xmin, ymin, xmax, ymax, cls, prob = bbox
            xmin = int(max(xmin, margin))
            ymin = int(max(ymin, margin))
            xmax = int(min(xmax, self.raw_width - margin))
            ymax = int(min(ymax, self.raw_height - margin))
            cls = int(cls)
            prob = round(prob, 3)
            name = category_map[cls]['category_name']
            if name is None:
                name = str(cls)
            color = tuple(np.array(np.array(cm.jet((
                prob - 0.25  # self.config.confidence_threshold
            ) / (
                1.0 - 0.25  # self.config.confidence_threshold
            ))) * 255, dtype=np.uint8).tolist())
            color = tuple(color[0:3])
            draw.rectangle(
                (xmin, ymin, xmax, ymax),
                fill=None, outline=color, width=2
            )
            draw.text(
                (xmin + 5, ymin + 5), name,
                fill=color, font=font
            )
            draw.text(
                (xmin + 5, ymin + 5 + fontsize), '%0.3f' % prob,
                fill=color, font=font
            )
        text = 'Elapsed Time: %d[ms]' % self.elapsed_ms
        draw.text((5, 5), text, fill=default_color, font=font)
        text = 'Detected Objects: %d' % self.pred_count
        draw.text((5, 5 + fontsize), text, fill=default_color, font=font)
        return buf

class Config(object):
    def __init__(
        self: Config,
        model: str,
        framework: str,
        quantize: str,
        image_dir: str,
        conf_threshold: float,
        iou_threshold: float,
        clarify_image: bool,
        use_superres: bool,
        disable_soft_nms: bool
    ) -> None:
        self.model = model
        self.framework = framework
        self.quantize = quantize
        self.image_dir = image_dir
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.clarify_image = clarify_image
        self.use_superres = use_superres
        self.disable_soft_nms = disable_soft_nms
        return

class Framework(object):
    def __init__(
        self: Framework,
        config: Config
    ) -> None:
        self.config = config
        self.input_blob = list()
        self.input_shape = [None, None, None, None]
        return

    def inference(self: Framework, sess: Session) -> np.ndarray:
        return np.empty((0, 6), dtype=float)

class Model(object):
    def __init__(
        self: Model,
        config: Config
    ) -> None:
        self.config = config
        # self.framework= Framework(config=config)
        self.framework = None
        self.category_map = self.read_labels()
        return

    def read_labels(self: Model) -> Dict:
        category_map = dict()
        with open('labels/coco_labels.txt', 'rt') as rf:
            for i, raw in enumerate(rf.read().strip().splitlines()):
                id_label = raw.strip().split(' ', 1)
                category_id = int(id_label[0])
                category_name = id_label[1].strip()
                category_map[i] = {
                    'coco_category_id': category_id + 1,
                    'category_name': category_name,
                }
        return category_map

    def prep_image(self: Model, sess: Session) -> None:
        return

    def inference(self: Model, sess: Session) -> np.ndarray:
        return self.framework.inference(sess=sess)

class Detector(object):
    def __init__(
        self: Detector,
        config: Config
    ) -> None:
        self.config = config
        # self.model = Model(config=config)
        self.model = None
        dataset_name = config.image_dir.split(os.sep)[-1]
        if self.config.framework == 'tflite':
            self.result_dir = 'results/%s/%s_%s_%s' % (
                dataset_name, config.model,
                config.framework, config.quantize
            )
        else:
            self.result_dir = 'results/%s/%s_%s' % (
                dataset_name, config.model, config.framework
            )
        os.makedirs(self.result_dir, exist_ok=True)
        self.sr = None
        if config.use_superres:
            self.sr = dnn_superres.DnnSuperResImpl_create()
            self.sr.readModel('superres/ESPCN_x4.pb')
            self.sr.setModel('espcn', 4)
        self.wf = open(os.path.join(
            self.result_dir, 'predictions.jsonl'
        ), 'wt')
        return

    def print_header(self: Detector) -> None:
        if self.config.framework == 'tflite':
            framework = f'{self.config.framework}_{self.config.quantize}'
        else:
            framework = self.config.framework
        print(
            '=== MODEL: %s, FRAMEWORK: %s ===' % (
                self.config.model, framework
            )
        )
        return

    def yield_session(self: Detector) -> Generator[Session, None, None]:
        for path in sorted(glob.glob(f'{self.config.image_dir}/*')):
            if not path.endswith(('.jpg', '.png')):
                continue
            yield Session(
                path=path,
                clarify_image=self.config.clarify_image,
                use_superres=self.config.use_superres
            )
        return

    def prep_image(self: Detector, sess: Session) -> None:
        sess.clarify_image(sr=self.sr)
        self.model.prep_image(sess=sess)
        return

    def inference(self: Detector, sess: Session) -> None:
        # image preprocessing
        self.prep_image(sess=sess)
        # inference
        start_time = time.perf_counter()
        pred = self.model.inference(sess=sess)
        pred = filter_bboxes(
            pred,
            conf_threshold=self.config.conf_threshold,
            iou_threshold=self.config.iou_threshold,
            disable_soft_nms=self.config.disable_soft_nms
        )
        end_time = time.perf_counter()
        # sort bounding boxes by confidence ascending
        pred = pred[np.argsort(pred[:, 5])]
        # results of this session
        sess.elapsed_ms = int(round((end_time - start_time) * 1000))
        sess.pred_count = pred.shape[0]
        sess.pred_bboxes = pred
        return

    def print_result(self: Detector, sess: Session) -> None:
        print(
            '%s: count=%d, time=%dms' % (
                sess.name, sess.pred_count, sess.elapsed_ms
            )
        )
        return

    def dump_result(self: Detector, sess: Session) -> None:
        if self.config.framework == 'tflite':
            framework = f'{self.config.framework}_{self.config.quantize}'
        else:
            framework = self.config.framework
        bboxes = list()
        for pbox in sess.pred_bboxes.tolist():
            box = [
                float(pbox[0]), float(pbox[1]),
                float(pbox[2] - pbox[0]), float(pbox[3] - pbox[1])
            ]
            bbox = {
                'category_id': self.model.category_map[
                    int(pbox[4])
                ]['coco_category_id'],
                'bbox': box,
                'score': float(pbox[5]),
            }
            bboxes.append(bbox)
        stats = {
            'model': self.config.model,
            'framework': framework,
            'filename': sess.name,
            'count': sess.pred_count,
            'msec': sess.elapsed_ms,
            'image_id': sess.image_id,
            'bboxes': bboxes,
        }
        self.wf.write(json.dumps(stats) + '\n')
        self.wf.flush()
        return

    def dump_image(self: Detector, sess: Session) -> None:
        image = sess.draw_prediction(category_map=self.model.category_map)
        path_image = os.path.join(self.result_dir, sess.name)
        image.save(path_image)
        return

    def close(self: Detector) -> None:
        self.wf.close()
        return


import os
import sys
import cv2
import argparse
import numpy as np

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import seaborn as sns
import SimpleITK as sitk
import json
from mean_average_precision.detection_map import DetectionMAP

sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'lib'))

import torch
import torch.utils.data

from nets.hourglass import get_hourglass
from utils.utils import load_model
from utils.image import transform_preds, get_affine_transform
from utils.post_process import ctdet_decode

# Training settings
parser = argparse.ArgumentParser(description='centernet')

parser.add_argument('--root_dir', type=str, default='./')
parser.add_argument('--ckpt_dir', type=str, default='./ckpt/hourglass/checkpoint.t7')
parser.add_argument('--arch', type=str, default='large_hourglass')

cfg = parser.parse_args()

os.chdir(cfg.root_dir)

def evaluation(model,validation_folder,visualization = True,map_save_name='map.png'):
  frames = []
  n_class = 1
  mAP = DetectionMAP(n_class)
  max_per_image = 100

  validation_label = r'L:\TB_Project\src\Detector\data\coco\annotations\instances_val2017.json'
  with open(validation_label,'r') as load_f:
    load_dict = json.load(load_f)
  file_dict = {}
  for file in load_dict['images']:
      file_dict[file['file_name']] = file['id']
  label_dict = {}
  for label in load_dict['annotations']:
      label_dict[label['image_id']] = []
  for label in load_dict['annotations']:
      label_dict[label['image_id']].append(label['bbox'])

  for img_file in os.listdir(validation_folder):
    img_file_path = os.path.join(validation_folder,img_file)
    image = np.load(img_file_path)

    gt_boxes = []
    for (x_0,y_0,bbox_width,bbox_height) in label_dict[file_dict[img_file]]:
      x_1 = x_0 + bbox_width
      y_1 = y_0 + bbox_height
      gt_boxes.append((x_0,y_0,x_1,y_1))

    height, width = image.shape[0:2]
    padding = 127 if 'hourglass' in cfg.arch else 31
    imgs = {}

    for scale in [1]:
      new_height = int(height * scale)
      new_width = int(width * scale)

      img_height = 512
      img_width = 512
      center = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
      scaled_size = max(height, width) * 1.0
      scaled_size = np.array([scaled_size, scaled_size], dtype=np.float32)

      img = cv2.resize(image, (new_width, new_height))
      trans_img = get_affine_transform(center, scaled_size, 0, [img_width, img_height])
      img = cv2.warpAffine(img, trans_img, (img_width, img_height))

      img = (img.astype(np.float32)+1024.) / 4096.
      img -= 0.09102
      img /= 0.3735
      img = np.stack((img,img,img),axis=2)
      img = img.transpose(2, 0, 1)[None, :, :, :]  # from [H, W, C] to [1, C, H, W]

      imgs[scale] = {'image': torch.from_numpy(img).float(),
                    'center': np.array(center),
                    'scale': np.array(scaled_size),
                    'fmap_h': np.array(img_height // 4),
                    'fmap_w': np.array(img_width // 4)}

    with torch.no_grad():
      detections = []
      for scale in imgs:
        imgs[scale]['image'] = imgs[scale]['image'].cuda()

        output = model(imgs[scale]['image'])[-1]
        dets = ctdet_decode(*output, K=100) #test_topk
        dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])[0]

        top_preds = {}
        dets[:, :2] = transform_preds(dets[:, 0:2],
                                      imgs[scale]['center'],
                                      imgs[scale]['scale'],
                                      (imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
        dets[:, 2:4] = transform_preds(dets[:, 2:4],
                                      imgs[scale]['center'],
                                      imgs[scale]['scale'],
                                      (imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
        cls = dets[:, -1]
        for j in range(1):
          inds = (cls == j)
          top_preds[j + 1] = dets[inds, :5].astype(np.float32)
          top_preds[j + 1][:, :4] /= scale

        detections.append(top_preds)

      bbox_and_scores = {}
      for j in range(1):
        temp= np.concatenate([d[j+1] for d in detections], axis=0)
        bbox_and_scores[j] = temp

      scores = np.hstack([bbox_and_scores[j][:, 4] for j in range(1)])

      if len(scores) > max_per_image:
        kth = len(scores) - max_per_image
        thresh = np.partition(scores, kth)[kth]
        for j in range(1):
          keep_inds = (bbox_and_scores[j][:, 4] >= thresh)
          bbox_and_scores[j] = bbox_and_scores[j][keep_inds]

      fig = plt.figure(0)
      matplotlib.use('QT5Agg')
      plt.imshow(image)
      count = 0

      pred_bb = []
      pred_cls = []
      pred_conf = []
      gt_bb = []
      gt_cls = []
      for lab in bbox_and_scores:
        for boxes in bbox_and_scores[lab]:
          x1, y1, x2, y2, score = boxes
          if score > 0.15:
            pred_bb.append([x1, y1, x2, y2])
            pred_cls.append(0)
            pred_conf.append(score)
            count = count+1
            plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1,
                                          linewidth=2, edgecolor='r', facecolor='none'))
            plt.text(x1 + 3, y1 + 3, '%.2f' % score,
                    bbox=dict(facecolor='r', alpha=0.3), fontsize=7, color='k')
      for boxes in gt_boxes:
        x1, y1, x2, y2 = boxes
        gt_bb.append(boxes)
        gt_cls.append(0)
        plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1,
                                          linewidth=2, edgecolor='r', facecolor='none'))

      frames.append((np.array(pred_bb), np.array(pred_cls), np.array(pred_conf), np.array(gt_bb), np.array(gt_cls)))
      fig.patch.set_visible(False)
      plt.axis('off')
      if visualization:
        save_path='./result/'
        save_name=save_path+img_file[:-3]+'png'
        plt.savefig(save_name, dpi=300, transparent=True)
      plt.close('all')

  for frame in frames:
      mAP.evaluate(*frame)
  mAP.plot()
  plt.savefig(map_save_name)

def main():
  cfg.device = torch.device('cuda')
  torch.backends.cudnn.benchmark = False
  validation_folder=r'L:\TB_Project\src\Detector\data\coco\val2017'
  print('Creating model...')
  if 'hourglass' in cfg.arch:
    model = get_hourglass[cfg.arch]
  else:
    raise NotImplementedError

  model = load_model(model, cfg.ckpt_dir)
  model = model.to(cfg.device)
  model.eval()
  evaluation(model,validation_folder)

if __name__ == '__main__':
  main()