app.py

"""Anomalib Gradio Script.

This script provide a gradio web interface
"""

from argparse import ArgumentParser, Namespace
from importlib import import_module
from pathlib import Path
from typing import Tuple, Union

import gradio as gr
import gradio.inputs
import gradio.outputs
import numpy as np
from omegaconf import DictConfig, ListConfig
from skimage.segmentation import mark_boundaries

from anomalib.config import get_configurable_parameters
from anomalib.deploy.inferencers.base import Inferencer
from anomalib.post_processing import compute_mask, superimpose_anomaly_map


def infer(
    image: np.ndarray, inferencer: Inferencer, threshold: float = 50.0
) -> Tuple[np.ndarray, float, np.ndarray, np.ndarray, np.ndarray]:
    """Inference fonction, return anomaly map, score, heat map, prediction mask ans visualisation.

    :param image: image
    :type image: np.ndarray
    :param inferencer: model inferencer
    :type inferencer: Inferencer
    :param threshold: threshold between 0 and 100, defaults to 50.0
    :type threshold: float, optional
    :return: anomaly_map, anomaly_score, heat_map, pred_mask, vis_img
    :rtype: Tuple[np.ndarray, float, np.ndarray, np.ndarray, np.ndarray]
    """
    # Perform inference for the given image.
    threshold = threshold / 100
    anomaly_map, anomaly_score = inferencer.predict(image=image, superimpose=False)
    heat_map = superimpose_anomaly_map(anomaly_map, image)
    pred_mask = compute_mask(anomaly_map, threshold)
    vis_img = mark_boundaries(image, pred_mask, color=(1, 0, 0), mode="thick")
    return anomaly_map, anomaly_score, heat_map, pred_mask, vis_img


def get_args() -> Namespace:
    """Get command line arguments.

    Returns:
        Namespace: List of arguments.
    """
    parser = ArgumentParser()
    parser.add_argument("--config", type=Path, default="/home/user/app/anomalib/models/padim/config.yaml", required=False, help="Path to a model config file")
    parser.add_argument("--weight_path", type=Path, default="/home/user/app/results/padim//mvtec/bottle/weights/model.ckpt", required=False, help="Path to a model weights")
    parser.add_argument("--meta_data", type=Path, required=False, help="Path to JSON file containing the metadata.")

    parser.add_argument(
        "--threshold",
        type=float,
        required=False,
        default=75.0,
        help="Value to threshold anomaly scores into 0-100 range",
    )

    parser.add_argument("--share", type=bool, required=False, default=False, help="Share Gradio `share_url`")

    args = parser.parse_args()

    return args


def get_inferencer(gladio_args: Union[DictConfig, ListConfig]) -> Inferencer:
    """Parse args and open inferencer."""
    config = get_configurable_parameters(config_path=gladio_args.config)
    # Get the inferencer. We use .ckpt extension for Torch models and (onnx, bin)
    # for the openvino models.
    extension = gladio_args.weight_path.suffix
    inferencer: Inferencer
    if extension in (".ckpt"):
        module = import_module("anomalib.deploy.inferencers.torch")
        TorchInferencer = getattr(module, "TorchInferencer")  # pylint: disable=invalid-name
        inferencer = TorchInferencer(
            config=config, model_source=gladio_args.weight_path, meta_data_path=gladio_args.meta_data
        )

    elif extension in (".onnx", ".bin", ".xml"):
        module = import_module("anomalib.deploy.inferencers.openvino")
        OpenVINOInferencer = getattr(module, "OpenVINOInferencer")  # pylint: disable=invalid-name
        inferencer = OpenVINOInferencer(
            config=config, path=gladio_args.weight_path, meta_data_path=gladio_args.meta_data
        )

    else:
        raise ValueError(
            f"Model extension is not supported. Torch Inferencer exptects a .ckpt file,"
            f"OpenVINO Inferencer expects either .onnx, .bin or .xml file. Got {extension}"
        )

    return inferencer


if __name__ == "__main__":
    session_args = get_args()

    gladio_inferencer = get_inferencer(session_args)

    iface = gr.Interface(
        fn=lambda image, threshold: infer(image, gladio_inferencer, threshold),
        inputs=[
            gradio.inputs.Image(
                shape=None, image_mode="RGB", source="upload", tool="editor", type="numpy", label="Image"
            ),
            gradio.inputs.Slider(default=session_args.threshold, label="threshold", optional=False),
        ],
        outputs=[
            gradio.outputs.Image(type="numpy", label="Anomaly Map"),
            gradio.outputs.Textbox(type="number", label="Anomaly Score"),
            gradio.outputs.Image(type="numpy", label="Predicted Heat Map"),
            gradio.outputs.Image(type="numpy", label="Predicted Mask"),
            gradio.outputs.Image(type="numpy", label="Segmentation Result"),
        ],
        title="Anomalib",
        description="Anomalib Gradio",
    )

    iface.launch(share=session_args.share)