648 lines
24 KiB
Python
648 lines
24 KiB
Python
import os
|
|
import argparse
|
|
import lzma
|
|
import pickle
|
|
from datetime import datetime
|
|
import numpy as np
|
|
import PIL.Image
|
|
|
|
import torch
|
|
|
|
import easyocr
|
|
|
|
# %%
|
|
def count_parameters(model):
|
|
return sum([param.numel() for param in model.parameters()])
|
|
|
|
def get_weight_norm(model):
|
|
with torch.no_grad():
|
|
return sum([param.norm() for param in model.parameters()]).cpu().item()
|
|
|
|
def replace(list_in, indices, values):
|
|
if not isinstance(indices, list):
|
|
indices = [indices]
|
|
if not isinstance(values, list):
|
|
values = [values]
|
|
assert len(indices) == len(values)
|
|
|
|
list_out = list_in.copy()
|
|
for index, value in zip(indices, values):
|
|
list_out[index] = value
|
|
|
|
return list_out
|
|
|
|
def get_easyocr(language):
|
|
if not isinstance(language, list):
|
|
language = [language]
|
|
return easyocr.Reader(language)
|
|
|
|
# %%
|
|
def main(args):
|
|
|
|
if args.output is None:
|
|
args.output = "EasyOcrUnitTestPackage_{}.pickle".format(datetime.now().strftime("%Y%m%dT%H%M"))
|
|
|
|
if args.data_dir is None:
|
|
data_dir = "./examples"
|
|
else:
|
|
data_dir = args.data_dir
|
|
|
|
image_preprocess = {
|
|
'english.png':{
|
|
"tiny": [540, 420, 690, 470],
|
|
"mini": [260, 90, 605, 160],
|
|
"small": [243, 234, 636, 360]
|
|
},
|
|
'french.jpg':{
|
|
"tiny": [184, 615, 425, 732]
|
|
},
|
|
'chinese.jpg':{
|
|
"tiny": [181, 78, 469, 157]
|
|
},
|
|
'korean.png':{
|
|
"tiny": [130, 84, 285, 180]
|
|
}
|
|
}
|
|
|
|
|
|
if any([file not in os.listdir(data_dir) for file in image_preprocess.keys()]):
|
|
raise FileNotFoundError("Cannot find {} in {}.").format(', '.join([file for file in image_preprocess.keys() if file not in os.listdir(data_dir)], data_dir))
|
|
|
|
easyocr_config = {"main_language": 'en'}
|
|
|
|
ocr = get_easyocr(easyocr_config["main_language"])
|
|
|
|
images = {os.path.splitext(file)[0]: {
|
|
key: np.asarray(PIL.Image.open(os.path.join(data_dir, file)).crop(crop_box))[:,:,::-1] for (key,crop_box) in page.items()
|
|
} for (file,page) in image_preprocess.items()}
|
|
|
|
|
|
english_mini_bgr, english_mini_gray = easyocr.utils.reformat_input(images['english']['mini'])
|
|
english_small_bgr, english_small_gray = easyocr.utils.reformat_input(images['english']['small'])
|
|
|
|
|
|
model_init_test = {'test01': {
|
|
'description': "Counting parameters of detector module.",
|
|
"method": "unit_test.count_parameters",
|
|
'input': ["unit_test.easyocr.ocr.detector"],
|
|
'output': count_parameters(ocr.detector),
|
|
'severity': "Error"
|
|
},
|
|
'test02': {
|
|
'description': "Calculating total norm of parameters in detector module.",
|
|
"method": "unit_test.get_weight_norm",
|
|
'input': ["unit_test.easyocr.ocr.detector"],
|
|
'output': get_weight_norm(ocr.detector),
|
|
'severity': "Warning"
|
|
},
|
|
'test03': {
|
|
'description': "Counting parameters of recognition module.",
|
|
"method": "unit_test.count_parameters",
|
|
'input': ["unit_test.easyocr.ocr.recognizer"],
|
|
'output': count_parameters(ocr.recognizer),
|
|
'severity': "Error"
|
|
},
|
|
'test04': {
|
|
'description': "Calculating total norm of parameters in recognition module.",
|
|
"method": "unit_test.get_weight_norm",
|
|
'input': ["unit_test.easyocr.ocr.recognizer"],
|
|
'output': get_weight_norm(ocr.recognizer),
|
|
'severity': "Warning"
|
|
},
|
|
}
|
|
|
|
|
|
get_textbox_test = {}
|
|
|
|
input0 = [ocr.detector,#detector
|
|
english_mini_bgr,#image
|
|
2560,#canvas_size
|
|
1.0,#mag_ratio
|
|
0.7,#text_threshold
|
|
0.4,#link_threshold
|
|
0.4,#low_text
|
|
False, #poly #Fixed
|
|
'cuda', #device #fixed ?
|
|
]
|
|
get_textbox_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.detection.get_textbox",
|
|
'input': replace(input0,
|
|
[0, 1],
|
|
["unit_test.easyocr.ocr.detector",
|
|
"unit_test.inputs.images.english.mini_bgr"
|
|
]),
|
|
'output': easyocr.detection.get_textbox(*input0),
|
|
'severity': "Error"
|
|
}})
|
|
|
|
input0 = [ocr.detector,#detector
|
|
english_mini_bgr,#image
|
|
1280,#canvas_size
|
|
1.2,#mag_ratio
|
|
0.6,#text_threshold
|
|
0.3,#link_threshold
|
|
0.3,#low_text
|
|
False, #poly #Fixed
|
|
'cuda', #device #fixed ?
|
|
]
|
|
|
|
get_textbox_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.detection.get_textbox",
|
|
'input': replace(input0,
|
|
[0, 1],
|
|
["unit_test.easyocr.ocr.detector",
|
|
"unit_test.inputs.images.english.mini_bgr"
|
|
]),
|
|
'output': easyocr.detection.get_textbox(*input0),
|
|
'severity': "Error"
|
|
}})
|
|
|
|
input0 = [ocr.detector,#detector
|
|
english_mini_bgr,#image
|
|
640,#canvas_size
|
|
0.8,#mag_ratio
|
|
0.8,#text_threshold
|
|
0.5,#link_threshold
|
|
0.5,#low_text
|
|
False, #poly #Fixed
|
|
'cuda', #device #fixed ?
|
|
]
|
|
|
|
get_textbox_test.update({'test03': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.detection.get_textbox",
|
|
'input': replace(input0,
|
|
[0, 1],
|
|
["unit_test.easyocr.ocr.detector",
|
|
"unit_test.inputs.images.english.mini_bgr"
|
|
]),
|
|
'output': easyocr.detection.get_textbox(*input0),
|
|
'severity': "Error"
|
|
}})
|
|
|
|
|
|
input0 = [ocr.detector,#detector
|
|
english_mini_bgr,#image
|
|
2560,#canvas_size
|
|
1.0,#mag_ratio
|
|
0.7,#text_threshold
|
|
0.4,#link_threshold
|
|
0.4,#low_text
|
|
False, #poly #Fixed
|
|
'cuda', #device #fixed ?
|
|
]
|
|
output0 = easyocr.detection.get_textbox(*input0)
|
|
polys = output0[0]
|
|
group_text_box_test = {}
|
|
|
|
input_ = [polys,
|
|
0.1,# slope_ths
|
|
0.5,#ycenter_ths
|
|
0.5,#height_ths
|
|
1.0,#width_ths
|
|
0.05,#add_margin
|
|
True#sort_output
|
|
]
|
|
group_text_box_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.utils.group_text_box",
|
|
'input': input_,
|
|
'output': easyocr.utils.group_text_box(*input_),
|
|
'severity': "Error"
|
|
}
|
|
})
|
|
input_ = [polys,
|
|
0.05,# slope_ths
|
|
0.3,#ycenter_ths
|
|
0.3,#height_ths
|
|
0.8,#width_ths
|
|
0.03,#add_margin
|
|
True#sort_output
|
|
]
|
|
group_text_box_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.utils.group_text_box",
|
|
'input': input_,
|
|
'output': easyocr.utils.group_text_box(*input_),
|
|
'severity': "Error"
|
|
}
|
|
})
|
|
input_ = [polys,
|
|
0.12,# slope_ths
|
|
0.7,#ycenter_ths
|
|
0.7,#height_ths
|
|
1.2,#width_ths
|
|
0.1,#add_margin
|
|
True#sort_output
|
|
]
|
|
group_text_box_test.update({'test03': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.utils.group_text_box",
|
|
'input': input_,
|
|
'output': easyocr.utils.group_text_box(*input_),
|
|
'severity': "Error"
|
|
}
|
|
})
|
|
|
|
input0 = [None,
|
|
20, #min_size
|
|
0.7, #text_threshold - fixed
|
|
0.4, #low_text - fixed
|
|
0.4, # link_threshold - fixed
|
|
2560, #canvas_size -fixed
|
|
1., #mag_ratio - fixed
|
|
0.1, #slope_ths - fixed
|
|
0.5, #ycenter_ths - fixed
|
|
0.5, #height_ths - fixed
|
|
0.5, #width_ths - fixed
|
|
0.1, #add_margin - fixed
|
|
True, #reformat - fixed
|
|
None #optimal_num_chars - fixed
|
|
]
|
|
|
|
detect_test = {}
|
|
|
|
input_ = replace(input0, [0,1], [english_mini_bgr, 20])
|
|
detect_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.ocr.detect",
|
|
'input': replace(input_, 0, "unit_test.inputs.images.english.mini_bgr"),
|
|
'output': ocr.detect(*input_),
|
|
'severity': "Error"
|
|
},
|
|
})
|
|
input_ = replace(input0, [0,1], [english_small_bgr, 20])
|
|
detect_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.ocr.detect",
|
|
'input': replace(input_, 0, "unit_test.inputs.images.english.small_bgr"),
|
|
'output': ocr.detect(*input_),
|
|
'severity': "Error"
|
|
},
|
|
})
|
|
input_ = replace(input0, [0,1], [english_small_bgr, 100])
|
|
detect_test.update({'test03': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.ocr.detect",
|
|
'input': replace(input_, 0, "unit_test.inputs.images.english.small_bgr"),
|
|
'output': ocr.detect(*input_),
|
|
'severity': "Error"
|
|
},
|
|
})
|
|
|
|
get_image_list_test = {}
|
|
output0 = ocr.detect(english_small_bgr)
|
|
input0 = [output0[0][0],
|
|
output0[1][0],
|
|
english_small_gray,
|
|
64, #model_height
|
|
True# sort_output
|
|
]
|
|
input_ = replace(input0, 2, "unit_test.inputs.images.english.small_gray")
|
|
get_image_list_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.utils.get_image_list",
|
|
'input': input_,
|
|
'output': easyocr.utils.get_image_list(*input0),
|
|
'severity': "Error"
|
|
},
|
|
})
|
|
|
|
output0 = ocr.detect(english_mini_bgr)
|
|
input0 = [output0[0][0],
|
|
output0[1][0],
|
|
english_mini_gray,
|
|
64, #model_height
|
|
True# sort_output
|
|
]
|
|
input_ = replace(input0, 2, "unit_test.inputs.images.english.mini_gray")
|
|
get_image_list_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.utils.get_image_list",
|
|
'input': input_,
|
|
'output': easyocr.utils.get_image_list(*input0),
|
|
'severity': "Error"
|
|
},
|
|
})
|
|
|
|
output0 = ocr.detect(english_mini_bgr)
|
|
input0 = [output0[0][0],
|
|
output0[1][0],
|
|
english_mini_gray,
|
|
64, #model_height
|
|
True# sort_output
|
|
]
|
|
image_list, max_width = easyocr.utils.get_image_list(*input0)
|
|
|
|
input0 = [ocr.character,
|
|
64, #imgH - fixed
|
|
int(max_width),
|
|
ocr.recognizer,
|
|
ocr.converter,
|
|
image_list[:2],
|
|
'', #ignore_char,
|
|
'greedy', #decoder,
|
|
5, #beamWidth,
|
|
1, #batch_size,
|
|
0.1, #contrast_ths,
|
|
0.5, #adjust_contrast,
|
|
0.003, #filter_ths,
|
|
1, #workers,
|
|
"cuda" #device
|
|
]
|
|
|
|
get_text_test = {}
|
|
|
|
output_ = easyocr.recognition.get_text(*input0)
|
|
input_ = replace(input0,
|
|
[0, 3, 4],
|
|
["unit_test.easyocr.ocr.character",
|
|
"unit_test.easyocr.ocr.recognizer",
|
|
"unit_test.easyocr.ocr.converter"]
|
|
)
|
|
|
|
get_text_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.recognition.get_text",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
},
|
|
})
|
|
|
|
input0 = [ocr.character,
|
|
64, #imgH - fixed
|
|
int(max_width),
|
|
ocr.recognizer,
|
|
ocr.converter,
|
|
image_list[:2],
|
|
'', #ignore_char,
|
|
'greedy', #decoder,
|
|
4, #beamWidth,
|
|
1, #batch_size,
|
|
0.05, #contrast_ths,
|
|
0.3, #adjust_contrast,
|
|
0.001, #filter_ths,
|
|
1, #workers,
|
|
"cuda" #device
|
|
]
|
|
|
|
output_ = easyocr.recognition.get_text(*input0)
|
|
input_ = replace(input0,
|
|
[0, 3, 4],
|
|
["unit_test.easyocr.ocr.character",
|
|
"unit_test.easyocr.ocr.recognizer",
|
|
"unit_test.easyocr.ocr.converter"]
|
|
)
|
|
get_text_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.recognition.get_text",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
}})
|
|
|
|
input0 = [ocr.character,
|
|
64, #imgH - fixed
|
|
int(max_width),
|
|
ocr.recognizer,
|
|
ocr.converter,
|
|
image_list[:2],\
|
|
'', #ignore_char,
|
|
'greedy', #decoder,
|
|
6, #beamWidth,
|
|
4, #batch_size,
|
|
0.2, #contrast_ths,
|
|
0.6, #adjust_contrast,
|
|
0.005, #filter_ths,
|
|
1, #workers,
|
|
"cuda" #device
|
|
]
|
|
|
|
output_ = easyocr.recognition.get_text(*input0)
|
|
input_ = replace(input0,
|
|
[0, 3, 4],
|
|
["unit_test.easyocr.ocr.character",
|
|
"unit_test.easyocr.ocr.recognizer",
|
|
"unit_test.easyocr.ocr.converter"]
|
|
)
|
|
get_text_test.update({'test03': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.recognition.get_text",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
}})
|
|
|
|
|
|
get_paragraph_test = {}
|
|
output0 = ocr.detect(english_mini_bgr)
|
|
input0 = [output0[0][0],
|
|
output0[1][0],
|
|
english_mini_gray,
|
|
64, #model_height
|
|
True# sort_output
|
|
]
|
|
image_list, max_width = easyocr.utils.get_image_list(*input0)
|
|
|
|
input0 = [ocr.character,
|
|
64, #imgH - fixed
|
|
int(max_width),
|
|
ocr.recognizer,
|
|
ocr.converter,
|
|
image_list[:2],
|
|
'', #ignore_char,
|
|
'greedy', #decoder,
|
|
5, #beamWidth,
|
|
1, #batch_size,
|
|
0.1, #contrast_ths,
|
|
0.5, #adjust_contrast,
|
|
0.003, #filter_ths,
|
|
1, #workers,
|
|
"cuda" #device
|
|
]
|
|
|
|
output0 = easyocr.recognition.get_text(*input0)
|
|
input_ = [output0,
|
|
1, #x_ths
|
|
0.5, #y_ths
|
|
'ltr' #mode
|
|
]
|
|
get_paragraph_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.utils.get_paragraph",
|
|
'input': input_,
|
|
'output': easyocr.utils.get_paragraph(*input_),
|
|
'severity': "Error"
|
|
}})
|
|
input_ = [output0,
|
|
0.5, #x_ths
|
|
0.3, #y_ths
|
|
'ltr' #mode
|
|
]
|
|
get_paragraph_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.utils.get_paragraph",
|
|
'input': input_,
|
|
'output': easyocr.utils.get_paragraph(*input_),
|
|
'severity': "Error"
|
|
}})
|
|
input_ = [output0,
|
|
1.5, #x_ths
|
|
1, #y_ths
|
|
'ltr' #mode
|
|
]
|
|
get_paragraph_test.update({'test03': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.utils.get_paragraph",
|
|
'input': input_,
|
|
'output': easyocr.utils.get_paragraph(*input_),
|
|
'severity': "Error"
|
|
}})
|
|
|
|
|
|
input_recog = [None,
|
|
None, #horizontal_list
|
|
None, #free_list
|
|
'greedy', #decoder
|
|
5, #beamWidth
|
|
1,#batch_size
|
|
0, #workers
|
|
None, #allowlist
|
|
None, #blocklist
|
|
1, #detail
|
|
None, #rotation_info
|
|
False,#paragraph
|
|
0.1,#contrast_ths
|
|
0.5, #adjust_contrast
|
|
0.003, #filter_ths
|
|
0.5, #y_ths
|
|
1.0, #x_ths
|
|
True, #reformat
|
|
'standard'#output_format
|
|
]
|
|
|
|
recognize_test = {}
|
|
|
|
h_list, f_list = ocr.detect(english_mini_bgr)
|
|
input_ = replace(input_recog,
|
|
[0, 1, 2],
|
|
[english_mini_gray, h_list[0], f_list[0]])
|
|
recognize_test.update({'test01': {
|
|
'description': "Testing with default input.",
|
|
"method": "unit_test.easyocr.ocr.recognize",
|
|
'input': replace(input_, 0, "unit_test.inputs.images.english.mini_gray"),
|
|
'output': ocr.recognize(*input_),
|
|
'severity': "Error"
|
|
}})
|
|
|
|
h_list, f_list = ocr.detect(english_small_bgr)
|
|
input_ = replace(input_recog,
|
|
[0, 1, 2],
|
|
[english_small_gray, h_list[0], f_list[0]])
|
|
recognize_test.update({'test02': {
|
|
'description': "Testing with custom input.",
|
|
"method": "unit_test.easyocr.ocr.recognize",
|
|
'input': replace(input_, 0, "unit_test.inputs.images.english.small_gray"),
|
|
'output': ocr.recognize(*input_),
|
|
'severity': "Error"
|
|
}})
|
|
|
|
readtext_test = {}
|
|
#english_tiny_bgr, _ = easyocr.utils.reformat_input(images['english']['tiny'])
|
|
input_ = ["unit_test.inputs.images.english.tiny", 'en']
|
|
ocr = get_easyocr('en')
|
|
_, pred, confidence = ocr.readtext(images['english']['tiny'])[0]
|
|
output_ = [pred, confidence]
|
|
readtext_test.update({'test01': {
|
|
'description': "Reading English text.",
|
|
"method": "unit_test.easyocr_read_as",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
}})
|
|
#french_tiny_bgr, _ = easyocr.utils.reformat_input(images['french']['tiny'])
|
|
input_ = ["unit_test.inputs.images.french.tiny", 'fr']
|
|
ocr = get_easyocr('fr')
|
|
_, pred, confidence = ocr.readtext(images['french']['tiny'])[0]
|
|
output_ = [pred, confidence]
|
|
readtext_test.update({'test02': {
|
|
'description': "Reading French text.",
|
|
"method": "unit_test.easyocr_read_as",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
}})
|
|
#chinese_tiny_bgr, _ = easyocr.utils.reformat_input(images['chinese']['tiny'])
|
|
input_ = ["unit_test.inputs.images.chinese.tiny", 'ch_sim']
|
|
ocr = get_easyocr('ch_sim')
|
|
_, pred, confidence = ocr.readtext(images['chinese']['tiny'])[0]
|
|
output_ = [pred, confidence]
|
|
readtext_test.update({'test03': {
|
|
'description': "Reading Chinese (simplified) text.",
|
|
"method": "unit_test.easyocr_read_as",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
}})
|
|
#korean_tiny_bgr, _ = easyocr.utils.reformat_input(images['korean']['tiny'])
|
|
input_ = ["unit_test.inputs.images.korean.tiny", 'ko']
|
|
ocr = get_easyocr('ko')
|
|
_, pred, confidence = ocr.readtext(images['korean']['tiny'])[0]
|
|
output_ = [pred, confidence]
|
|
readtext_test.update({'test04': {
|
|
'description': "Reading Korean text.",
|
|
"method": "unit_test.easyocr_read_as",
|
|
'input': input_,
|
|
'output': output_,
|
|
'severity': "Error"
|
|
}})
|
|
|
|
|
|
|
|
solution_book = {
|
|
'inputs':{'images': image_preprocess,
|
|
'easyocr_config': easyocr_config
|
|
},
|
|
'tests':{
|
|
"model initialization": model_init_test,
|
|
"get_textbox function": get_textbox_test,
|
|
"group_text_box function": group_text_box_test,
|
|
"detect method": detect_test,
|
|
"get_image_list function": get_image_list_test,
|
|
"get_text_test function": get_text_test,
|
|
"get_paragraph_test function": get_paragraph_test,
|
|
"recognize method": recognize_test,
|
|
"readtext method": readtext_test,
|
|
}
|
|
}
|
|
|
|
|
|
|
|
with lzma.open(args.output, 'wb') as fid:
|
|
pickle.dump(solution_book, fid)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Script to pack EasyOCR weight.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument("-o", "--output", default=None, help="output path.")
|
|
parser.add_argument("-d", "--data_dir", default=None, help="data directory")
|
|
args = parser.parse_args()
|
|
main(args)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|