import gradio as gr
import h2o
import numpy as np
import pandas as pd
import cv2
from skimage.color import rgb2hsv
from skimage.measure import shannon_entropy
from scipy.ndimage import generic_filter

# Initialize H2O and load the saved model
h2o.init()
model_path = "GBM_grid_1_AutoML_2_20241228_54907_model_10"  # Replace with your H2O model path
h2o_model = h2o.load_model(model_path)

# Feature extraction function
def extract_features(image):
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Extract RGB means
    meanr = np.mean(image[:, :, 0])
    meang = np.mean(image[:, :, 1])
    meanb = np.mean(image[:, :, 2])

    # Convert to HSI and compute HHR
    hsv_image = rgb2hsv(image)
    hue = hsv_image[:, :, 0]
    high_hue_pixels = np.sum(hue > 0.95)
    total_pixels = hue.size
    HHR = high_hue_pixels / total_pixels

    # Convert to Grayscale
    gray_image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Compute Entropy and Brightness
    Ent = shannon_entropy(gray_image)
    B = np.mean(gray_image)

    # Sliding window filters
    def g1_filter(window):
        return window[4] - np.min(window)

    def g2_filter(window):
        return np.max(window) - window[4]

    def g3_filter(window):
        return window[4] - np.mean(window)

    def g4_filter(window):
        return np.std(window)

    def g5_filter(window):
        return window[4]

    g1 = generic_filter(gray_image, g1_filter, size=3).mean()
    g2 = generic_filter(gray_image, g2_filter, size=3).mean()
    g3 = generic_filter(gray_image, g3_filter, size=3).mean()
    g4 = generic_filter(gray_image, g4_filter, size=3).mean()
    g5 = generic_filter(gray_image, g5_filter, size=3).mean()

    return {
        "meanr": meanr,
        "meang": meang,
        "meanb": meanb,
        "HHR": HHR,
        "Ent": Ent,
        "B": B,
        "g1": g1,
        "g2": g2,
        "g3": g3,
        "g4": g4,
        "g5": g5,
    }

# Prediction function
def predict(image, gender, age):
    # Extract image features
    features = extract_features(image)
    features["gender"] = gender
    features["age"] = age

    # Convert features to DataFrame
    features_df = pd.DataFrame([features])
    features_h2o = h2o.H2OFrame(features_df)

    # Predict using the model
    prediction = h2o_model.predict(features_h2o)
    return prediction.as_data_frame().iloc[0, 0]

# Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(label="Upload Image"),
        gr.Dropdown(choices=["Male", "Female"], label="Gender"),
        gr.Slider(0, 100, step=1, label="Age"),
    ],
    outputs="label",
    title="Image-based Prediction App",
    description="Upload an image, enter your gender and age, and get predictions using the pre-trained model."
)

# Launch the app
interface.launch()