import json
import numpy as np
import gradio as gr
import tensorflow as tf

MODEL_PATH = 'best_pneumonia_model.keras'
META_PATH = 'model_meta.json'

model = tf.keras.models.load_model(MODEL_PATH)
with open(META_PATH, 'r') as f:
    meta = json.load(f)

label_names = meta['label_names']
img_size = int(meta['img_size'])
best_model_name = meta.get('best_model_name', '')

def preprocess(img):
    x = img.convert('L').resize((img_size, img_size))
    arr = np.array(x).astype(np.float32) / 255.0
    arr = arr[..., None]

    if best_model_name == 'Scratch_CNN':
        pass
    else:
        arr = np.repeat(arr, 3, axis=-1) * 255.0
        if 'MobileNetV2' in best_model_name:
            arr = tf.keras.applications.mobilenet_v2.preprocess_input(arr)
        elif 'EfficientNetB0' in best_model_name:
            arr = tf.keras.applications.efficientnet.preprocess_input(arr)
        arr = np.array(arr, dtype=np.float32)

    return np.expand_dims(arr, axis=0)

def predict(image):
    x = preprocess(image)
    probs = model.predict(x, verbose=0)[0]
    return {label_names[i]: float(probs[i]) for i in range(len(label_names))}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type='pil'),
    outputs=gr.Label(num_top_classes=3),
    title='Pneumonia Classification',
    description='Upload chest image and get predicted class with probability.'
)

if __name__ == '__main__':
    demo.launch(server_name='0.0.0.0', server_port=7860)