diff --git a/exp/hakugen.py b/exp/hakugen.py new file mode 100644 --- /dev/null +++ b/exp/hakugen.py @@ -0,0 +1,49 @@ +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.hakugenModel) + + +def locateGrid(img): + t1=time.time() + (width,height)=img.size + normedImg=img.convert("L") + 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)))