# -*- coding: utf-8 -*-
"""
Created on Sat Jul 11 21:24:06 2020

@author: MANAV

"""

import tkinter as tk
from tkinter import *
import webbrowser
from tkinter import ttk
from tkinter import font, colorchooser, filedialog, messagebox
from tkinter import Toplevel
from tkinter.filedialog import asksaveasfile
from PIL import Image, ImageTk
import os
import numpy as np
import cv2 as cv
import os.path

numpy_file = np.load('./pts_in_hull.npy')
Caffe_net = cv.dnn.readNetFromCaffe("./colorization_deploy_v2.prototxt", "./colorization_release_v2.caffemodel")
numpy_file = numpy_file.transpose().reshape(2, 313, 1, 1)


        
class Window(Frame):
    def __init__(self, master=None):
        Frame.__init__(self, master)
        
        self.master = master
 
        self.pos = []
        self.master.title("Automatic Image Colorization")
        self.pack(fill=BOTH, expand=1)

        menu = Menu(self.master)
        self.master.config(menu=menu)

        file = tk.Menu(menu,tearoff=0)
        

        file.add_command(label="New",  compound=tk.LEFT )
        file.add_command(label="Upload Image",  compound=tk.LEFT, command=self.uploadImage)
        file.add_command(label="Color Image", command=self.color)
        file.add_command(label="Save", compound=tk.LEFT, command=self.save)
        file.add_separator()
        file.add_command(label="Exit", compound=tk.LEFT, command=self.exitProgram)
        menu.add_cascade(label="File", menu=file)

        editMenu = Menu(menu,tearoff=0 )
        editMenu.add_command(label="Undo")
        editMenu.add_command(label="Redo")
         
        menu.add_cascade(label="Edit", menu=editMenu )

        optionMenu = Menu(menu,tearoff=0)
        
        optionMenu.add_command(label="change color", command=self.col)
        menu.add_cascade(label="Option", menu=optionMenu)

        helpMenu = Menu(menu,tearoff=0)
        helpMenu.add_command(label="About", command=self.About)
        helpMenu.add_command(label="Examples", command=self.example)
        menu.add_cascade(label="Help", menu=helpMenu)

        quitMenu = Menu(menu,tearoff=0)
        exit_icon = tk.PhotoImage(file='icons2/exit.png')
        quitMenu.add_command(label="Exit", command=self.exitProgram)
        menu.add_cascade(label="Quit", menu=quitMenu)
    
       
        self.canvas = tk.Canvas(self)
        self.canvas.pack(fill=tk.BOTH, expand=True)
        self.image = None
        self.image2 = None
        label1=Label(self,image=img)
        label1.image=img
        label1.place(x=350,y=430)
 
    def col(self):
        color='green'

    def new(self):
        new.pack(fill="both", expand=0)

    def save(self): 
        filename = filedialog.asksaveasfile(initialdir = "/", defaultextension=".jpg", title = "Save file",filetypes = (('JPG', '*.jpg'),('PNG', '*.png')))  
        print (filename)

    def example(self):
         webbrowser.open("C:/Users/ab/Desktop/models/Examples.rtf")

    def About(self):
        webbrowser.open("C:/Users/ab/Desktop/models/About.txt")

    def exitProgram(self):
        self.master.destroy()

    def hide_all_frames(self):
        file_new_frame.pack_forget()


   
        
    def uploadImage(self):
        filename = filedialog.askopenfilename(initialdir=os.getcwd())
        if not filename:
            return
        load = Image.open(filename)
        load = load.resize((480, 360), Image.ANTIALIAS)
        if self.image is None:
            w, h = load.size
            width, height = root.winfo_width(), root.winfo_height()
            self.render = ImageTk.PhotoImage(load)
            self.image = self.canvas.create_image((w / 2, h / 2), image=self.render)
           
        else:
            self.canvas.delete(self.image3)
            w, h = load.size
            width, height = root.winfo_screenmmwidth(), root.winfo_screenheight()
            self.render2 = ImageTk.PhotoImage(load)
            self.image2 = self.canvas.create_image((w / 2, h / 2), image=self.render2)
        frame = cv.imread(filename)
    
        Caffe_net.getLayer(Caffe_net.getLayerId('class8_ab')).blobs = [numpy_file.astype(np.float32)]
        Caffe_net.getLayer(Caffe_net.getLayerId('conv8_313_rh')).blobs = [np.full([1, 313], 2.606, np.float32)]
        input_width = 224
        input_height = 224
        rgb_img = (frame[:,:,[2, 1, 0]] * 1.0 / 255).astype(np.float32)
        lab_img = cv.cvtColor(rgb_img, cv.COLOR_RGB2Lab)
        l_channel = lab_img[:,:,0] 
        l_channel_resize = cv.resize(l_channel, (input_width, input_height)) 
        l_channel_resize -= 50 
        Caffe_net.setInput(cv.dnn.blobFromImage(l_channel_resize))
        ab_channel = Caffe_net.forward()[0,:,:,:].transpose((1,2,0)) 
        (original_height,original_width) = rgb_img.shape[:2] 
        ab_channel_us = cv.resize(ab_channel, (original_width, original_height))
        lab_output = np.concatenate((l_channel[:,:,np.newaxis],ab_channel_us),axis=2) 
        bgr_output = np.clip(cv.cvtColor(lab_output, cv.COLOR_Lab2BGR), 0, 1)
  
        cv.imwrite("./Output/result.png", (bgr_output*255).astype(np.uint8))

    def color(self):
        load = Image.open("./Output/result.png")
        load = load.resize((480, 360), Image.ANTIALIAS)
        if self.image is None:
            w, h = load.size
            self.render = ImageTk.PhotoImage(load)
            self.image = self.canvas.create_image((w / 2, h/2), image=self.render)
            root.geometry("%dx%d" % (w, h))
        else:
            w, h = load.size
            width, height = root.winfo_screenmmwidth(), root.winfo_screenheight()
            self.render3 = ImageTk.PhotoImage(load)
            self.image3 = self.canvas.create_image((w / 2, h / 2), image=self.render3)
            self.canvas.move(self.image3, 500, 0)



root = tk.Tk()
root.geometry("%dx%d" % (980, 600))
root.title("B&W Image Colorization GUI") 
img = ImageTk.PhotoImage(Image.open("logo2.png"))
app = Window(root)
app.pack(fill=tk.BOTH, expand=1)
root.resizable(False, False)
root.mainloop()
