import base64
import json
import re
import sys
import time
import paddlehub as hub
import cv2
from PIL import Image
import logging
import numpy as np

use_model = "mobile"
use_gpu = False

# img_data = "./1.jpg"
img_data = "D:/Project/PaddleOCR-release-2.0/train_data/bidi_data/orgs_data/train/text_304.jpg"
np_images = [cv2.imread(img_data)]


def only_recognize():
    ocr = hub.Module(name="chinese_text_detection_db_"+use_model)
    results = ocr


def only_detect():
    ocr = hub.Module(name="chinese_text_detection_db_"+use_model)
    results = ocr.detect_text(
        images=np_images,           # 图片数据，ndarray.shape 为 [H, W, C]，BGR格式；
        use_gpu=use_gpu,            # 是否使用 GPU；若使用GPU，请先设置CUDA_VISIBLE_DEVICES环境变量
        output_dir='../ocr_result',  # 图片的保存路径，默认设为 ocr_result；
        visualization=True,         # 是否将识别结果保存为图片文件；
        box_thresh=0.5              # 检测文本框置信度的阈值；
        )                           # 识别中文文本置信度的阈值；
    for result in results:
        print(results)


def detect_and_recognize():
    ocr = hub.Module(name="chinese_ocr_db_crnn_"+use_model)
    results = ocr.recognize_text(
        images=np_images,         # 图片数据，ndarray.shape 为 [H, W, C]，BGR格式；
        use_gpu=use_gpu,            # 是否使用 GPU；若使用GPU，请先设置CUDA_VISIBLE_DEVICES环境变量
        output_dir='../ocr_result',  # 图片的保存路径，默认设为 ocr_result；
        visualization=True,       # 是否将识别结果保存为图片文件；
        box_thresh=0.5,           # 检测文本框置信度的阈值；
        text_thresh=0.0)          # 识别中文文本置信度的阈值；

    for result in results:
        data = result['data']
        save_path = result['save_path']
        for infomation in data:
            if infomation['text'] == "":
                print("no text")
                continue
            print('text: ', infomation['text'], '\nconfidence: ',
                  infomation['confidence'], '\ntext_box_position: ',
                  infomation['text_box_position'])


def image_bigger():
    img = cv2.imread("./1.jpg", -1)
    height, weight = img.shape[:2]
    print(height, weight)
    fx = 0.7
    fy = 1.2
    enlarge = cv2.resize(img, (0, 0), fx=fx, fy=fy, interpolation=cv2.INTER_CUBIC)
    print(enlarge.shape[:2])
    cv2.imwrite("./1_1.jpg", enlarge)
    # img.save("./1_1.jpg", "jpeg")


if __name__ == '__main__':
    only_detect()
    # image_bigger()
    # detect_and_recognize()
