Spaces:
Sleeping
Sleeping
from __future__ import annotations | |
from importlib import import_module | |
from pathlib import Path | |
import gradio as gr | |
import gradio.inputs | |
import gradio.outputs | |
import numpy as np | |
import os | |
from anomalib.deploy import Inferencer | |
def get_inferencer(weight_path: Path, metadata_path: Path | None = None) -> Inferencer: | |
"""Parse args and open inferencer. | |
Args: | |
weight_path (Path): Path to model weights. | |
metadata_path (Path | None, optional): Metadata is required for OpenVINO models. Defaults to None. | |
Raises: | |
ValueError: If unsupported model weight is passed. | |
Returns: | |
Inferencer: Torch or OpenVINO inferencer. | |
""" | |
inferencer: Inferencer | |
module = import_module("anomalib.deploy") | |
openvino_inferencer = getattr(module, "OpenVINOInferencer") | |
print(f"weight path: {weight_path}") | |
print(f"metadata path: {metadata_path}") | |
inferencer = openvino_inferencer(path=weight_path, metadata_path=metadata_path) | |
return inferencer | |
def infer(radio: str, image: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: | |
"""Inference function, return anomaly map, score, heat map, prediction mask ans visualisation. | |
Args: | |
image (np.ndarray): image to compute | |
inferencer (Inferencer): model inferencer | |
Returns: | |
tuple[np.ndarray, float, np.ndarray, np.ndarray, np.ndarray]: | |
heat_map, pred_mask, segmentation result. | |
""" | |
# Perform inference for the given image. | |
print(f"Radio Value: {radio.lower()}") | |
print(f"{os.getcwd()}") | |
weight_path = f"cfa/mvtec/{radio.lower()}/run/weights/openvino/model.onnx" | |
metadata_path = f"cfa/mvtec/{radio.lower()}/run/weights/openvino/metadata.json" | |
inferencer = get_inferencer(weight_path, metadata_path) | |
predictions = inferencer.predict(image=image) | |
return (predictions.heat_map, predictions.pred_mask, predictions.segmentations) | |
if __name__ == "__main__": | |
interface = gr.Interface( | |
fn=lambda radio, image: infer(radio, image), | |
inputs=[ | |
gr.Radio( | |
[ | |
"Bottle", | |
"Cable", | |
"Capsule", | |
"Carpet", | |
"Grid", | |
"Hazelnut", | |
"Leather", | |
"Metal_nut", | |
"Pill", | |
"Screw", | |
"Tile", | |
"Toothbrush", | |
"Transistor", | |
"Wood", | |
"Zipper", | |
], | |
label="MVTEC Class Name", | |
value="Bottle", | |
).style(height=400), | |
gradio.inputs.Image( | |
shape=None, image_mode="RGB", source="upload", tool="editor", type="numpy", label="Image" | |
).style(height=350), | |
], | |
outputs=[ | |
gradio.outputs.Image(type="numpy", label="Predicted Heat Map").style(height=200), | |
gradio.outputs.Image(type="numpy", label="Predicted Mask").style(height=200), | |
gradio.outputs.Image(type="numpy", label="Segmentation Result").style(height=200), | |
], | |
examples=[ | |
["Bottle", "sample_images/bottle.png"], | |
["Cable", "sample_images/cable.png"], | |
["Capsule", "sample_images/capsule.png"], | |
["Carpet", "sample_images/carpet.png"], | |
["Grid", "sample_images/grid.png"], | |
["Hazelnut", "sample_images/hazelnut.png"], | |
["Leather", "sample_images/leather.png"], | |
["Metal_nut", "sample_images/metal_nut.png"], | |
["Pill", "sample_images/pill.png"], | |
["Screw", "sample_images/screw.png"], | |
["Tile", "sample_images/tile.png"], | |
["Toothbrush", "sample_images/toothbrush.png"], | |
["Transistor", "sample_images/transistor.png"], | |
["Wood", "sample_images/wood.png"], | |
["Zipper", "sample_images/zipper.png"], | |
], | |
title="Anomaly Detection", | |
description="Anomlay Detection on Industrial Images", | |
css=".output-image, .image-preview {height: 300px !important}", | |
allow_flagging="never", | |
) | |
interface.launch(share=False) | |