Please read till the END of this file BEFORE executing any code. Go through this presentation for a deep dive into the problem statement, datasets, methods etc. Thanks.
For running the code in this repository, using a virtual environment is strongly recommended. The codebase was developed on Python 3.8.5. Install the following packages in the virtual environment-
- pandas
- numpy
- matplotlib
- torch
- transformers
- datasets
- rasterio
- jiwer
- spacy
- tqdm
For NER to work correctly, the spaCy module en_core_web_md must also be installed in the same environment. To install it, run the following command after activating the environment:
python -m spacy download en_core_web_mdDownload the dataset from here. Unzip the zipped file inside the directory which the repository was cloned into. The data to run the code will now be available under the /cvl-database-1-1 directory. To run the code on the CVL dataset, activate the environment, and run the different code files in the following order:
python make_splits.pyThe make_splits.py script parses the XML files in the CVL dataset to generate line-level and page-level string labels for the samples in the CVL datasets. It creates four .csv files-
line_trainset.csvline_testset.csvpage_trainset.csvpage_testset.csvThese files are used by the training script. If they are not produced correctly, the training script might throw errors.
python training.pyThis file contains the main training and evaluation code. It defines the dataset class for the CVL dataset and allows fine-tuning for the pre-trained TrOCR models from HuggingFace. The PAGE parameter in the script controls whether the model is being fine-tuned on page-level or line-level data. This has to be changed accordingly. The pre-trained model that is being loaded can also be changed. By default it is set to microsoft/trocr-base-handwritten. Other parameters such as batch_size, learning rate can be changed as per requirement. This code was developed on an NVIDIA A5000 GPU with 24GB VRAM. A similar or better GPU is strongly recommended.
The script creates a pickle file recognized_text.pickle. On unpickling this, it yields a list of strings which is the text recognized by the model of the input samples in the test set. The following scripts are dependent on this file.
python spellcheck_ner.pyThis script loads the pickle file created by training.py and performs spelling correction and NER on the outputs of the OCR model. The spelling correction pipeline is based on a model from HF and is quite simple. The NER pipeline has two variants-- (1) spaCy, and (2) Transformer-based. In general, the spaCy model is preferred because it recognizes a much larger variety of entities, is lightweight, and much faster. The extract_entities function in the script has a parameter reduced. If set to True, the entities extracted by the spaCy model will be mapped to a smaller set of broader entities defined as a dictionary within the script itself. This mapping can be easily changed by directly modifying the dictionary. This behaviour is False by default.
This script creates two pickle files:
corrected_text.picklewhich contains the outputs of the spelling correction model.entities_list.picklewhich contains the set of extracted entities in a specific format (described in the presentation).
The Transformer-based pipeline is right after spaCy's. It is commented out to save execution time since the results are quite poor. However, simply uncommenting that part of the code will run the pipeline and print the results. However, the pickle file will still contain the results returned by the spaCy pipeline.
python build_schema.pyThis script loads the pickle files generated by spellcheck_ner.py. Using those files, it creates a data schema (described in the presentation) to effectively organize and store the recognized text, extracted entities, etc. for each sample. It outputs a json file containing the data organized that particular format for each sample.
The checkpoints directory contains the file-tuned model weights. Two files are available, each using a different set of pre-trained weights from HuggingFace but both fine-tuned on theh CVL dataset using the procedure described above.
The plots directory contains some useful plots that give some insight about the characteristics of the dataset and fine-tuning performance of the pre-trained models. The code to generate these plots is not provided in a separate script, but they are available in the CVL Handwriting Recognition.ipynb notebook. This notebook also contains the whole, unified codebase for ease of use. Execution outputs are also provided for ease of understanding. Please go through the notebook BEFORE executing the scripts.
Finally, the test_images directory contains a few self-written sample input texts to test the validity of the model on real-world images (poor lighting, image resolution very different from the training data, tilted/improperly cropped text). The result/output of this can also be viewed in the 'Sample Testing' section of the jupyter notebook. To run the model on custom test images, please create a new directory containing line-by-line crops (necessary) of the input image, and replace the image_folder variable in the script with path/to/new/dir. Then run the inference_example.py script. This will internally call all the other scripts and produce the json file.
However, please note that running inference_example.py will OVERRIDE the pickle files previously created in other runs of the code. This behaviour can be modified by simply adding arguments for file names and (hyper)parameters instead of hard-coding them in each of these scripts. It will be updated at a later stage.