unit_test/make_test_solution.py

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)