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

# Load the trained OCR model
model = load_model("model_ocr.h5")

# Define a function to preprocess the image
def preprocess_image(image):
    image = np.array(image.convert('L'))  # Convert to grayscale
    image = cv2.resize(image, (width, height))  # Resize to expected input shape
    image = image / 255.0  # Normalize pixel values
    image = np.expand_dims(image, axis=-1)  # Add channel dimension if necessary
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Define the prediction function
def predict(image):
    preprocessed_image = preprocess_image(image)
    prediction = model.predict(preprocessed_image)
    predicted_text = decode_prediction(prediction)  # Implement your decoding function here
    return predicted_text

# Create a Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),  # Updated to use gr.Image directly
    outputs="text",
    title="Captcha OCR",
    description="An OCR model to read captcha text from images."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()