import streamlit as st
from transformers import pipeline as pip
from PIL import Image

# set page setting
st.set_page_config(page_title='Smoke & Fire Detection')

# set history var
if 'history' not in st.session_state:
    st.session_state.history = []

@st.cache(persist=True)
def loadModel():
    pipeline = pip(task="image-classification", model="EdBianchi/vit-fire-detection")
    return pipeline

# PROCESSING
def compute(image):
    predictions = pipeline(image)

    with st.container():
        st.image(image, use_column_width=True)

    with st.container():
        st.write("### Classification Outputs:")
        col1, col2, col6 = st.columns(3)
        col1.metric(predictions[0]['label'], str(round(predictions[0]['score']*100, 1))+"%", delta = 50)
        col2.metric(predictions[1]['label'], str(round(predictions[1]['score']*100, 1))+"%", delta = 50)
        col6.metric(predictions[2]['label'], str(round(predictions[2]['score']*100, 1))+"%", delta = 50)
    return None

# INIT
with st.spinner('Loading the model, this could take some time...'):
    pipeline = loadModel()

# TITLE
st.write("# Fire in Forest Environments")
st.write("""Wildfires or forest fires are *unpredictable catastrophic and destructive* events that affect *rural areas*. 
The impact of these events affects both *vegetation and wildlife*. 

This application showcases the *vit-fire-detection* model, a version of google **vit-base-patch16-224-in21k** vision transformer *fine-tuned for smoke and fire detection*. In particular, we can imagine a setup in which webcams, drones, or other recording devices *take pictures of a wild environment every t seconds or minutes*. The proposed system is then able to classify the current situation as *normal, smoke, or fire*. 
""")

st.write("### Upload an image to see the classifier in action")
# INPUT IMAGE
file_name = st.file_uploader("")
if file_name is not None:
    image = Image.open(file_name)
    compute(image)

# DEMO IMAGE
demo_img = Image.open("./demo.jpg")
compute(demo_img)

# SIDEBAR
#st.sidebar.write("""""")