import cv2
import numpy as np
import gradio as gr
import tensorflow as tf
from huggingface_hub import hf_hub_download

# Download the entire model directory
model_dir = hf_hub_download(repo_id="Par24/sign_language", filename="saved_model", repo_type="model")

# Load the model correctly
model = tf.saved_model.load(model_dir)
infer = model.signatures["serving_default"]

# Define class labels
class_labels = ['Hello', 'Yes', 'No', 'Thank You', 'Please']

def predict_sign(frame):
    # Convert BGR (OpenCV) to RGB (TensorFlow format)
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # Preprocess the frame
    img = cv2.resize(frame, (224, 224))  # Resize
    img = img / 255.0  # Normalize
    img = np.expand_dims(img, axis=0)  # Add batch dimension
    img = tf.convert_to_tensor(img, dtype=tf.float32)
    
    # Make prediction
    predictions = infer(tf.constant(img))
    output_tensor_name = list(predictions.keys())[0]  # Get the output tensor name
    predictions = predictions[output_tensor_name].numpy()
    
    # Get predicted class and confidence
    predicted_class = class_labels[np.argmax(predictions)]
    confidence = np.max(predictions)
    
    return predicted_class, confidence

def process_frame(frame):
    pred, conf = predict_sign(frame)

    # Convert RGB back to BGR for OpenCV
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    # Overlay prediction text
    cv2.putText(frame, f"{pred} ({conf:.2f})", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

    return frame

# Gradio Live Webcam Interface
gui = gr.Interface(
    fn=process_frame,  # Function to process frames
    inputs="webcam",   # Use webcam as input
    outputs="image",   # Output is an image
    live=True
)

# Launch Gradio App
if __name__ == "__main__":
    gui.launch(server_name="0.0.0.0", server_port=7860)