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()