import gradio as gr

from timeit import default_timer as timer
from typing import Tuple , Dict
import tensorflow as tf 
import numpy as np

from PIL import Image
import os

# 1.Import and class names setup
class_names = ['CNV','DME','DRUSEN','NORMAL']


# 2. Model annd transforms prepration
model = tf.keras.models.load_model(
    'best_model_lg.keras', custom_objects=None, compile=True, safe_mode=True
)

# Load save weights

# 3.prediction function (predict())

def load_and_prep_imgg(filename, img_shape=224, scale=True):
    img = tf.io.read_file(filename)
    img = tf.io.decode_image(img)
    img = tf.image.resize(img, size=[img_shape, img_shape])
    if scale:
        return img / 255
    else:
        return img
        
def predict(img) -> Tuple[Dict,float] :

  start_time = timer()

  image =  load_and_prep_imgg(img)
  image = Image.open(image)  

  pred_img = model.predict(tf.expand_dims(img, axis=0))
  pred_class = class_names[pred_img.argmax()]
  st.write(f"Predicted brain tumor is: {pred_class} with probability: {pred_img.max():.2f}")

  

  
  end_time = timer()
  pred_time = round(end_time - start_time , 4)

  return pred_class , pred_time

### 4. Gradio app - our Gradio interface + launch command

title = 'FoodVision Big'
description = 'Feature Extraxtion VGG  model to classifiy Macular Diseases by OCT '
article = 'created at Tensorflow Model Deployment'

# Create example list

example_list = [['examples/'+ example] for example in os.listdir('examples')]
example_list

# create a gradio demo
demo = gr.Interface(fn=predict ,
                    inputs=gr.Image(type='pil'),
                    outputs=[gr.Label(num_top_classes = 3 , label= 'prediction'),
                            gr.Number(label= 'Prediction time (s)')],
                    examples = example_list,
                    title = title,
                    description = description,
                    article= article)

# Launch the demo
demo.launch(debug= False)