import gradio as gr
import cv2
import numpy as np
from anomalib.deploy import OpenVINOInferencer
from huggingface_hub import hf_hub_download
from anomalib import TaskType
from PIL import Image

# Load the model
model_path = hf_hub_download(repo_id="suidinpa/bottle-anomaly-detection", filename="model.bin")
metadata_path = hf_hub_download(repo_id="suidinpa/bottle-anomaly-detection", filename="metadata.json")

# Initialize OpenVINO inferencer
inferencer = OpenVINOInferencer(
    path=model_path,
    metadata=metadata_path,
    device="CPU",
    task="classification"
)

def predict(image):
    """Function to process the image and predict with the model"""
    # Convert image to a format suitable for the model
    img_array = np.array(image)
    img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

    # Predict anomaly
    predictions = inferencer.predict(img_bgr)

    # Return results
    return predictions.pred_label, predictions.pred_score

# Define the Gradio interface
interface = gr.Interface(
    fn=predict,  # The prediction function
    inputs=gr.Image(type="numpy", label="Upload Image"),  # Image input from the user
    outputs=["text", "number"],  # Text (for label) and number (for prediction score)
    live=True  # Make the app live for real-time predictions
)

# Launch the app
interface.launch()