Model Demo: Optical Character Recognition (OCR) Zoo - MMOCR

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/zoo/ocr_zoo/mmocr.ipynb.

MMOCR

MMOCR is an open-source toolbox based on PyTorch and mmdetection for text detection, text recognition, and the corresponding downstream tasks including key information extraction. It is part of the OpenMMLab project.

MMOCR hosts text detection models and text recognition ones. They work as a pipeline to recognize characters in images.

The text detection models detect polygons containing characters in images. The following is an example of model inference results by converting polygons to bounding boxes.

The text recognition models recognize characters from each detected image portion.

In this notebook, we will explore the text detection models.

Environment Setup

!pip3 install -r ../../../requirements.txt -q
!mim install -r ../../../requirements_mim.txt -q

import gc
import os
import pickle
import sys
from pathlib import Path
import shutil
import zipfile

import numpy as np
import onnx
from onnxruntime import InferenceSession
from PIL import Image

from examples.models.zoo.zoo_utils import (
    download_file,
    get_mmdeploy_parameters_and_model_info,
    get_transforms_and_dataset,
    install_openmmlab_github,
    onnx_check_and_simplify,
    split_onnx,
)
from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from sdk_cli.lib.inference import InferenceEngine, batch_inference
from sdk_cli.lib.quantize import QuantizedONNXModel, quantize_onnx_model
from sdk_cli.utils.models.mmdeploy import MMDeployVariant
from sdk_cli.utils.models.mmocr import MMOCRModelVariant
from sdk_cli.utils.networks import ChimeraNetwork
from sdk_cli.visualizers.layouter import (
    DetectorsLayouter,
    VisualizeHandle,
    VisualizeObject,
    get_visualize_handle,
)
from ocr_utils import (
    get_mmocr_parameters,
    mmocr_textdet_postprocessing,
    MMOCRInfo,
    setup_mmocr_code,
)

CUR_DIR = os.path.abspath(os.getcwd())

sys.path.insert(0, ".")

Install MMOCR and MMDeploy

github_names = [
    ("mmocr", "1.0.1", "mmocr-main"),
    ("mmdeploy", "1.3.1", None),
    ("mmengine", "0.10.7", "mmengine-main"),
]
symlink_list = [
    ("configs", "mmocr-main"),
    ("demo", "mmocr-main"),
    ("mmengine", "mmengine-main"),
    ("mmocr", "mmocr-main"),
    ("model-index.yml", "mmocr-main"),
]
install_openmmlab_github(github_names, symlink_list)
setup_mmocr_code()

Model Lowering

Select Model

This notebook can experiment the following models.

• Text Detector
  • dbnet_resnet18_fpnc_icdar2015
  • dbnet_resnet18_fpnc_totaltext
  • dbnet_resnet50_fpnc_icdar2015
  • dbnetpp_resnet50_fpnc_icdar2015
  • fcenet_resnet50_fpn_icdar2015
  • fcenet_resnet50_fpn_totaltext
  • panet_resnet18_fpem_ffm_ctw1500
  • panet_resnet18_fpem_ffm_icdar2015
  • psenet_resnet50_fpnf_ctw1500
  • psenet_resnet50_fpnf_icdar2015
  • textsnake_resnet50_fpn_unet_ctw1500

Now we are going to use the following model as an example.

• fcenet_resnet50_fpn_totaltext

## Text Detector
text_det = MMOCRInfo(MMOCRModelVariant.fcenet_resnet50_fpn_totaltext)

text_det.mmocr_params = get_mmocr_parameters(text_det.name)
text_det.mmdeploy_params, text_det.model_info = get_mmdeploy_parameters_and_model_info(
    MMDeployVariant.mmocr, text_det.mmocr_params
)
text_det.onnx_file = f"{text_det.name}.onnx"

Load Pretrained Model

weights = Path(text_det.mmdeploy_params.weights_url).name

if not Path(weights).exists():
    download_file(text_det.mmdeploy_params.weights_url, weights)
    print(f"{weights} is downloaded.")

Export ONNX

import examples.models.zoo.fix_registry

os.chdir(Path(CUR_DIR) / "mmdeploy")

from mmdeploy.apis import torch2onnx
import torch

os.chdir(CUR_DIR)

if Path(text_det.onnx_file).exists():
    print(f"{text_det.onnx_file} is already exported. Skip exporting.")
else:
    Image.open("demo/demo_text_ocr.jpg").resize(text_det.mmocr_params.model_input_size).save(
        "image.jpg"
    )

    # Convert the model to ONNX
    torch2onnx(
        "image.jpg",
        work_dir=".",
        save_file=text_det.onnx_file,
        deploy_cfg=text_det.mmdeploy_params.deploy_config,
        model_cfg=text_det.mmdeploy_params.config,
        model_checkpoint=Path(text_det.mmdeploy_params.weights_url).name,
        device="cpu",
    )

    onnx_model = onnx_check_and_simplify(onnx.load(text_det.onnx_file))
    onnx.save(onnx_model, text_det.onnx_file)

print(f"ONNX is exported and simplified as {text_det.onnx_file}")

Split ONNX

## Get input edges and output edges
session = InferenceSession(onnx.load(text_det.onnx_file).SerializeToString())
input_edges = [inp.name for inp in session.get_inputs()]
output_edges = [oup.name for oup in session.get_outputs()]

## Split onnx
text_det.target_onnx_file, text_det.second_onnx_file = split_onnx(
    text_det.onnx_file,
    text_det.mmocr_params,
    input_edges=input_edges,
    output_edges=output_edges,
)

Quantize and Compile

## Get transforms and dataset for the target model
text_det.transforms, dataset = get_transforms_and_dataset(text_det.mmdeploy_params)

## Quantization
print(f"\nStart quantization for {text_det.name}")
quantized_onnx_model = quantize_onnx_model(
    text_det.target_onnx_file,
    dataset,
    asymmetric_activation=True,
)
print(
    f"quantized onnx: {quantized_onnx_model.model_path}, tranges file: {quantized_onnx_model.tensor_ranges_path}"
)

## Compilation
print(f"\nStart compilation for {text_det.name}")
text_det.cgc_job = ChimeraJob(
    model_p=str(quantized_onnx_model.model_path),
    trange_file=str(quantized_onnx_model.tensor_ranges_path),
)
text_det.cgc_job.compile(quiet=True)
print(f"Compilation completed for {text_det.name}\n")

Demo

Text Detection

Inference

all_image_paths = [
    "demo/demo_text_det.jpg",
]

all_images = []
for image_path in all_image_paths:
    image = np.array(Image.open(image_path))
    transformed_image = text_det.transforms(image)
    all_images.append(np.expand_dims(transformed_image.numpy(), axis=0).astype(np.float32))

engines = {
    InferenceEngine.CHIMERA_ORT_INT8: text_det.cgc_job,
    InferenceEngine.CHIMERA_ISS_INT8: text_det.cgc_job,
}

head_session = (
    None
    if text_det.second_onnx_file is None
    else InferenceSession(onnx.load(text_det.second_onnx_file).SerializeToString())
)

outputs_per_inference_engine = {}
THREADS = min(len(all_images), 6)
for inference_engine, engine in engines.items():
    all_outputs = batch_inference(
        inference_engine,
        engine,
        all_images,
        head_session=head_session,
        threads=THREADS,
    )
    outputs_per_inference_engine[inference_engine] = all_outputs

Run Statistics

print(text_det.cgc_job)
text_det.cgc_job.plot_run_statistics()

Visualized Results

from mmocr.apis.inferencers import MMOCRInferencer

%matplotlib inline

## Get visualize handle for visualization.
visualize_handle = get_visualize_handle(
    text_det.mmocr_params.visualize_object, text_det.mmocr_params.visual_threshold
)

inferencer = MMOCRInferencer(text_det.model_info["Name"])

for image_idx, image_file in enumerate(all_image_paths):
    layouter = DetectorsLayouter(
        visualize_handle,
        image=image_file,
        show_class=False,
    )
    for inference_engine, all_outputs in outputs_per_inference_engine.items():
        detections, _ = mmocr_textdet_postprocessing(
            outputs=all_outputs[image_idx],
            image_file=image_file,
            model_input_size=text_det.mmocr_params.model_input_size,
            model_postprocessor=inferencer.textdet_inferencer,
            pre_postprocess=text_det.mmocr_params.pre_postprocess,
        )
        layouter.add_data(
            bboxes=detections,
            title=f"{str(inference_engine).upper()}\nModel: {text_det.name}",
        )

    layouter.display()

Citation

@article{mmocr2022,
    title={MMOCR:  A Comprehensive Toolbox for Text Detection, Recognition and Understanding},
    author={MMOCR Developer Team},
    howpublished = {\url{https://github.com/open-mmlab/mmocr}},
    year={2022}
}

@misc{=mmdeploy,
    title={OpenMMLab's Model Deployment Toolbox.},
    author={MMDeploy Contributors},
    howpublished = {\url{https://github.com/open-mmlab/mmdeploy}},
    year={2021}
}