import numpy as np
import cv2 as cv
import joblib
import gradio as gr

# Load the trained model
loaded_model = joblib.load('model.pkl')

# Define the class names
class_names = ['glioma_tumor', 'meningioma_tumor', 'normal', 'pituitary_tumor']

def preprocess_image(image, target_size=(30, 10)):
    """Preprocess the input image for model prediction."""
    img = np.array(image)
    # Resize the image
    img_resized = cv.resize(image, target_size)
    # Convert image to grayscale
    gray_img = cv.cvtColor(img_resized, cv.COLOR_BGR2GRAY)
    # Flatten the image
    img_flattened = gray_img.flatten()
    # Make sure the feature length matches what the model expects
    if len(img_flattened) != 300:  # Adjust this number to match the expected feature length
        raise ValueError(f"Expected 300 features, but got {len(img_flattened)} features.")
    return gray_img, img_flattened

def predict_image(image):
    """Predict the class of the input image using the trained model and return the processed image."""
    # Preprocess the image
    processed_image, img_preprocessed = preprocess_image(image)
    # Make prediction
    prediction = loaded_model.predict([img_preprocessed])
    # Map the prediction to class names
    predicted_class = class_names[prediction[0]]
    # Convert processed_image to RGB format for display
    processed_image_rgb = cv.cvtColor(processed_image, cv.COLOR_GRAY2RGB)
    return predicted_class, processed_image_rgb

# Gradio interface
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy"),
    outputs=[gr.Text(), gr.Image()]
)

if __name__ == "__main__":
    demo.launch()