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

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

def predict_digit(image):
    if image is None:
        return {"error": "No image provided"}
    
    try:
        if len(image.shape) == 3:
            img = Image.fromarray(image.astype('uint8')).convert('L')
        else:
            img = Image.fromarray(image.astype('uint8'))
        
        img = img.resize((28, 28))
        img_array = np.array(img)
        img_array = 255 - img_array
        img_array = img_array / 255.0
        img_array = img_array.reshape(1, 28, 28, 1)
        
        prediction = model.predict(img_array, verbose=0)
        digit = int(np.argmax(prediction))
        confidence = float(np.max(prediction) * 100)
        probabilities = [float(p) for p in prediction[0]]
        
        return {
            "digit": digit,
            "confidence": confidence,
            "probabilities": probabilities
        }
        
    except Exception as e:
        return {"error": str(e)}

iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(type="numpy", label="Upload digit image"),
    outputs=gr.JSON(label="Prediction Result"),
    title="Digit Detector API"
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)