DET is an end-to-end tool for extracting Key-Value pairs from a variety of documents, built entirely on PyTorch and served using TorchServe.
Try it live on the web-app! 👋
The Document Extraction tool is composed of two main components, namely,
1. OCR (Optical Character Recognition) and
2. Document NER (Named Entity Recognition)
The OCR component comprises of a Detection and Recognition module which work sequentially to produce results which is then consumed by the NER part for training/prediction.
For the OCR part, we have deployed a torchserve model server on GCP using Vertex AI service. Using torchserve, we deploy a complex workflow in the form of a DAG, comprising of pre_processing, detection and recognition models.
For the NER part, we created a training module using transformers library which requires the text and bounding-box results from OCR output to train/predict documents.
The architectural flow of the two modules is shown 
- Optical Character Recognition (OCR):
- Detection:
- Recognition:
- Named Entity Recognition (NER):
- Receipt Dataset (SROIE):
- Forms Dataset (FUNSD):
deploy # INSTRUCTIONS/SCRIPTS for deploying model(s) on GPU/CPU |---GPU # For Serving models on GPU |------jit-models # contains jitted models created after tracing the model using torch.jit.trace |------------craft_ts.pt # DOWNLOAD the torchscript files from here: https://drive.google.com/drive/folders/1NBSZIZzSzIVOUqnxu0PHgmy-_Tvvp2hY?usp=sharing |------------crnn_ts.pt |------------sroie_ts.pt |------------funsd_ts.pt |------model-archive # stores models along with configuration in .mar format as required by "torchserve" |------------model_store # GENERATE standalone .mar files from torch-model-archiver command given below |------------------craft.mar |------------------crnn.mar |------------------sroie.mar |------------------funsd.mar |------------wf_store # GENERATE workflow .mar files from torch-workflow-archiver command given below |------------------ocr.mar |------------------ner_sroie.mar |------------------ner_funsd.mar |------config.properties # config.properties for storing serving related configurations. | For details, refer this: https://github.com/pytorch/serve/blob/master/docs/configuration.md |------detection_handler.py # handler for text detection pipeline |------rec_handler.py # handler for text recognition pipeline |------workfow_handler.py # handler for configuring end2end ocr serving pipeline |------workflow_ocr.yaml # define the pipeline here, also specify batch-size, num workers etc. |------workflow_ner_sroie.yaml |------workflow_ner_funsd.yaml |------Dockerfile # template for creating GPU docker image for UCR |------.dockerignore # files to ignore in docker image |------cloud_deploy.ipynb # scripts to deploy the file on GCP Vertex ai and run predictions. |------sample_b64.json # sample file to send request on inference api |------index_sroie.json # contains mapping of indexes to human-readable labels for SROIE dataset |------index_funsd.json # contains mapping of indexes to human-readable labels for FUNSD dataset |---CPU # Same as above except it's CPU (currently INCOMPLETE) |------jit-models |------------craft_ts.pt |------------crnn_ts.pt |------model-archive |------------model_store |------------------craft.mar |------------------crnn.mar |------------wf_store |------------------ocr.mar |------config.properties | |------detection_handler.py |------rec_handler.py |------workfow_handler.py |------workflow.yaml |------Dockerfile train # NOTEBOOKS for training models on GPU/CPU. More training scripts COMING SOON! |---NER.ipynb # Jupyter Notebook to train,test Document NER models and convert it to torchscript format. Try it live here! 👋
Install Docker and NVIDIA Container Toolkit: See this link for help!
Download and start docker:
docker run -d --gpus all -p 7080:7080 -p 7081:7081 -p 7082:7082 abhigyanr/det-gpu:latest
For testing API, follow these steps! Note: This method requires NVIDIA GPU and driver to be present!
Install docker and nvidia container toolkit: See this link for help!
Clone this repository and change directory:
git clone https://github.com/RamanHacks/pytorch-hackathon.git cd pytorch-hackathon && cd deploy cd GPU # OR (cd CPU) Build Docker Image
docker build -f Dockerfile -t det . Run Docker container
docker run -d --rm --name det-cpu -p 7080:7080 -p 7081:7081 -p 7082:7082 det # (For CPU) docker run -d --rm --name det-gpu --gpus all -p 7080:7080 -p 7081:7081 -p 7082:7082 det # (For GPU, use --gpus '"device=0,1"' to specify device) Optional: Check Status
docker logs $(docker ps -l -q) # to check if the docker container is running fine curl localhost:7080/ping # to check if the network is accessible from localhost, should return Healthy Register Models
curl -X POST "localhost:7081/workflows?url=ocr.war" curl -X POST "localhost:7081/workflows?url=ner_sroie.war" curl -X POST "localhost:7081/workflows?url=ner_funsd.war" Optional: Stop and Remove Container
docker stop $(docker ps -l -q) docker rm $(docker ps -l -q) For testing API, follow these steps!
Install torch from official link: PyTorch Official
Install torchserve from official repo: TorchServe Official
Clone this repository and install dependencies:
git clone https://github.com/RamanHacks/pytorch-hackathon.git cd pytorch-hackathon && cd deploy pip install -r requirements.txt cd GPU # OR (cd CPU) Download pretrained torchscript models from Google Drive and move it inside jit-models folder.
Generate .mar files:
# create model archives torch-model-archiver -f --model-name craft --version 1.0 --serialized-file jit-models/craft_ts.pt --handler det_handler.py --export-path model-archive/model_store/ torch-model-archiver -f --model-name crnn --version 1.0 --serialized-file jit-models/crnn_ts.pt --handler rec_handler.py --export-path model-archive/model_store/ cp index_sroie.json index.json torch-model-archiver -f --model-name sroie --version 1.0 --serialized-file jit-models/sroie_ts.pt --handler ext_handler.py --export-path model-archive/model_store/ --extra-files index.json cp index_funsd.json index.json torch-model-archiver -f --model-name funsd --version 1.0 --serialized-file jit-models/funsd_ts.pt --handler ext_handler.py --export-path model-archive/model_store/ --extra-files index.json rm index.json # create workflow archives torch-workflow-archiver -f --workflow-name ocr --spec-file workflow_ocr.yaml --handler workflow_handler.py --export-path model-archive/wf_store/ torch-workflow-archiver -f --workflow-name ner_sroie --spec-file workflow_ner_sroie.yaml --handler workflow_handler.py --export-path model-archive/wf_store/ torch-workflow-archiver -f --workflow-name ner_funsd --spec-file workflow_ner_funsd.yaml --handler workflow_handler.py --export-path model-archive/wf_store/ Start Model Server
torchserve --start --model-store model-archive/model_store/ --workflow-store model-archive/wf_store/ --ncs --ts-config config.properties Register Models
curl -X POST "localhost:7081/workflows?url=ocr.war" curl -X POST "localhost:7081/workflows?url=ner_sroie.war" curl -X POST "localhost:7081/workflows?url=ner_funsd.war" Optional: Stop TorchServe
torchserve --stop For testing API, follow these steps!
Request format: Create a sample
jsonfile containing base64 values of image
{ 'b64': '<base64 value of an image>' }Response format of OCR(only), i.e. when hitting "/wfpredict/ocr":
[ { 'bbox': [[<top-left>],[<top-right>],[<bottom-left>],[<bottom-right>]], 'ocr': [<value>, <confidence>], } ]Response format of OCR+NER, i.e. when hitting "/wfpredict/ner_sroie" or "/wfpredict/ner_funsd":
[ { 'bbox': [<top-left-x>,<top-left-y>,<bottom-right-x>,<bottom-right-y>], 'ocr': <value>, 'key': <value>, } ]Sample CURL Request
curl -X POST -H "Content-Type: application/json; charset=utf-8" -d @sample_b64.json localhost:7080/wfpredict/ner_sroie Python function to convert an image into base64, send request and return predictions
import base64 import requests def sample_request(image_file_path) def convert_b64(image_file): """Open image and convert it to Base64""" with open(image_file, "rb") as input_file: jpeg_bytes = base64.b64encode(input_file.read()).decode("utf-8") return jpeg_bytes req = {"data": convert_b64(image_file_path)} res = requests.post("http://localhost:7080/wfpredict/ner_sroie", json=req) return res.json()Jump to NER.ipynb for details on training and testing Document NER models!
-----Coming-Soon-----
-----Coming Soon-----
-----If you like our work, do not forget to ⭐ this repository and follow us on twitter, linkedin-----
-----If you have got any specific feature request, contact us at admin@docyard.ai-----