import os
os.system('pip install tf_keras')
import tf_keras #Keras3 no longer supports hub.KerasLayer()

import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
import gradio as gr

model_path = "./4-17_ViT-B16_Lung3_v4_E30/"
model = tf_keras.models.load_model(model_path) #Load model with old version (tf_keras instead of tf.keras)

# Define preprocessing function
def preprocess_image(image_path):
    # Load the image from file
    img = image.load_img(image_path, target_size=(224, 224))
    
    # Convert the image to a numpy array
    img_array = image.img_to_array(img)
    
    # Rescale the pixel values to the range [0, 1]
    img_array /= 255.0

    return img_array

# Define inference function
def predict(image_path):
    # Preprocess the image
    preprocessed_image = preprocess_image(image_path)

    # Add batch dimension to input image
    input_tensor = tf.expand_dims(preprocessed_image, axis=0)
    
    # Perform inference
    predictions = model.predict(input_tensor)

    predicted_probabilities = np.squeeze(predictions)
    predicted_class = 'Pneumonia' if predicted_probabilities > 0.5 else 'Normal' 
    confidence_score = predicted_probabilities if predicted_class == 'Pneumonia' else 1 - predicted_probabilities

    return "Prediction: " + predicted_class, {'Confidence: ': confidence_score}

image_component = gr.Image(type="filepath", label="Chest X-ray Image")
label_component = gr.Label(label="Classification result", num_top_classes=1)
confidence_component = gr.Label(label="Confidence Score (%)")

description = ""
title = "Pneumonia Classifier"

demo = gr.Interface(fn=predict, 
                    inputs=image_component, 
                    outputs=[label_component, confidence_component], 
                    description=description, 
                    title=title,
                    examples=["demo_example_imgs/NORMAL-0001-0001.jpeg",
                              "demo_example_imgs/NORMAL2-IM-0337-0001.jpeg",
                              "demo_example_imgs/PNEUMONIA_bacteria_434.jpeg",
                              "demo_example_imgs/PNEUMONIA_virus_140.jpeg"
                             ])
    
demo.launch()