
# Skript pro augmentaci galerie dat. Mozne augmentovat ruznymi sadami vzoru, popr. rotaci obrazku o zvoleny uhel

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
import os

from scipy import ndimage


def pridej_index(soubor, nazev):                # metoda pro pridavani indexu k nazvu, index je cten z textoveho souboru
    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)) + "_aug" + ".png"
    return vystupninazev

def vyberkanalu(img):                           # metoda pro vyber nejkontrastnejsiho kanalu rgb snimku
    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 = 40                                   # hystereze, ktera omezuje velikost hlavniho peaku histogramu

    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,:])   # nalezeni indexu hlavniho peaku

        if idxmax[vyska]+hyst > 255:            # nulovani histogramu v okoli definovanem hysterezi, aby se mohl najit druhy nejvetsi peak
            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,:])  # nalezeni druheho nejvetsiho peaku histogramu

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

    ratio = abs(idxmax-idxmax2)                 # vypocet vzdalenosti mezi peaky
    Best = np.argmax(ratio)                     # nejvetsi vzdalenst znamena nejvetsi kontrast, ktery chceme

    return Best                                 # vracen je index nejlepsiho kanalu rgb snimku



def natoc(img):                                         # metoda pro natoceni vstupniho obrazku dle nahodne generovane hodnoty v danem rozsahu 
    rozsah = 2.5                                        #
    while True:
        uhel = 2*rozsah*np.random.rand()-rozsah
        if abs(uhel)>0.5:
            break                                       # kdyz nahodne generovany uhel splnuje pozadavek je cyklus ukoncen
    img = ndimage.rotate(img, uhel, reshape=False)      # rotace snimku
    img = np.uint8(img)
    return img                                          # metoda vraci zrotovany snimek




if __name__ == "__main__":                          # funkce main

    category = 'nok'                                # kategorie, ktera je spustenim resena - bud ok nebo nok

    # nacteni jednotlivych referencnich obrazku
    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)


    for index in range(650):                                        # cyklus pres vsechny indexy snimku + rezerva navic

        print(index)

        nazev = 'img_' + str(index).zfill(4) + '.png'               # synteza nazvu

        img =cv2.imread('apuvodni/' + category + '/' + nazev,1)     # snimek je nacten, pokud snimek o danem nazvu neni do promenne img je ulozen datovy typ notype<>



        try:                                                        # snimek je rotovan, pokud se dany snimek nenacetl a chceme s nim provadet otoceni je vyhozena vyjimka
            #img = natoc(img)
            img = np.uint8(img)
        except (TypeError, ValueError):
            print(index)
            continue                                                # vyjimka je resena prejitim na dalsi otacku cyklu

        Best = vyberkanalu(img)                                     # volani metody pro vyber nejlepsiho kanalu daneho snimku



        # je provedeno srovnani snimku se vzorem loga
        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 je ulozena k identifikatoru daneho potisku propisky nejvyssi hodnota prislusnosti daneho vzoru
        # zjistenim maxima slovniku je urceno o jaky typ potisku fixy se jedna a podle toho se dale voli dalsi loga
        srovnani = {'power':np.max(res_power), 'trio':np.max(res_trio), 'cina':np.max(res_cina), 'tucnak':np.max(res_tucnak)}
        fixa = max(srovnani.items(), key=operator.itemgetter(1))[0]                                 #najde nejvetsi prislusnost
        
        # srovnani dalsich vzoru log se snimkem
        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,:])
            

            # 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,:])
            

            # 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]

            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,:])
            
            # vytisknuti ramecku kolem loga
            cv2.rectangle(img[:,:,:],pt_power, (pt_power[0]+w_power, pt_power[1]+h_power), (255,255,255), 2)
            cv2.rectangle(img[:,:,:],pt_powerstabilo, (pt_powerstabilo[0]+w_powerstabilo, pt_powerstabilo[1]+h_powerstabilo), (255,255,255), 2)
            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]

            # 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]

            # 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 = copy.deepcopy(img[pt_trio[1][0]:pt_trio[1][0]+h_trio,pt_trio[0][0]:pt_trio[0][0]+w_trio,:])
            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,:])
            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, (pt_trio[0]+w_trio, pt_trio[1]+h_trio), (255,255,255), 2)
            cv2.rectangle(img[:,:,:],pt_triostabilo, (pt_triostabilo[0]+w_triostabilo, pt_triostabilo[1]+h_triostabilo), (255,255,255), 2)
            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,:])

            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,:])

            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_cina, (pt_cina[0]+w_cina, pt_cina[1]+h_cina), (0,0,0), 2)
            cv2.rectangle(img[:,:,:],pt_czech, (pt_czech[0]+w_czech, pt_czech[1]+h_czech), (0,0,0), 2)
            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,:])


            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_tucnak, (pt_tucnak[0]+w_tucnak, pt_tucnak[1]+h_tucnak), (0,0,0), 2)
            cv2.rectangle(img[:,:,:],pt_czech, (pt_czech[0]+w_czech, pt_czech[1]+h_czech), (0,0,0), 2)



        # ulozeni jednotlivych vzoru do prislusnych mist adresarove struktury
        
        if fixa == 'power':
            os.chdir("./stabilo_power")                                 # zmena aktualni cesty do prislusneho adresare
            ulozit = np.uint8(matched_powerstabilo)
            for i in range(3):                                          # cyklus pro ulozeni vsech tri rgb kanalu
                nazev = pridej_index("index_stabilo.txt","stab_p")      # generovani nazvu s indexem

                (Image.fromarray(ulozit[:,:,i])).save(nazev)            # ulozeni
            os.chdir("../")                                             # navrat do puvodniho adresare, tento proces je proveden pro vsechna loga daneho typu fixy

            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("../")

            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':
            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("../")

            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("../")

            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':
            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("../")

            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("../")

            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':
            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("../")

            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("../")

                