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("#### Different classification outputs at different threshold values:")
        col1, col2, col6 = st.columns(3)
        col1.metric(predictions[0]['label'], round(predictions[0]['score']+100, 1)+"%")
        col2.metric(predictions[1]['label'], round(predictions[1]['score']+100, 1)+"%")
        col6.metric(predictions[2]['label'], 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("# Smoke & Fire Detection in Forest Environments")
st.write("#### Upload an Image to see the classifier in action")

# INPUT IMAGE
file_name = st.file_uploader("Upload an image")
if file_name is not None:
    image = Image.open(file_name)
    compute(image)

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