Model Demo: Segmentors Zoo - MMSegmentation

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/segmentors_zoo/mmsegmentation.ipynb.

MMSegmentation

MMSegmentation is an open source object detection toolbox based on PyTorch. It is a part of the OpenMMLab project.

Please refer the MMSegmentation github for more details.

Environment Setup

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

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

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

from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from tvm.contrib.epu.chimera_job.hw_config import HWConfig
from sdk_cli.lib.inference import InferenceEngine, batch_inference
from sdk_cli.lib.quantize import quantize_onnx_model
from sdk_cli.utils.models.mmdeploy import MMDeployVariant
from sdk_cli.utils.models.mmseg import MMSegModelVariant
from sdk_cli.utils.networks import ChimeraNetwork
from sdk_cli.visualizers.layouter import (
    DetectorsLayouter,
    VisualizeHandle,
    VisualizeObject,
    get_visualize_handle,
)
import examples.models.zoo.fix_registry
from examples.models.zoo.zoo_utils import (
    download_file,
    get_mmdeploy_parameters,
    get_transforms_and_dataset,
    install_openmmlab_github,
    onnx_check_and_simplify,
    split_onnx,
)
from segmentors_utils import get_mmseg_parameters

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

Install MMSegmentation and MMDeploy

github_names = [
    ("mmsegmentation", "1.2.2", None),
    ("mmdeploy", "1.3.1", None),
]
symlink_list = [
    ("mmseg", "mmsegmentation"),
    ("configs", "mmsegmentation"),
    ("model-index.yml", "mmsegmentation"),
]
install_openmmlab_github(github_names, symlink_list)

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

from mmdeploy.apis import torch2onnx
import torch

os.chdir(CUR_DIR)

Create configuration files to export onnx files for input-size variants.

config_text = """_base_ = ['./segmentation_static.py', '../_base_/backends/ascend.py']

onnx_config = dict(input_shape=[%s, %s])
backend_config = dict(
    model_inputs=[dict(input_shapes=dict(input=[1, 3, %s, %s]))])
"""

config_path = "./mmdeploy/configs/mmseg/segmentation_onnxruntime_static-%sx%s.py"

for model_input_size in [
    (size, size) for size in sorted([64, 128, 256, 480, 640, 680, 768, 769, 832, 1024])
] + [(1024, 512)]:
    if not Path(config_path % model_input_size[::-1]).exists():
        with Path(config_path % model_input_size[::-1]).open("w") as fd:
            fd.writelines(config_text % (*model_input_size, *model_input_size[::-1]))
            print(f"{config_path % model_input_size[::-1]} is stored")

Download Datasets

We are going to download HRF dataset. CHASE DB1, DRIVE, STARE can be downloaded from kaggle.

dataset_urls = {
    "hrf": "https://www5.cs.fau.de/fileadmin/research/datasets/fundus-images/healthy.zip",
}

dataset_dir = Path("./datasets")
dataset_dir.mkdir(exist_ok=True)

## Download datasets.
for dataset_name, dataset_url in dataset_urls.items():
    if (dataset_dir / dataset_name).exists():
        print(f"Existing {dataset_dir / dataset_name} is used.")
    else:
        filename = Path(dataset_url).name
        if dataset_url.startswith("https"):
            download_file(dataset_url, filename)
        else:
            shutil.copy(dataset_url, filename)
        with zipfile.ZipFile(filename, "r") as fd:
            fd.extractall(dataset_dir / dataset_name)

        print(f"{dataset_name} is installed.")
        os.remove(filename)

Select Model

This notebook can experiment the following models.

Now we are going to experiment fcn_hr18s as an example.

## Select target model.
MODEL_NAME = MMSegModelVariant.fcn_hr18s

mmseg_parameters = get_mmseg_parameters(MODEL_NAME)
model_input_size = mmseg_parameters.model_input_size
mmdeploy_variant = MMDeployVariant.mmsegmentation

mmdeploy_parameters = get_mmdeploy_parameters(mmdeploy_variant, mmseg_parameters)

Load Pretrained Model

weights = Path(mmdeploy_parameters.weights_url).name

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

Export ONNX

onnx_file = f"{MODEL_NAME}.onnx"

## Convert the model to ONNX
torch2onnx(
    "../../../common/calibration/face/daniel_maverick.png",
    work_dir=".",
    save_file=onnx_file,
    deploy_cfg=mmdeploy_parameters.deploy_config,
    model_cfg=mmdeploy_parameters.config,
    model_checkpoint=weights,
    device="cpu",
)

onnx_model = onnx_check_and_simplify(onnx.load(onnx_file), mmseg_parameters.skip_shape_inference)
onnx.save(onnx_model, onnx_file)

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

gc.collect();

Split ONNX

session = InferenceSession(onnx.load(onnx_file).SerializeToString())
input_edges = [inp.name for inp in session.get_inputs()]
output_edges = [oup.name for oup in session.get_outputs()]

target_onnx_file, second_onnx_file = split_onnx(
    onnx_file,
    mmseg_parameters,
    input_edges=input_edges,
    output_edges=output_edges,
)

del session
gc.collect();

Quantization

transforms, coco_like_subset_of_dataset = get_transforms_and_dataset(mmdeploy_parameters)

quantized_onnx_model = quantize_onnx_model(
    target_onnx_file,
    coco_like_subset_of_dataset,
    asymmetric_activation=True,
)
print(
    f"quantized onnx: {quantized_onnx_model.model_path}, tranges file: {quantized_onnx_model.tensor_ranges_path}"
)

del coco_like_subset_of_dataset
gc.collect();

Compilation

ocm_size = "32MB" if MODEL_NAME in ["fcn_hr48", "mobilenetv2_fcn", "upernet_r50"] else "16MB"

hw_config = HWConfig(ocm_size=ocm_size)
cgc_job = ChimeraJob(
    model_p=str(quantized_onnx_model.model_path),
    hw_config=hw_config,
    trange_file=str(quantized_onnx_model.tensor_ranges_path),
)
cgc_job.compile(quiet=True)
print(cgc_job)

Demo

Inference

NUM_IMAGES = 3
if MODEL_NAME.endswith("_hrf"):
    all_image_paths = glob.glob(f'{dataset_dir / "hrf" / "*.jpg"}')
    all_image_paths = random.sample(all_image_paths, NUM_IMAGES)
else:
    all_image_paths = [
        "../../../common/calibration/coco-like/33823288584_1d21cf0a26_k.jpeg",
        "../../../common/calibration/coco-like/17790319373_bd19b24cfc_k.jpeg",
        "../../../common/calibration/coco-like/19064748793_bb942deea1_k.jpeg",
    ]

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

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

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

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

Run Statistics

print(cgc_job)
cgc_job.plot_run_statistics()

Display Inference Results

%matplotlib inline

visualize_object = VisualizeObject.SemanticSegmentation
visualize_handle = get_visualize_handle(
    visualize_object, mmseg_parameters.visual_threshold, random_seed=17
)

for i in range(len(all_images)):
    layouter = DetectorsLayouter(
        visualize_handle,
        image=all_image_paths[i],
        classes=mmseg_parameters.classes,
    )
    for inference_engine, all_outputs in outputs_per_inference_engine.items():
        layouter.add_data(
            bboxes=None,
            masks=all_outputs[i][0][0],
            title=f"{MODEL_NAME}: {str(inference_engine).upper()}",
        )

    layouter.display()

Citation

@misc{mmseg2020,
    title={{MMSegmentation}: OpenMMLab Semantic Segmentation Toolbox and Benchmark},
    author={MMSegmentation Contributors},
    howpublished = {\url{https://github.com/open-mmlab/mmsegmentation}},
    year={2020}
}

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

@dataset{High Resolution Fundus,
  author={Attila Budai and Jan Odstrcilik},
  title={High Resolution Fundus (HRF) Image Database},
  year={2013},
  url={https://www5.cs.fau.de/research/data/fundus-images/}
}={2021}