Spaces:
Running
Running
File size: 2,303 Bytes
098c98c 32f1589 098c98c 32f1589 098c98c 32f1589 098c98c 32f1589 098c98c 32f1589 098c98c 508121b 32f1589 098c98c 32f1589 098c98c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from typing import *
import importlib
from pathlib import Path
import cv2
import gradio as gr
import numpy as np
available_modes = []
inferencer = None
last_model_path: str = None
if importlib.util.find_spec("openvino"):
available_modes.append("OpenVINO")
from inference.openvino import OpenVINOInferencer
def predict_openvino(
image: Union[str, Path, np.ndarray],
model_path: Union[str, Path],
device: str,
threshold: float) -> Dict[str, np.ndarray]:
global inferencer, last_model_path
if not isinstance(inferencer, OpenVINOInferencer) or last_model_path != str(model_path):
inferencer = OpenVINOInferencer(
path = model_path,
device=device.upper()
)
last_model_path = str(model_path)
inferencer.metadata["pixel_threshold"] = threshold
return inferencer(image)
def draw_contour(image, mask, color=(255,0,0), thickness=3):
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
return cv2.drawContours(image.copy(), contours, -1, color, thickness)
def convert_to_heatmap(heatmap):
heatmap = cv2.applyColorMap((heatmap*255).astype("uint8"), cv2.COLORMAP_HSV)
return heatmap
def predict(image, threshold, device, model_dir):
outputs = predict_openvino(image, model_dir, device, threshold)
out_image = draw_contour(image, outputs["pred_mask"])
heatmap = convert_to_heatmap(outputs["anomaly_map"])
return out_image, heatmap
def launch():
input_image = gr.Image(
label="Input image",
value="images/171436008_Fail.jpeg"
)
threshold = gr.Slider(value=5.5, step=0.1, label="Threshold")
devices = gr.Radio(
label="Device",
choices=["AUTO", "CPU", "GPU"],
value="CPU",
interactive=False
)
output_image = gr.Image(label="Output image")
output_heatmap = gr.Image(label="Heatmap")
model = gr.Text(label="Model", interactive=False, value="models/glass-tr-5-ov")
intf = gr.Interface(
title="Anomaly Detection",
fn=predict,
inputs=[input_image, threshold, devices, model],
outputs=[output_image, output_heatmap]
)
intf.launch()
if __name__ == "__main__":
launch() |