import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
import json
import os

# Load metadata
CLASSES = ['Minor', 'Serious', 'Fatal']
IMG_SIZE = 224

if os.path.exists("metadata.json"):
    with open("metadata.json") as f:
        meta = json.load(f)
        CLASSES = meta.get("classes", CLASSES)
        IMG_SIZE = meta.get("img_size", IMG_SIZE)

# Load model
model = tf.keras.models.load_model("model.keras")

def preprocess(image):
    img = image.resize((IMG_SIZE, IMG_SIZE))
    img = np.array(img, dtype=np.float32) / 255.0
    return np.expand_dims(img, 0)

def predict(image):
    img_array = preprocess(image)
    
    preds = model.predict(img_array)[0]
    idx = int(np.argmax(preds))

    probs = {
        CLASSES[i]: float(preds[i] * 100)
        for i in range(len(CLASSES))
    }

    return {
        "severity": CLASSES[idx],
        "confidence": f"{preds[idx]*100:.2f}%",
        "probabilities": probs
    }

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="Accident Severity Prediction"
)

iface.launch()