
# Toto je skript pro porizeni datasetu k trenovani cnn

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import copy
from PIL import Image
import operator
import subprocess
import matplotlib.pyplot as plt



def pridej_index(soubor, nazev):
	f = open(str(soubor),"r+")
	cislo = int(f.read())
	f.seek(0)
	f.truncate()
	cislo = str(cislo+1)
	f.write(cislo)
	f.close()

	vystupninazev = nazev + "_" + str(cislo.zfill(4)) + ".png"
	return vystupninazev


def parsuj(puv0):
	#puv = puv0[:,49:650]	# orez (binning 3)
	puv = puv0 				# resp. pro jiz upravene

	rozmery_obr = puv.shape 


	#----------------------------------- Rozdeleni do tri obrazku------------------------------

	img = np.zeros((int(rozmery_obr[0]/3),rozmery_obr[1],3), dtype=int)
	img = np.uint(img)
	radek = 0;
	for i in range(rozmery_obr[0]):
		#print(i)

		mod = int(i%3)
		if mod == 0:
			img[radek,:,0] = puv[i,:]
		if mod == 1:
			img[radek,:,1] = puv[i,:]
		if mod == 2:
			img[radek,:,2] = puv[i,:]
			radek = radek+1
	img = np.rot90(img,2)
	#BestOne = img[:,:,1]
	return img				# Vraci RGB snimek






def vyberkanalu(img):					# Vyber na zaklade porovnani vzdalenosti dvou nejvetsich peaku histogramu
	h = np.zeros((3,256),dtype=int)
	h = np.uint(h)

	idxmax = np.zeros((3), dtype=int)
	idxmax = np.uint(idxmax)

	idxmax2 = np.zeros((3), dtype=int)
	idxmax2 = np.uint(idxmax2)

	hyst = 30							# Zda se jedna o jeden peak urcuje velikost hystereze

	for vyska in range(img.shape[2]):
		a = np.histogram(img[:,:,vyska].ravel(), 256, [0, 256])
		h[vyska,:] = a[0]



		idxmax[vyska] = np.argmax(h[vyska,:])

		if idxmax[vyska]+hyst > 255:
			h[vyska,(idxmax[vyska]-hyst):] = 0
		elif (idxmax[vyska]-hyst)<0:
			h[vyska,:idxmax[vyska]+hyst] = 0
		else:
			h[vyska, idxmax[vyska]-hyst:idxmax[vyska]+hyst] = 0

		idxmax2[vyska] = np.argmax(h[vyska,:])


	idxmax = np.int16(idxmax)
	idxmax2 = np.int16(idxmax2)

	ratio = abs(idxmax-idxmax2)

	Best = np.argmax(ratio)

	return Best							# Vraci se index nejlepsiho kanalu


















if __name__ == "__main__":

	#img =cv2.imread('im2_bin3_g.png',0)
	#img =cv2.imread('../im5_bin3.png',0)
	


	
	os.chdir("./../Akvizice")		# Zmena adresare
	subprocess.call("./Akvizice")	# Spusteni c++ programu k akvizici snimku
	os.chdir("./../Dataset")		# Navrat do puvodniho adresare
	img =cv2.imread("../Akvizice/rawsnimek.png",0)	# Nacteni porizeneho snimku
	img = img[:,60:650]				# Oriznuti snimku z prava a z leva
	

	pars = parsuj(img)				# Funkce pro parsovani suroveho snimku do tri dle spektra osvetleni

	
	#----------------------- Zobrazeni RGB snimku pro posouzeni ------------
	#cv2.imshow("pars",cv2.cvtColor(np.uint8(pars), cv2.COLOR_BGR2RGB))
	#cv2.waitKey(0)
	plt.imshow(pars)
	plt.title("Nahled barevneho obrazku - zkontroluj ho")
	plt.show()
	# ----------------------------------------------------------------------


	print("Je obrazek v poradku? Ok = y or z or 4, Nok = ostatni")
	vpohode = input()
	if vpohode == "4" or vpohode == "y" or vpohode == "z":		

		img = pars
		img = np.uint8(img)
		Best = vyberkanalu(img)		# Vyber nejkontrastnejsiho kanalu
		
		prvni = False
		druha = False
		treti = False

		# Nacitani hlavnich referencnich snimku
		template_power = cv2.imread('ref_power.png',0)
		template_trio = cv2.imread('ref_trio.png',0)
		template_cina = cv2.imread('ref_cina.png',0)
		template_tucnak = cv2.imread('ref_tucnak.png',0)
		

		
		# Srovnani se vzorem - porizeny snimek s referencnimi vyrezy
		res_power = cv2.matchTemplate(img[:,:,Best], template_power, cv2.TM_CCOEFF_NORMED)
		res_trio = cv2.matchTemplate(img[:,:,Best], template_trio, cv2.TM_CCOEFF_NORMED)
		res_cina = cv2.matchTemplate(img[:,:,Best], template_cina, cv2.TM_CCOEFF_NORMED)
		res_tucnak = cv2.matchTemplate(img[:,:,Best], template_tucnak, cv2.TM_CCOEFF_NORMED)

		# Do slovniku "srovnani" je vlozeno oznaceni fixy, spolu s nejvyssi hodnotou shody vzoru s danym snimkem
		srovnani = {'power':np.max(res_power), 'trio':np.max(res_trio), 'cina':np.max(res_cina), 'tucnak':np.max(res_tucnak)}
		# Nalezenim maxima je zjisteno o jaky druh potisku fixy se jedna 
		fixa = max(srovnani.items(), key=operator.itemgetter(1))[0]									#najde nejvetsi prislusnost
		

		# Jelikoz je uz je zjisten druh potisku fixy, mohou se nacist prislusna dalsi loga a porovnat je
		if fixa == 'power':
			h_power,w_power = template_power.shape
			
			loc_power = np.where(res_power==np.max(res_power))
			pt_power = loc_power[::-1]

			matched_power = copy.deepcopy(img[pt_power[1][0]:pt_power[1][0]+h_power,pt_power[0][0]:pt_power[0][0]+w_power,:])
			cv2.rectangle(img[:,:,:],pt_power, (pt_power[0]+w_power, pt_power[1]+h_power), (255,255,255), 2)

			#stabilo
			template_powerstabilo = cv2.imread('ref_powerstabilo.png',0)
			h_powerstabilo,w_powerstabilo = template_powerstabilo.shape
			res_powerstabilo = cv2.matchTemplate(img[:,:,Best], template_powerstabilo, cv2.TM_CCOEFF_NORMED)
			loc_powerstabilo = np.where(res_powerstabilo==np.max(res_powerstabilo))
			pt_powerstabilo = loc_powerstabilo[::-1]

			matched_powerstabilo = copy.deepcopy(img[pt_powerstabilo[1][0]:pt_powerstabilo[1][0]+h_powerstabilo,pt_powerstabilo[0][0]:pt_powerstabilo[0][0]+w_powerstabilo,:])
			cv2.rectangle(img[:,:,:],pt_powerstabilo, (pt_powerstabilo[0]+w_powerstabilo, pt_powerstabilo[1]+h_powerstabilo), (255,255,255), 2)

			#CE
			template_power_CE = cv2.imread('ref_power_CE.png',0)
			h_power_CE,w_power_CE = template_power_CE.shape
			res_power_CE = cv2.matchTemplate(img[:,:,Best], template_power_CE, cv2.TM_CCOEFF_NORMED)
			loc_power_CE = np.where(res_power_CE==np.max(res_power_CE))
			pt_power_CE = loc_power_CE[::-1]

			# Vykresleni ramecku templatu do puvodniho obrazku
			matched_power_CE = copy.deepcopy(img[pt_power_CE[1][0]:pt_power_CE[1][0]+h_power_CE,pt_power_CE[0][0]:pt_power_CE[0][0]+w_power_CE,:])
			cv2.rectangle(img[:,:,:],pt_power_CE, (pt_power_CE[0]+w_power_CE, pt_power_CE[1]+h_power_CE), (255,255,255), 2)

		elif fixa == 'trio':
			h_trio,w_trio = template_trio.shape

			loc_trio = np.where(res_trio==np.max(res_trio))
			pt_trio = loc_trio[::-1]

			matched_trio = copy.deepcopy(img[pt_trio[1][0]:pt_trio[1][0]+h_trio,pt_trio[0][0]:pt_trio[0][0]+w_trio,:])
			cv2.rectangle(img[:,:,:],pt_trio, (pt_trio[0]+w_trio, pt_trio[1]+h_trio), (255,255,255), 2)
			
			#stabilo
			template_triostabilo = cv2.imread('ref_triostabilo.png',0)
			h_triostabilo,w_triostabilo = template_triostabilo.shape
			res_triostabilo = cv2.matchTemplate(img[:,:,Best], template_triostabilo, cv2.TM_CCOEFF_NORMED)
			loc_triostabilo = np.where(res_triostabilo==np.max(res_triostabilo))
			pt_triostabilo = loc_triostabilo[::-1]

			matched_triostabilo = copy.deepcopy(img[pt_triostabilo[1][0]:pt_triostabilo[1][0]+h_triostabilo,pt_triostabilo[0][0]:pt_triostabilo[0][0]+w_triostabilo,:])
			cv2.rectangle(img[:,:,:],pt_triostabilo, (pt_triostabilo[0]+w_triostabilo, pt_triostabilo[1]+h_triostabilo), (255,255,255), 2)

			# CE
			template_trio_CE = cv2.imread('ref_trio_CE.png',0)
			h_trio_CE,w_trio_CE = template_trio_CE.shape
			res_trio_CE = cv2.matchTemplate(img[:,:,Best], template_trio_CE, cv2.TM_CCOEFF_NORMED)
			loc_trio_CE = np.where(res_trio_CE==np.max(res_trio_CE))
			pt_trio_CE = loc_trio_CE[::-1]

			matched_trio_CE = copy.deepcopy(img[pt_trio_CE[1][0]:pt_trio_CE[1][0]+h_trio_CE,pt_trio_CE[0][0]:pt_trio_CE[0][0]+w_trio_CE,:])
			cv2.rectangle(img[:,:,:],pt_trio_CE, (pt_trio_CE[0]+w_trio_CE, pt_trio_CE[1]+h_trio_CE), (255,255,255), 2)



		elif fixa == 'cina':

			h_cina,w_cina = template_cina.shape
			

			loc_cina = np.where(res_cina==np.max(res_cina))

			pt_cina = loc_cina[::-1]

			matched_cina = copy.deepcopy(img[pt_cina[1][0]:pt_cina[1][0]+h_cina,pt_cina[0][0]:pt_cina[0][0]+w_cina,:])
			cv2.rectangle(img[:,:,:],pt_cina, (pt_cina[0]+w_cina, pt_cina[1]+h_cina), (0,0,0), 2)



			template_czech = cv2.imread('ref_czech.png',0)
			h_czech,w_czech = template_czech.shape
			
			res_czech = cv2.matchTemplate(img[:,:,Best], template_czech, cv2.TM_CCOEFF_NORMED)
			loc_czech = np.where(res_czech==np.max(res_czech))

			pt_czech = loc_czech[::-1]

			matched_czech = copy.deepcopy(img[pt_czech[1][0]:pt_czech[1][0]+h_czech,pt_czech[0][0]:pt_czech[0][0]+w_czech,:])
			cv2.rectangle(img[:,:,:],pt_czech, (pt_czech[0]+w_czech, pt_czech[1]+h_czech), (0,0,0), 2)



			template_company = cv2.imread('ref_company.png',0)
			h_company,w_company = template_company.shape
			
			res_company = cv2.matchTemplate(img[:,:,Best], template_company, cv2.TM_CCOEFF_NORMED)
			loc_company = np.where(res_company==np.max(res_company))

			pt_company = loc_company[::-1]

			matched_cina2 = copy.deepcopy(img[pt_company[1][0]:pt_company[1][0]+h_company,pt_company[0][0]:pt_company[0][0]+w_company,:])
			cv2.rectangle(img[:,:,:],pt_company, (pt_company[0]+w_company, pt_company[1]+h_company), (0,0,0), 2)



		elif fixa == 'tucnak':
			h_tucnak,w_tucnak = template_tucnak.shape
			

			loc_tucnak = np.where(res_tucnak==np.max(res_tucnak))

			pt_tucnak = loc_tucnak[::-1]

			matched_tucnak = copy.deepcopy(img[pt_tucnak[1][0]:pt_tucnak[1][0]+h_tucnak,pt_tucnak[0][0]:pt_tucnak[0][0]+w_tucnak,:])
			cv2.rectangle(img[:,:,:],pt_tucnak, (pt_tucnak[0]+w_tucnak, pt_tucnak[1]+h_tucnak), (0,0,0), 2)



			template_czech = cv2.imread('ref_czech.png',0)
			h_czech,w_czech = template_czech.shape
			
			res_czech = cv2.matchTemplate(img[:,:,Best], template_czech, cv2.TM_CCOEFF_NORMED)
			loc_czech = np.where(res_czech==np.max(res_czech))

			pt_czech = loc_czech[::-1]

			matched_czech = copy.deepcopy(img[pt_czech[1][0]:pt_czech[1][0]+h_czech,pt_czech[0][0]:pt_czech[0][0]+w_czech,:])
			cv2.rectangle(img[:,:,:],pt_czech, (pt_czech[0]+w_czech, pt_czech[1]+h_czech), (0,0,0), 2)



		plt.imshow(img[:,:,Best],cmap='gray',vmin=0,vmax=255)
		plt.title("Nahled nalezenych regionu - zkontroluj")
		plt.show()


		print("Ktere chces ulozit?") 			#Vyber vyrezu pro ulozeni (muze se napr. stat, ze je nalezen vyrez, ktery neodpovida realite - nechceme ho ulozit)
		# 1 - stabilo, resp.washable; 2 - power, trio, resp. czech, 3 - ce, resp. company
		# Zadavani bez ohledu na poradi napr. 213 pro ulozeni vsech tri vyrezu, pro ulozeni druheho a tretiho napr. 23

		ktere = input()
		ktere = "{:<3}".format(ktere)

		if (ktere[0] == '1' or ktere[1] == '1' or ktere[2] == '1'):
			prvni = True
		if (ktere[0] == '2' or ktere[1] == '2' or ktere[2] == '2'):
			druha = True
		if (ktere[0] == '3' or ktere[1] == '3' or ktere[2] == '3'):
			treti = True

		# Vypis ukladanych objektu
		print("\nUkladam: ")
		print("[1]",prvni)
		print("[2]",druha)
		print("[3]",treti)

		# Pokud je ulozen nektery z vyrezu, uklada se i puvodni seskladany RGB snimek
		if prvni == True or druha == True or treti == True:
			os.chdir("./apuvodni")
			nazev = pridej_index("index_puvodni.txt","img")					# Pridava index k nazvu snimku
			ulozit = np.uint8(pars)
			ulozit = Image.fromarray(ulozit)
			ulozit.save(nazev)
			os.chdir("../")

		# Dle typu vzoru potisku se jednotlive vyrezy ulozi do prislusnych adresaru

		if fixa == 'power':
			if prvni == True:
				os.chdir("./stabilo_power")
				ulozit = np.uint8(matched_powerstabilo)
				for i in range(3):
					nazev = pridej_index("index_stabilo.txt","stab_p")		

					(Image.fromarray(ulozit[:,:,i])).save(nazev)
				os.chdir("../")

			if druha == True:
				os.chdir("./power")
				ulozit = np.uint8(matched_power)
				for i in range(3):
					nazev = pridej_index("index_power.txt","pow")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)
				os.chdir("../")

			if treti == True:
				os.chdir("./ce_power")
				ulozit = np.uint8(matched_power_CE)
				for i in range(3):
					nazev = pridej_index("index_ce_power.txt","ce_pow")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)
				os.chdir("../")


		if fixa == 'trio':
			if prvni == True:
				os.chdir("./stabilo_trio")
				ulozit = np.uint8(matched_triostabilo)
				for i in range(3):
					nazev = pridej_index("index_stabilo.txt","stab_t")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)
				os.chdir("../")

			if druha == True:
				os.chdir("./trio")
				ulozit = np.uint8(matched_trio)
				for i in range(3):
					nazev = pridej_index("index_trio.txt","trio")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)
				os.chdir("../")

			if treti == True:
				os.chdir("./ce_trio")
				ulozit = np.uint8(matched_trio_CE)
				for i in range(3):
					nazev = pridej_index("index_ce_trio.txt","ce_trio")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)
				os.chdir("../")

		if fixa == 'cina':
			if prvni == True:
				os.chdir("./washable")
				ulozit = np.uint8(matched_cina)
				for i in range(3):
					nazev = pridej_index("index_washable.txt","wash")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)

				os.chdir("../")

			if druha == True:
				os.chdir("./czech")
				ulozit = np.uint8(matched_czech)
				for i in range(3):
					nazev = pridej_index("index_czech.txt","czech")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)

				os.chdir("../")

			if treti == True:
				os.chdir("./company")
				ulozit = np.uint8(matched_cina2)
				for i in range(3):
					nazev = pridej_index("index_company.txt","comp")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)

				os.chdir("../")

		if fixa == 'tucnak':
			if prvni == True:
				os.chdir("./tucnak")
				ulozit = np.uint8(matched_tucnak)
				for i in range(3):
					nazev = pridej_index("index_tucnak.txt","tucnak")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)

				os.chdir("../")

			if druha == True:
				os.chdir("./czech")
				ulozit = np.uint8(matched_czech)
				for i in range(3):
					nazev = pridej_index("index_czech.txt","czech")

					(Image.fromarray(ulozit[:,:,i])).save(nazev)

				os.chdir("../")



