app.py

import gradio as gr
import tensorflow as tf
import joblib
import numpy as np
import zipfile
import os
import re

# Step 1: Unzip models only once
unzip_dir = "unzipped_models"
zip_file = "Models.zip"  # Make sure this is the filename inside your Space repo

if not os.path.exists(unzip_dir):
    print("Extracting model zip file...")
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        zip_ref.extractall(unzip_dir)
    print("Extraction complete.")

# Step 2: Parse folders to dynamically populate dropdowns

model_root = os.path.join(unzip_dir, 'Models')  # Adjust this if your ZIP structure is different

activations = []
seeds_dict = dict()
neurons_dict = dict()

for act in os.listdir(model_root):
    act_path = os.path.join(model_root, act)
    if os.path.isdir(act_path) and not act.startswith("linear_models"):
        activations.append(act)
        seeds = []
        for seed_folder in os.listdir(act_path):
            seed_path = os.path.join(act_path, seed_folder)
            if os.path.isdir(seed_path):
                seeds.append(seed_folder)
                neuron_list = []
                for model_file in os.listdir(seed_path):
                    match = re.match(r"model_(\d+)\.keras", model_file)
                    if match:
                        neuron_list.append(int(match.group(1)))
                neurons_dict[(act, seed_folder)] = sorted(neuron_list)
        seeds_dict[act] = sorted(seeds)

activations = sorted(activations)

def update_seeds(activation):
    return {"choices": seeds_dict[activation], "value": seeds_dict[activation][0]}

def update_neurons(activation, seed):
    neurons = neurons_dict[(activation, seed)]
    return {"choices": neurons, "value": neurons[0]}

# Step 3: Load linear models
# linear_rgb_path = os.path.join(unzip_dir, "linear_models/linear_rgb.joblib")
# linear_grey_path = os.path.join(unzip_dir, "linear_models/linear_grey.joblib")

# linear_rgb = joblib.load(linear_rgb_path)
# linear_grey = joblib.load(linear_grey_path)

# Step 4: Prediction function
def predict(r, g, b, activation, seed, neurons):
    try:
        X = np.array([[r, g, b]])
        
        # Linear predictions
        lin_pred_rgb = (1.9221 * r) - (1.3817 * g) + (1.4058 * b) - 0.1318

        # ANN prediction
        keras_path = os.path.join(unzip_dir, activation, seed, f"model_{neurons}.keras")
        if not os.path.exists(keras_path):
            raise FileNotFoundError(f"Model not found: {keras_path}")
        
        model = tf.keras.models.load_model(keras_path)
        ann_pred = model.predict(X)[0][0]

        return ann_pred, lin_pred_rgb

    except Exception as e:
        return f"Error: {str(e)}", "", ""

# Dynamic components for UI

def update_seeds(activation):
    return gr.Dropdown.update(choices=seeds_dict[activation], value=seeds_dict[activation][0])

def update_neurons(activation, seed):
    neurons = neurons_dict[(activation, seed)]
    return gr.Dropdown.update(choices=neurons, value=neurons[0])

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# ANN vs Linear Model Predictor")
    gr.Markdown("Dynamically select models and predict cholesterol concentration.")
    
    with gr.Row():
        r = gr.Number(label="R")
        g = gr.Number(label="G")
        b = gr.Number(label="B")

    with gr.Row():
        activation = gr.Dropdown(choices=activations, label="Activation Function")
        seed = gr.Dropdown(label="Seed")
        neurons = gr.Dropdown(label="Neurons")

    activation.change(update_seeds, inputs=[activation], outputs=[seed])
    seed.change(update_neurons, inputs=[activation, seed], outputs=[neurons])

    with gr.Row():
        btn = gr.Button("Predict")
    
    with gr.Row():
        ann_output = gr.Text(label="ANN Model Prediction")
        lin_rgb_output = gr.Text(label="Linear RGB Prediction")
        # lin_grey_output = gr.Text(label="Linear Grey Prediction")

    btn.click(
        fn=predict, 
        inputs=[r, g, b, activation, seed, neurons],
        outputs=[ann_output, lin_rgb_output]
    )

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