import os

import time

import argparse

import logging as log

import numpy as np

import keras

from keras.models import load_model

from PIL import Image,ImageDraw

from epoint import EPoint

import exp_config as cfg

from kerokero.k_util import averageDistance

keras.losses.averageDistance=averageDistance

keras.metrics.averageDistance=averageDistance

model=load_model(cfg.dochiModel)

def locateGrid(img):

	t1=time.time()

	(width,height)=img.size

	normedImg=img.convert("L").resize((224,224),resample=Image.BILINEAR)

	npImg=np.array(normedImg.getdata()).reshape((224,224,1)).astype(np.float32)

	npImg=npImg/128-1

	label=model.predict(np.reshape(npImg,(1,224,224,1)))

	points=[]

	for i in range(4):

		points.append(EPoint((label[0][i][0]+1)*(width/2),(label[0][i][1]+1)*(height/2)))

	t=time.time()-t1

	log.info("grid located in {0:.3}s".format(t))

	return points

if __name__=="__main__":

	parser=argparse.ArgumentParser()

	parser.add_argument("-i","--input",nargs="+")

	parser.add_argument("-o","--output_dir",required=True)

	args=parser.parse_args()

	for image_path in args.input:

		image=Image.open(image_path)

		points=locateGrid(image)

		drawer=ImageDraw.Draw(image)

		for p in points:

			drawer.ellipse((p.x-2,p.y-2,p.x+2,p.y+2),fill="#00ff00")

		image.save(os.path.join(args.output_dir,os.path.basename(image_path)))

dochiModel=cfg.get("dochiModel")