from . import preprocess
from . import sheetocr
import time
import os
import cv2

# sheetpath = r'C:\Users\Administrator\Desktop\sheet'  # 预处理前的试卷目录
# testpath = r'C:\Users\Administrator\Desktop\test'  # 预处理后的试卷目录
# resultpath = r'C:\Users\Administrator\Desktop\result'  # 结果生成目录

# Parameter Sets
langs = {'ce': '-l chi_sim+eng', 'ec': '-l eng+chi_sim', 'c': '-l chi_sim', 'e': '-l eng', 'eq': '-l eng+equ'}  # 语言选项
psms = {'block': '--psm 6', 'default': '--psm 3'}  # , '_line': ' --psm 7'}  # Page segmentation modes
# oems = {'legacy': '--oem 0', 'lstm': '--oem 1', 'lstm+legacy': '--oem 2'}  # OCR Engine modes
langs_py = {'ce': 'chi_sim+eng', 'ec': 'eng+chi_sim', 'c': 'chi_sim', 'e': 'eng', 'eq': 'eng+equ'}  # 语言选项
psms_py = {'block': '--psm 6', 'default': '--psm 3'}  # , '_line': ' --psm 7'}  # Page segmentation modes
scales = (0, 0.5, 2)
dilates = (0, 1, 3, 5)
blurs = (0, 1, 3, 5, 7)


# 用默认最佳参数处理图片， 返回文本
def ocr_py(picture, lang='ce', psm='block', scale=0, dilate=1, blur=5):
    image = preprocess.preprocess(picture, scale=scale, dilate=dilate, blur=blur)
    words = sheetocr.sheetocr_py(image, lang=langs_py[lang], psm=psms_py[psm])
    return words


# 用默认最佳参数处理图片， 返回文本文件
def ocr(picture, output, lang='ce', psm='block', scale=0, dilate=1, blur=5):
    image = preprocess.preprocess(picture, scale=scale, dilate=dilate, blur=blur)
    cv2.imwrite('tmp_pic', image)
    sheetocr.sheetocr('tmp_pic', output, lang=langs[lang], psm=psms[psm])
    os.remove('tmp_pic')


# 测试最佳参数
def test_parameters(picture_path, output=0):
    start = time.time()
    for root, dirs, files in os.walk(picture_path):
        for file in files:
            picture = os.path.join(root, file)
            if output == 0:  # 屏幕显示
                for s in scales:
                    for d in dilates:
                        for b in blurs:
                            print('Parameters:' + 's' + str(s) + 'd' + str(d) + 'b' + str(b) + '\n')
                            words = ocr_py(picture, scale=s, dilate=d, blur=b)
                            print(words)
            else:  # 输出到路径为output的文件中
                for s in scales:
                    for d in dilates:
                        for b in blurs:
                            save = file + 's' + str(s) + 'd' + str(d) + 'b' + str(b)
                            save = os.path.join(output, save)
                            with open(save, 'r', encoding='UTF-8') as f:
                                words = ocr_py(picture, scale=s, dilate=d, blur=b)
                                f.write(words)
    end = time.time()
    print('running time:', end - start, 's')


# test_parameters(sheetpath)
# print('OCR done!\n')