import tensorflow as tf
from tensorflow.keras.applications import EfficientNetB0

efficient_net = EfficientNetB0(weights='imagenet',include_top=False,input_shape=(150, 150, 3))

model = efficient_net.output
model = tf.keras.layers.GlobalAveragePooling2D()(model)
model = tf.keras.layers.Dense(64, activation='relu')(model)
model = tf.keras.layers.Dropout(rate=0.1)(model)
model = tf.keras.layers.Dense(32, activation='relu')(model)
model = tf.keras.layers.Dropout(rate=0.1)(model)
model = tf.keras.layers.Dense(2, activation='sigmoid')(model)
model = tf.keras.models.Model(inputs=efficient_net.input, outputs=model)

model.compile(loss='binary_crossentropy',optimizer = 'Adam', metrics= ['accuracy'])

model.load_weights('./checkpoint')

import gradio as gr

def cardiomegaly(img):
    img = img.reshape(1, 150, 150, 3)
    prediction = model.predict(img).tolist()[0]
    class_names = ["False", "True"]
    return {class_names[i]: prediction[i] for i in range(2)}
#set the user uploaded image as the input array
#match same shape as the input shape in the model
im = gr.inputs.Image(shape=(150, 150), image_mode='RGB', invert_colors=False, source="upload")
#setup the interface
gr.Interface(
    fn = cardiomegaly, 
    inputs = im, 
    outputs = gr.outputs.Label(),
).launch()