import numpy as np
import cv2
import time
from PIL import Image

from paddle.inference import Config
from paddle.inference import create_predictor

from handwriting_util import ResizeImage,CropImage,ToTensor,NormalizeImage

# paddle模型文件
MODEL_FILE = "output_2/inference/inference.pdmodel"
MODEL_PARAM = "output_2/inference/inference.pdiparams"
# 分配显存
GPU_MEM = 8000

# 前k个候选
top_k = 5

# 测试文件
img_file = "data/test/test5.png"
#标签文件
label_file = "data/trade/labels.txt"

# 建立Paddle预测器
def create_paddle_perdictor():
    config=Config(MODEL_FILE, MODEL_PARAM)
    config.enable_use_gpu(GPU_MEM, 0)
    config.set_cpu_math_library_num_threads(10)
    config.disable_glog_info()
    config.switch_ir_optim(True)
    config.enable_memory_optim()
    # use zero copy
    config.switch_use_feed_fetch_ops(False)
    predictor = create_predictor(config)
    return predictor

# 图像预处理，转为Tensor张量
def preprocess(img):
    resize_op = ResizeImage(resize_short=256)
    img = resize_op(img)
    crop_op = CropImage(size=(224, 224))
    img = crop_op(img)
    img_mean = [0.485, 0.456, 0.406]
    img_std = [0.229, 0.224, 0.225]
    img_scale = 1.0 / 255.0
    normalize_op = NormalizeImage(
        scale=img_scale, mean=img_mean, std=img_std)
    img = normalize_op(img)
    tensor_op = ToTensor()
    img = tensor_op(img)
    return img

# 结果分析处理
def postprocess(output):
    output = output.flatten()
    classes = np.argpartition(output, -top_k)[-top_k:]
    classes = classes[np.argsort(-output[classes])]
    scores = output[classes]
    return classes, scores

# 读取标签文件
def readLabels():
    index = 0
    labels=[]
    with open(label_file) as file_obj:
        for line in file_obj:
            labels.append(line.strip())
            index=index+1
    return labels

def dealTransform():
    img=Image.open('data/test/test.png') 
    img=transparence2white(img)  # 将图片传入，改变背景色后，返回
    img=img.resize((64,64))
    img.save(img_file)  # 保存图片

def transparence2white(img):
#     img=img.convert('RGBA')  # 此步骤是将图像转为灰度(RGBA表示4x8位像素，带透明度掩模的真彩色；CMYK为4x8位像素，分色等)，可以省略
    sp=img.size
    width=sp[0]
    height=sp[1]
    print(sp)
    for yh in range(height):
        for xw in range(width):
            dot=(xw,yh)
            color_d=img.getpixel(dot)  # 与cv2不同的是，这里需要用getpixel方法来获取维度数据
            if(color_d[3]==0):
                color_d=(255,255,255,255)
                img.putpixel(dot,color_d)  # 赋值的方法是通过putpixel
    return img

# 减少图像白边
def cutImg(img):
    dst = 255- img
    gray = cv2.cvtColor(dst,cv2.COLOR_BGR2GRAY)  
    ret, binary = cv2.threshold(gray,127,255,cv2.THRESH_BINARY)  
    
    contours, hierarchy = cv2.findContours(binary,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

    minx = 65523
    miny = 65523
    maxx = 0
    maxy = 0

    for contour in contours:
        boundRect = cv2.boundingRect(contour)
        (x,y,w,h)=boundRect
        if x<minx:
            minx=x
        if y<miny:
            miny=y
        if x+w>maxx:
            maxx=x+w
        if y+h>maxy:
            maxy=y+h
    cutw=maxx-minx
    cuth=maxy-miny
    # w>h
    # if cutw>cuth:
    #     cuty1=int((maxy+miny)/2-cutw/2)
    #     cuty2=int((maxy+miny)/2+cutw/2)
    #     img = img[cuty1:cuty2,minx:maxx]
    # if cutw<cuth:
    #     cutx1=int((maxx+minx)/2-cuth/2)
    #     cutx2=int((maxx+minx)/2+cuth/2)
    #     img = img[miny:maxy,cutx1:cutx2] 
    img = img[miny:maxy,minx:maxx]
    return img

all_label = readLabels()
# dealTransform()

predictor = create_paddle_perdictor();
input_names = predictor.get_input_names()
input_tensor = predictor.get_input_handle(input_names[0])

output_names = predictor.get_output_names()
output_tensor = predictor.get_output_handle(output_names[0])

img = cv2.imread(img_file)[:, :, ::-1]
img=cutImg(img)
cv2.imwrite("dealcut.png",img)
img = cv2.resize(img, (64,64))
inputs = preprocess(img)
inputs = np.expand_dims(
    inputs, axis=0).repeat(
        1, axis=0).copy()
input_tensor.copy_from_cpu(inputs)
predictor.run()
output = output_tensor.copy_to_cpu()
classes, scores = postprocess(output)
print(classes)
for index in classes:
    print(all_label[index])