import streamlit as st
from fastai.vision.all import *

class Pipeline:
    IS_IT_A_SINK_CLASSIFIER = load_learner("is_it_a_sink_classifier_11:34_11:26:2022.pkl")
    DIRTY_OR_CLEAN_SINK_CLASSIFIER = load_learner("dirty_or_clean_classifier_11:33_11:26:2022.pkl")

    def __init__(self):
        pass

    def classify(self, image):
        classification, class_idx, prediction_probs = Pipeline.IS_IT_A_SINK_CLASSIFIER.predict(image)

        if classification == "not_a_sink":
            return classification, class_idx, prediction_probs

        return Pipeline.DIRTY_OR_CLEAN_SINK_CLASSIFIER.predict(image)

st.title("Clean sink? Or dirty sink?")

file_name = st.file_uploader("Upload a picture of a sink")

classification_pipeline = Pipeline()

classification_statements = {
    "not_a_sink": "This isn't a sink silly.",
    "dirty_sink": "This is a dirty sink.",
    "clean_sink": "This is a clean sink."
}

probability_modifier = lambda prob : col2.write("Though I might be wrong...") if prob < 0.8 else None

if file_name is not None:
    col1, col2 = st.columns(2)

    image = Image.open(file_name)
    col1.image(image, use_column_width=True)
    
    classification, class_idx, prediction_probs = classification_pipeline.classify(PILImage.create(file_name))

    print(f"classification={classification}")
    print(f"class_idx={class_idx}")
    print(f"prediction_probs={prediction_probs}")

    col2.header("Classification")
    col2.subheader(classification_statements[classification])
    probability_modifier(prediction_probs[class_idx].numpy())