import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.models import load_model 
import keras

# 1. Load your model
@keras.saving.register_keras_serializable()
def scaling(x, scale=255.0, **kwargs):
    return x * scale

@keras.saving.register_keras_serializable()
def l2_normalize(x, axis=-1, epsilon=1e-10):
    return tf.nn.l2_normalize(x, axis=axis, epsilon=epsilon)

# Pass the registered functions explicitly
custom_objects = {"scaling": scaling, "l2_normalize": l2_normalize}

# Load the model with custom_objects
model = load_model("facenet_real_fake_classifier_final.keras", custom_objects=custom_objects)

# 2. Define your inference function
def detect_forgery(image):
    # Preprocess the image to match your model’s input requirements
    img = tf.image.resize(image, (160, 160))  # Example size; adjust for your model
    img = tf.expand_dims(img, axis=0)
    img = img / 255.0  # Example normalization; adapt as needed

    # Run inference
    predictions = model.predict(img)[0]  
    
    pred_value = float(predictions)  # or pred_value.item()
    # Suppose predictions = [prob_real, prob_fake]


    # Format output
    # You can return a dictionary or a string. For example:
    if pred_value >= 0.5:
        # It's real, use p as the confidence
        confidence = pred_value
        label = "REAL"
    else:
        # It's fake, confidence is (1 - p)
        confidence = 1.0 - pred_value
        label = "FAKE"

    return f"{label} with confidence {confidence:.2f}"

# 3. Build your Gradio interface
demo = gr.Interface(
    fn=detect_forgery,
    inputs=gr.Image(type="numpy"),  # 'type="numpy"' gives a NumPy array
    outputs="text",
    title="Face Forgery Detector",
    description="Upload a face image to check if it's likely forged or real."
)

# 4. Launch the app (Gradio handles the rest)
if __name__ == "__main__":
    _ = model.predict(np.zeros((1, 160, 160, 3), dtype=np.float32))
    demo.launch()