import gradio as gr
import numpy as np
import cv2
import tensorflow as tf

model = tf.keras.models.load_model('best_model.keras')

class_names = [
    'among us', 'apex legends', "baldur's gate 3", 'btd6', 'content warning',
    'csgo', 'cyber punk 2077', 'darkest dungeon', 'doom eternal', 'fallout 3',
    'fallout 4', 'fallout new vegas', 'fnaf security breach', 'fortnite',
    'genshin impact', 'guilty gear strive', 'honkaistarrail', 'minecraft',
    'overwatch 2', 'phasmophobia', 'rainbow six siege', 'resident evil village',
    'slay the spire', 'stardew valley', 'street fighter 6', 'subnautica',
    'terraria', 'valorant', 'wizard 101', 'wuthering waves', 'yugioh master duel'
]

def predict_image(image):
    if image is None:
        return {"Error": 0.0}
    img = cv2.resize(image, (96, 96))
    img_array = np.array(img, dtype=np.float32) / 255.0
    img_batch = np.expand_dims(img_array, axis=0)
    predictions = model.predict(img_batch)
    confidences = {class_names[i]: float(predictions[0][i]) for i in range(len(class_names))}
    return confidences

interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy", label="Upload a Gameplay Screenshot"),
    outputs=gr.Label(num_top_classes=3, label="Top 3 Predictions"),
    title="🎮 Game Environment Classifier",
    description="Drag and drop a gameplay screenshot, and the AI will predict which game it is from!"
)

interface.launch()