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))+"%")
        col2.metric(predictions[1]['label'], str(round(predictions[1]['score']*100, 1))+"%")
        col6.metric(predictions[2]['label'], str(round(predictions[2]['score']*100, 1))+"%")
    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.image("./demo.jpg", use_column_width=True)
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_img = Image.open("./demo.jpg")
compute(demo_img)

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