streamlit_app.py

import json
import os

import streamlit as st
from streamlit.logger import get_logger
from PIL import Image

from butt_or_bread.core import ButtBreadClassifier
from butt_or_bread.utils import health_check

# Create Streamlit logger
st_logger = get_logger(__name__)
st.set_option("deprecation.showfileUploaderEncoding", False)

# Load Streamlit configuration file
with open("streamlit_app.json") as cfg_file:
    st_app_cfg = json.load(cfg_file)


ui_cfg = st_app_cfg["ui"]
model_cfg = st_app_cfg["model"]
image_cfg = st_app_cfg["image"]

st.set_page_config(
    layout="centered",
    page_title=ui_cfg["title"],
    page_icon=ui_cfg["icon"],
)


@st.cache(allow_output_mutation=True, suppress_st_warning=True, max_entries=3, ttl=300)
def get_classifier():
    """Allow butt_bread model caching"""
    classifier = ButtBreadClassifier(model_url=model_cfg["url"])
    classifier.download()
    classifier.initialize()

    return classifier


if __name__ == "__main__":

    image_file, image, prediction = None, None, None
    classifier = get_classifier()

    st_logger.info("[DEBUG] %s", health_check(), exc_info=0)
    st_logger.info("[INFO] Initialize %s model successfully", "buttbread_resnet152_3.h5", exc_info=0)

    st.title(ui_cfg["title"])
    st.markdown(f'{ui_cfg["markdown"]["release"]} {ui_cfg["markdown"]["star"]} {ui_cfg["markdown"]["visitor"]}', unsafe_allow_html=True)

    mode = st.radio("", [ui_cfg["mode"]["upload"]["main_label"], ui_cfg["mode"]["select"]["main_label"]])

    if mode == ui_cfg["mode"]["upload"]["main_label"]:
        image_file = st.file_uploader(mode, accept_multiple_files=False)
    elif mode == ui_cfg["mode"]["select"]["main_label"]:
        class_label = st.selectbox(ui_cfg["mode"]["select"]["class_label"], model_cfg["label"].values())

        if class_label == model_cfg["label"]["corgi"]:
            image_label = st.selectbox(ui_cfg["mode"]["select"]["corgi_label"], [*image_cfg["corgi"]])
            image_file = os.path.join(image_cfg["base_path"], image_cfg["corgi"][image_label])
        elif class_label == model_cfg["label"]["bread"]:
            image_label = st.selectbox(ui_cfg["mode"]["select"]["bread_label"], [*image_cfg["bread"]])
            image_file = os.path.join(image_cfg["base_path"], image_cfg["bread"][image_label])

    if image_file:
        try:
            image = Image.open(image_file)

            if image.mode != "RGB":
                temporary_format = image.format
                image = image.convert("RGB")
                image.format = temporary_format

            if mode == ui_cfg["mode"]["upload"]["main_label"]:
                image.filename = image_file.name
            elif mode == ui_cfg["mode"]["select"]["main_label"]:
                image.filename = os.path.basename(image_file)

            prediction = classifier.predict(image)

            st_logger.info("[DEBUG] %s", health_check(), exc_info=0)
            st_logger.info("[INFO] Predict %s image successfully", image.filename, exc_info=0)

        except Exception as ex:
            st.error("ERROR: Unable to predict {} ({}) !!!".format(image_file.name, image_file.type))
            st_logger.error("[ERROR] Unable to predict %s (%s) !!!", image_file.name, image_file.type, exc_info=0)
            image_file, image, prediction = None, None, None

    if image is not None or prediction is not None:
        st.header("Here is the image you've chosen")
        resized_image = image.resize((400, 400))
        st.image(resized_image)
        st.write("Prediction:")
        st.json(prediction)