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

if not os.path.exists("keras_model.h5"):
    raise FileNotFoundError("Model file not found!")

model = tf.keras.models.load_model("keras_model.h5")

# Load Teachable Machine model
model = tf.keras.models.load_model("./keras_model.h5")
# Labels from your model — change if needed
labels = ["Good Posture", "Bad Posture"]

def predict_from_webcam(frame):
    # Convert to RGB if needed (Gradio usually provides RGB)
    if frame.shape[-1] == 4:
        frame = frame[:, :, :3]  # remove alpha channel

    # Resize to 224x224 (Teachable Machine default)
    img = cv2.resize(frame, (224, 224))
    img = img / 255.0  # normalize to [0, 1]
    img = np.expand_dims(img, axis=0)  # add batch dimension

    # Prediction
    prediction = model.predict(img)[0]
    label = labels[np.argmax(prediction)]
    confidence = np.max(prediction)

    return f"{label} ({confidence * 100:.2f}%)"

# Gradio interface
demo = gr.Interface(
    fn=predict_from_webcam,
    inputs=gr.Image(source="webcam", streaming=True),
    outputs=gr.Text(),
    live=True,
    title="🪑 Posture Detector",
    description="Detects good or bad posture from webcam using a Teachable Machine model."
)

demo.launch()