<reponame>hzitoun/TripletLossFace
import tensorflow as tf
import sys
import json
import numpy as np
import cv2
from tqdm import tqdm

sys.path.append("../")

from deep_learning.make_better_dataset_for_deepfake.main_data_creator import FaceExtractor

class Claudia:
	def load_full(self, path):
		image = tf.io.read_file(path)
		image = tf.image.decode_jpeg(image, channels=3)

		return image.numpy()

	def set_image(self, image):
		image = tf.image.resize(image, (self.input_shape[0], self.input_shape[1]), method="nearest")

		return tf.cast(image, tf.float32)/255.

	def save_json(self):
		with open(self.json_path, 'w') as f:
			json.dump(self.json, f)

	def create_color_map(self):
		for key in self.json:
			try:
				self.colors[key]
			except KeyError:				
				color = tuple(np.random.choice(range(256), size=3))
				color = (int(color[0]), int(color[1]), int(color[2]))

				self.colors[key] = color

	def __init__(self, model_path: str):
		self.model = tf.keras.models.load_model(model_path)
		self.faceExtractor = FaceExtractor()

		self.input_shape = self.model.layers[0].input_shape[0][1:]

		self.json_path = "my_data.json"

		try:
			with open(self.json_path, 'rb') as f:
				self.json = json.loads(f.read())
		except:
			with open(self.json_path, 'w+') as f:
				self.json = {}		

		self.colors = {}
		self.create_color_map()

		self.video_writer = None
		self.cosine_loss = tf.keras.losses.CosineSimilarity()

	def get_output_from_image(self, path, get_face: bool = True, l2: bool = False):
		if type(path) == str:
			image = self.load_full(path)
		else:
			image = path

		if get_face:
			faces, all_frames = self.faceExtractor.extract([image])
			faces = faces[0]
			all_frames = all_frames[0]

		else:
			faces = [image]
			all_frames = [(0,image.shape[1], 0, image.shape[0])]

		outputs = []
		for face in faces:
			face = self.set_image(face)
			output = self.model(tf.expand_dims(face, axis=0))

			if l2:
				output = tf.nn.l2_normalize(output, 1, 1e-10)

			outputs.append(output)

		return image, outputs, all_frames

	def add_to_json(self, output):
		print("ADDED")
		i = len(self.json.keys()) + 1
		self.json[str(i)] = list(output.numpy().tolist())
		self.save_json()
		self.create_color_map()

		return str(i)

	def index_image(self, path, get_face: bool = True, print_out: bool = False, th: float = -0.60):
		image, output, all_frames = self.get_output_from_image(path, get_face, l2=True)

		mins = []
		assert len(all_frames) == len(output)

		for i in range(len(output)):
			founded = False

			for key in self.json:
				my_min = (100000, "")
				oo = self.json[key]
				dist = self.cosine_loss(tf.convert_to_tensor(oo), output[i]).numpy()

				if dist <= th:
					if dist < my_min[0]:
						my_min = (dist, key)

						founded = True

			if not founded:
				new_key = self.add_to_json(output[i])
				my_min = (-1., new_key)

			mins.append(my_min)

		return mins, image, all_frames

	def mark(self, image, min_im, all_frames):
		for (confidance, who), frame in zip(min_im, all_frames):
			try:
				color = self.colors[str(who)]
				x1, x2, y1, y2 = frame
				cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
				cv2.putText(image, f"id: {str(who)}- conf:{abs(round(float(confidance), 2))}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA) # -{round(float(confidance), 2)}
			except KeyError:
				continue

		return image


	def show_who_in_image(self, path, get_face: bool = True, show: bool = True, turn_rgb: bool = True):
		min_im, image, all_frames = self.index_image(path, get_face)

		for (confidance, who), frame in zip(min_im, all_frames):
			try:
				color = self.colors[str(who)]
				x1, x2, y1, y2 = frame
				cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
				cv2.putText(image, f"id: {str(who)}- conf:{abs(round(float(confidance), 2))}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA) # -{round(float(confidance), 2)}
			except KeyError:
				continue

		if turn_rgb:
			image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

		if show:
			cv2.imshow("a", image)
			cv2.waitKey(1)

		return image, min_im, all_frames

	def go_for_video(self, path, i: int = -1):
		cap = cv2.VideoCapture(path)
		n = 0

		if i != -1:
			bar = tqdm(total=i)
		else:
			bar = tqdm()

		min_im, all_frames = None, None
		while True:
			try:
				result, frame = cap.read()
				if n % 6 == 0:
					if self.video_writer is None:
						h, w, c = frame.shape
						self.video_writer = cv2.VideoWriter('result.avi', cv2.VideoWriter_fourcc(*"MJPG"), 30,(w,h))

					if not result:
						break

					frame, min_im, all_frames = self.show_who_in_image(frame, True, False, turn_rgb=False)
					self.video_writer.write(frame)

				else:
					frame = self.mark(frame, min_im, all_frames)
					self.video_writer.write(frame)

				n += 1
				bar.update()
			except Exception as e:
				print(e)
				continue


		self.video_writer.release()



if __name__ == '__main__':
	claudia = Claudia("../deep_learning/models/triplet_inception_resnet_v1_0.h5")

	claudia.go_for_video("bbt_test1.mp4")
import tensorflow as tf
import sys
import json
import numpy as np
import cv2
from tqdm import tqdm

sys.path.append("../")

from deep_learning.make_better_dataset_for_deepfake.main_data_creator import FaceExtractor

class Claudia:
	def load_full(self, path):
		image = tf.io.read_file(path)
		image = tf.image.decode_jpeg(image, channels=3)

		return image.numpy()

	def set_image(self, image):
		image = tf.image.resize(image, (self.input_shape[0], self.input_shape[1]), method="nearest")

		return tf.cast(image, tf.float32)/255.

	def save_json(self):
		with open(self.json_path, 'w') as f:
			json.dump(self.json, f)

	def create_color_map(self):
		for key in self.json:
			try:
				self.colors[key]
			except KeyError:				
				color = tuple(np.random.choice(range(256), size=3))
				color = (int(color[0]), int(color[1]), int(color[2]))

				self.colors[key] = color

	def __init__(self, model_path: str):
		self.model = tf.keras.models.load_model(model_path)
		self.faceExtractor = FaceExtractor()

		self.input_shape = self.model.layers[0].input_shape[0][1:]

		self.json_path = "my_data.json"

		try:
			with open(self.json_path, 'rb') as f:
				self.json = json.loads(f.read())
		except:
			with open(self.json_path, 'w+') as f:
				self.json = {}		

		self.colors = {}
		self.create_color_map()

		self.video_writer = None
		self.cosine_loss = tf.keras.losses.CosineSimilarity()

	def get_output_from_image(self, path, get_face: bool = True, l2: bool = False):
		if type(path) == str:
			image = self.load_full(path)
		else:
			image = path

		if get_face:
			faces, all_frames = self.faceExtractor.extract([image])
			faces = faces[0]
			all_frames = all_frames[0]

		else:
			faces = [image]
			all_frames = [(0,image.shape[1], 0, image.shape[0])]

		outputs = []
		for face in faces:
			face = self.set_image(face)
			output = self.model(tf.expand_dims(face, axis=0))

			if l2:
				output = tf.nn.l2_normalize(output, 1, 1e-10)

			outputs.append(output)

		return image, outputs, all_frames

	def add_to_json(self, output):
		print("ADDED")
		i = len(self.json.keys()) + 1
		self.json[str(i)] = list(output.numpy().tolist())
		self.save_json()
		self.create_color_map()

		return str(i)

	def index_image(self, path, get_face: bool = True, print_out: bool = False, th: float = -0.60):
		image, output, all_frames = self.get_output_from_image(path, get_face, l2=True)

		mins = []
		assert len(all_frames) == len(output)

		for i in range(len(output)):
			founded = False

			for key in self.json:
				my_min = (100000, "")
				oo = self.json[key]
				dist = self.cosine_loss(tf.convert_to_tensor(oo), output[i]).numpy()

				if dist <= th:
					if dist < my_min[0]:
						my_min = (dist, key)

						founded = True

			if not founded:
				new_key = self.add_to_json(output[i])
				my_min = (-1., new_key)

			mins.append(my_min)

		return mins, image, all_frames

	def mark(self, image, min_im, all_frames):
		for (confidance, who), frame in zip(min_im, all_frames):
			try:
				color = self.colors[str(who)]
				x1, x2, y1, y2 = frame
				cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
				cv2.putText(image, f"id: {str(who)}- conf:{abs(round(float(confidance), 2))}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA) # -{round(float(confidance), 2)}
			except KeyError:
				continue

		return image


	def show_who_in_image(self, path, get_face: bool = True, show: bool = True, turn_rgb: bool = True):
		min_im, image, all_frames = self.index_image(path, get_face)

		for (confidance, who), frame in zip(min_im, all_frames):
			try:
				color = self.colors[str(who)]
				x1, x2, y1, y2 = frame
				cv2.rectangle(image, (x1, y1), (x2, y2), color, 4)
				cv2.putText(image, f"id: {str(who)}- conf:{abs(round(float(confidance), 2))}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3, cv2.LINE_AA) # -{round(float(confidance), 2)}
			except KeyError:
				continue

		if turn_rgb:
			image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

		if show:
			cv2.imshow("a", image)
			cv2.waitKey(1)

		return image, min_im, all_frames

	def go_for_video(self, path, i: int = -1):
		cap = cv2.VideoCapture(path)
		n = 0

		if i != -1:
			bar = tqdm(total=i)
		else:
			bar = tqdm()

		min_im, all_frames = None, None
		while True:
			try:
				result, frame = cap.read()
				if n % 6 == 0:
					if self.video_writer is None:
						h, w, c = frame.shape
						self.video_writer = cv2.VideoWriter('result.avi', cv2.VideoWriter_fourcc(*"MJPG"), 30,(w,h))

					if not result:
						break

					frame, min_im, all_frames = self.show_who_in_image(frame, True, False, turn_rgb=False)
					self.video_writer.write(frame)

				else:
					frame = self.mark(frame, min_im, all_frames)
					self.video_writer.write(frame)

				n += 1
				bar.update()
			except Exception as e:
				print(e)
				continue


		self.video_writer.release()



if __name__ == '__main__':
	claudia = Claudia("../deep_learning/models/triplet_inception_resnet_v1_0.h5")

	claudia.go_for_video("bbt_test1.mp4")
# -{round(float(confidance), 2)} # -{round(float(confidance), 2)}