import streamlit as st
import tensorflow as tf
import zipfile
from tensorflow.keras.models import load_model

# Unzipping model 
zip_ref = zipfile.ZipFile("L10_fine_tuned_model.zip", "r")
zip_ref.extractall()
zip_ref.close()

# Loading model
model_1=load_model("L10_fine_tuned_model")
model_2=load_model("L10_fine_tuned_model")

@st.cache(allow_output_mutation=True)
def load_and_predict_image(model,Input,scale=False):
  classes=['bacterial_leaf_blight', 'bacterial_leaf_streak', 'bacterial_panicle_blight', 'blast', 'brown_spot', 'dead_heart', 'downy_mildew', 'hispa', 'normal', 'tungro']
  #img=tf.io.read_file(Input)
  img=tf.io.decode_image(Input,channels=3)
  img=tf.image.resize(img,size=(224,224))
  if scale:
    img/255.
  else:
    img
  pred_probs=model_1.predict(tf.expand_dims(img,axis=0))
  pred_labels=tf.argmax(pred_probs,axis=1)[0].numpy()
  pred_class=classes[pred_labels]
  return pred_class
  
def main():
  st.set_page_config(page_title="Paddy Disease Detector",
                     page_icon="🌱",
                     layout="wide",
                     initial_sidebar_state="expanded")
  st.title("Paddy Disease Detector")
  st.write("A Computer vision model that classifies a given paddy leaf images into one of the nine disease categories or a normal leaf accurately.")
  col1, col2 = st.columns(2)
  with col1:
    uploaded_file=st.file_uploader("Click to upload paddy leaf")
    selected_model = st.selectbox("Select Model",("model 1", "model 2"), index=0)
    if uploaded_file is not None:
      uploaded_img=uploaded_file.read()
      
      col2.image(uploaded_file,width=500)
    predict=st.button("Predict!")
  if predict:
        if uploaded_file is not None:
            with st.spinner('Please Wait 👩‍🍳'):
                # setting model and rescalling 
                if selected_model == 'model 1':
                    pred_model = model_1 
                    pred_rescale = False
                else:
                    pred_model = model_2 
                    pred_rescale = False
                # makeing prediction 
                pred = load_and_predict_image(model=pred_model, Input=uploaded_img, scale=pred_rescale)
                
                col1.success(f"It's an {pred}")
        else:
          st.warning('Please Upload Image')

if __name__=='__main__': 
    main()