app.py

"""
Streamlit app
"""
import sys

import streamlit as st

from constants import (CLAS_FILEPATH, CLAS_THRESHOLD, CLASSES, DET_FILEPATH,
                       DET_NAME, DET_THRESHOLD, DEVICE, OUTPUT_IMG_FILEPATH)

sys.path.append("./efficientdet")

from PIL import Image

from efficientdet.efficientdet import plot_results
from trash_detector import detect_trash


def initial_config():
    """
    Initial configuration of streamlit page
    """
    st.set_page_config(
        page_title="Waste Classifier",
        page_icon="♻️",
    )


def render():
    """
    Render the streamlit app
    """
    st.title("Waste classifier")
    st.markdown("""Classify your waste into different classes""")

    # Image loader and button
    uploaded_file = st.file_uploader(
        "Upload image with trash", type=["jpg", "jpeg", "png", "gif", "bmp"]
    )
    classify_button = st.button("Classify trash")

    if classify_button:
        if not uploaded_file:
            st.error("Upload an image")
        else:
            # Create two columns
            col1, col2 = st.columns(2)

            # Column 1: Uploaded image
            with col1:
                st.write("Uploaded image")
                st.image(
                    uploaded_file, caption="Uploaded Image.", use_column_width=True
                )

            # Column 2: Classified image
            with col2:
                with st.spinner(text="Classifying the trash..."):
                    img = Image.open(uploaded_file).convert("RGB")
                    cls_prob, bboxes_final = detect_trash(
                        img,
                        DET_NAME,
                        DET_FILEPATH,
                        CLAS_FILEPATH,
                        DEVICE,
                        DET_THRESHOLD,
                        CLAS_THRESHOLD,
                    )
                    # plot and save demo image
                    plot_results(
                        img, cls_prob, bboxes_final, CLASSES, OUTPUT_IMG_FILEPATH
                    )
                    output_img = Image.open(OUTPUT_IMG_FILEPATH)
                    st.write("Classified image")
                    st.image(
                        output_img, caption="Classified Image.", use_column_width=True
                    )


if __name__ == "__main__":
    initial_config()
    render()