Package: efficient-ocr
EffOCRThis repository contains the official implementation ofEffOCR. EffOCR is an OCR architecture that emphasizes sample (and computational) efficiency without compromising on accuracy! EffOCR accomplishes this by treating text recognition not as a sequence/language modeling task, but as a metric learning/image retrieval task.EffOCR is comprised of two modules, trained separately:alocalizer, an object detector that is fine-tuned for the purpose of detecting characters/letters/glyphs/words in an image of text; andarecognizer, which classifies/assigns identities to localized glyphs.In this repo, we provide tools for training your own EffOCR models for your particular documents of interest, as well as for running fast and accurate inference on new samples.InstallationClone this repo, then installrequirements.txtin an environment withPython >= 3.7. (To accommodate downloading GPU-compatible FAISS, as referenced below, it can be useful to first set up and activate acondaenvironment before installing these required dependencies.)gitclonegit@github.com:dell-research-harvard/effocr.gitcdeffocr
pipinstall-rrequirements.txtSeparately, you'll also need to installbothFAISSandPyTorchas well as eitherDetectron2orMMDetectionorYOLOv5as is appropriate for your use case.Dataset FormattingFor the purposes of training and evaluating EffOCR localizers and recognizers, we require that your dataset of interest is transcribed into theCOCO formatwith annotations at the character level, and with image level text fields that denote the transcription of the text line image, e.g.,...
  "images": [
    {
      "id": 13,
      "width": 958,
      "height": 46,
      "file_name": "194906.png",
      "text": "they arrived overland from San Fran-"
    },
...Thescriptsdirectory contains a few useful utilities for combining, splitting, and subsetting COCO-formatted JSON files.This COCO-formatted dataset can be used directly for localizer training, and can be converted into another format suitable for recognizer training with a single script, described below.By default, English and Japanese language datasets are supported by EffOCR. Other languages are certainly compatible with the EffOCR framework, but it might require that you fork and modify this codebase as appropriate for your application of interest.LoggingBy default, EffOCR supports logging withWeights & Biases. We encourage you therefore to set up a W&B account and initiate awandb logincommand before starting out on the training process.TrainingLocalizerPresently, EffOCR localizer models can be trained either via the providedtrain_effocr_localizer_d2.pyscriptorvia theMMDetectionorYOLOv5object detection libraries. We aim to provide built-in support for MMDetection and YOLOv5 localizer training in the near future.To use the Detectron2-backend training localizer script, a character localization dataset must first belocally registered. As one example, calling thetrain_effocr_localizer_d2.pyscript could look like:python./train_effocr_localizer_d2.py\--name<W&Brunname>\--config<absolutepathtoD2config>\--dataset_name<nameofDetectron2datasettoberegistered>\--dataset_root<absolutepathtorootdirforCOCOformatteddata>\--train_file<nameofCOCOJSONtrainingfile>\--val_file<nameofCOCOJSONvalfile>\--test_file<nameofCOCOJSONtestfile>\--images_dir<namedirinrootwithimagesfortraining,testing,validation>\--init_ckpt<absolutepathtosavedmodelcheckpoint,ifapplicable>\--num_epochs<e.g.100>\--lr<e.g.1e-5>\--batch_size<e.g.2>\--num-gpus<N>RecognizerThe scriptformat_effocr_recognizer_dataset.pytakes in a COCO-style dataset that would be used for e.g. localizer training and outputs a dataset amenable to recognizer training, which is essentially a dataset with a file structure compatible with a slightly modified version of a PyTorchImageFolderdataset class. Specifically, the script creates a training dataset by combining synthetic charaters generated from a wide collection of fonts with user-provided character crops from source documents.pythonformat_effocr_recognizer_dataset.py\--image_dir<absolutepathtoCOCOimagedir>\--coco_jsons<comma-separatedabsolutepathstoCOCOJSONs,e.g.,/path/to/train.json,/path/to/test.json,/path/to/val.json>\--crops_save_dir<absolutepathforsavingglyphcrops,anintermediateoutput>\--cat_id<COCOcategoryIDforcropping,e.g.,0>\--clip_to_top_and_bottom--font_dir<absolutepathtofontsforsyntheticdatamixin,e.g.,./english_font_files>\--charset_dir<pathtocharactersetsforrenderingsyntheticdata,e.g.,./english_charsets>\--dataset_save_dir<absolutepathforsavingrecognizerdataset>\--padding<e.g.0.05>\--spacesRecognizer models are trained via thetrain_effocr_recognizer.pyscript. Example usage is as follows:python./train_effocr_recognizer.py\--root_dir_path<abspathtorecognizerdataset,specifiedindataset_save_dirabove>\--train_ann_path<abspathtotrainannotationsfromCOCOdataset>\--val_ann_path<abspathtovalannotationsfromCOCOdataset>\--test_ann_path<abspathtotestannotationsfromCOCOdataset>\--run_name<W&Brunnameforlogging>\--auto_model_timm<timmmodelname(seebelow)>\--batch_size<e.g.128>\--num_epochs<e.g.60>\--num_passes<e.g.1forJapanese,10forEnglish>\--lang<e.g.,"en"forEnglishcharacteror"jp"forJapanese>\--lr<e.g.2e-5>\--imsize<mustbecompatiblewithtimmmodelinput,e.g.,224>\--infer_hardneg_k<e.g.8>\--hns_txt_path<usedinasecondrun,orifyouhavealreadycreatedahns.txtfile,seebelow>To view a list of available timm models that one can use as encoders for the EffOCR recognizer, one can run:importtimmprint(timm.list_models(pretrained=True))We recommend many of the models used in theEffOCR paper, including:mobilenetv3_small_050convnext_tinyxcit_small_12_p8_224Hard Negative Mining and TrainingTypically, recognizer models are trained in two steps: a first training run is essentially just used to mine hard negatives and a second training run trains a model incorporating those hard negatives (N.B. this second model does not need to be fine-tuned from the model output in the first run, such that if you have already gathered info on hard negatives, you can jump right to hard negative training). Adding the--infer_hardneg_kargument generates a list of $k$ hard negatives (most visually similar charaters) for each character in the dataset and saves this list to a text file calledhns.txt. Passing in the absolute path of this text file of generated hard negatives via the--hns_txt_pathargument will utilize the enclosed hard negative information in a training run, typically increasing recognizer performance substantially.InferenceText line inference is done via either theinfer_effocr.pyorinfer_effocr_onnx_multi.pyscripts.infer_effocr.pysupports Detectron2- or MMDetection-backend localizers. It doesnotcurrently support YOLOv5-backend localizers.An example call toinfer_effocr.pywould look like:python./infer_effocr.py\--image_dir<pathtotextlineimagestobeinferenced>\--infer_over_img_dir\--auto_model_timm<timmmodelname(sameasusedinrecognizertraining)\--recognizer_dir<directorywithsavedrecognizermodelas"enc_best.pth">\--localizer_dir<directorywithsavedlocalizermodelas"best_bbox_mAP.pth"andconfigas"*.yaml"or"*.py">\--lang<e.g.,"en"forEnglishor"jp"forJapanese>--rcnn_score_thr<e.g.0.3>--save_output<optional,directorytosavepredictionsto>Theinfer_effocr_onnx_multi.pyscript accepts localizer models in ONNX format from any of the backends mentioned above. It also expects a recognizer in ONNX format. Aconversion scriptis provided for convenience.infer_effocr_onnx_multi.pyrunsonlyon the CPU and implements multithreading in an attempt to minimize inference time. An example call to `infer_effocr_onnx_multi.pywould look like:python./infer_effocr_onnx_multi.py\--image_dir<pathtotextlineimagestobeinferenced>\--recognizer_dir<directorywithsavedrecognizermodelas"enc_best.onnx">\--localizer_dir<directorywithsavedlocalizermodelas"best_bbox_mAP.onnx">\--localier_backend<wherethelocalizermodelwastrained,oneof"mmdetection","yolo",or"detectron2"\--lang<e.g.,"en"forEnglishor"jp"forJapanese>\--localizer_iou_thresh<e.g.0.01>\--localizer_conf_thresh<e.g.0.35>\--num_threads<N>\--save_output<optional,directorytosavepredictionsto>Synthetic Data GenerationIf you're interested in synthetic data generation for, e.g., localizer pre-training, then you may also be interested in the companion repo to EffOCR,EffSynth.