Spaces:

iBrokeTheCode
/

Image_Classifier_with_CNN

Sleeping

Image_Classifier_with_CNN / src /streamlit_app.py

fix: Clear session state after predicting images

a2c2878 about 1 year ago

6.25 kB

import os

import streamlit as st
from PIL import Image

from predictor import predict_image

APP_DIR = os.path.dirname(os.path.abspath(__file__))
ASSETS_DIR = os.path.join(APP_DIR, "assets")

# 📌 PAGE SETUP
st.set_page_config(page_title="Image Classifier App", page_icon="🤖", layout="centered")
st.html("""
<style>
    .stMainBlockContainer {
        max-width: 70rem;
    }
</style>
""")

# 📌 INITIALIZE SESSION STATE
if "uploaded_image" not in st.session_state:
    st.session_state["uploaded_image"] = None
if "example_selected" not in st.session_state:
    st.session_state["example_selected"] = False
if "prediction_result" not in st.session_state:
    st.session_state["prediction_result"] = None

# 📌 MAIN APP LAYOUT
with st.container():
    st.title(
        body="🖼️ Image Classifier with CNN",
        help="An interactive application to classify images into over 1000 categories.",
    )
    st.html("<br>")

    # Use tabs for different sections of the app
    tab_app, tab_description = st.tabs(["**App**", "**Description**"])

    # 📌 APP TAB
    with tab_app:
        # Create a two-column layout for the app interface
        col_upload, col_results = st.columns(2, gap="large")

        # 📌 IMAGE UPLOAD & EXAMPLE SELECTION
        with col_upload:
            st.header("Upload an Image", divider=True)

            # File uploader widget
            uploaded_file = st.file_uploader(
                label="Drag and drop an image here or click to browse",
                type=["jpg", "jpeg", "png"],
                help="Maximum file size is 200MB",
                key="image_uploader",
            )

            st.html("<br>")
            st.subheader("Or Try an Example", divider=True)

            # Segmented control for selecting example images
            selected_example = st.segmented_control(
                label="Categories",
                options=["Animal", "Vehicle", "Object", "Building"],
                default=None,
                help="Select one of the pre-loaded examples",
            )

            st.html("<br>")

            # --- THE SINGLE CLASSIFY BUTTON ---
            classify_button = st.button(
                label="Classify Image",
                key="classify_btn",
                type="primary",
                icon="✨",
            )

        # --- LOGIC FOR IMAGE SELECTION & PREDICTION ---
        # Clear the previous prediction result if a new input is selected
        if uploaded_file or selected_example:
            st.session_state.prediction_result = None

        image_to_process = None

        if uploaded_file:
            image_to_process = Image.open(uploaded_file)

        elif selected_example:
            try:
                img_path = os.path.join(
                    APP_DIR, "assets", f"{selected_example.lower()}.jpg"
                )
                image_to_process = Image.open(img_path)
            except FileNotFoundError:
                st.error(
                    f"Error: The example image '{selected_example.lower()}.jpg' was not found."
                )
                st.stop()

        # 📌 PREDICTION RESULTS
        with col_results:
            st.header("Results", divider=True)

            # Display a "get started" message if no image is selected
            if not image_to_process and not st.session_state.prediction_result:
                st.info("Choose an image or an example to get a prediction.")

            # Display the image if one is selected
            if image_to_process:
                st.image(image_to_process, caption="Image to be classified")

            # If the button is clicked, run the prediction logic
            if classify_button and image_to_process:
                with st.spinner("Analyzing image..."):
                    try:
                        from predictor import predict_image

                        predicted_label, predicted_score = predict_image(
                            image_to_process
                        )
                        st.session_state.prediction_result = {
                            "label": predicted_label.replace("_", " ").title(),
                            "score": predicted_score,
                        }
                    except Exception as e:
                        st.error(f"An error occurred during prediction: {e}")

            # Display the prediction result if available in session state
            if st.session_state.prediction_result:
                st.metric(
                    label="Prediction",
                    value=st.session_state.prediction_result["label"],
                    delta=f"{st.session_state.prediction_result['score'] * 100:.2f}%",
                    help="The predicted category and its confidence score.",
                    delta_color="normal",
                )
                st.balloons()

            elif image_to_process:
                st.info("Click 'Classify Image' to see the prediction.")


# 📌 DESCRIPTION TAB
with tab_description:
    st.header("About This Project", divider=True)
    st.markdown(
        """
        This project showcases a Convolutional Neural Network (CNN) model that automatically
        classifies images into over 1000 different categories.
        ### Original Architecture
        The original project was built as a multi-service architecture, featuring:
        * **Streamlit:** For the web user interface.
        * **FastAPI:** As a RESTful API to handle image processing and model serving.
        * **Redis:** A message broker for communication between the services.
        ### Portfolio Adaptation
        For a live and cost-effective demo, this application has been adapted into a single-service
        solution. The core logic of the FastAPI backend has been integrated directly into
        the Streamlit app. This demonstrates the ability to adapt a solution for
        specific deployment and resource constraints.
        ### Technologies Used
        * **Streamlit:** For the interactive web interface.
        * **TensorFlow:** For loading and running the pre-trained CNN model.
        * **Pre-trained Model:** ResNet50 with weights trained on the ImageNet dataset.
        """
    )

	import os

	import streamlit as st
	from PIL import Image

	from predictor import predict_image

	APP_DIR = os.path.dirname(os.path.abspath(__file__))
	ASSETS_DIR = os.path.join(APP_DIR, "assets")

	# 📌 PAGE SETUP
	st.set_page_config(page_title="Image Classifier App", page_icon="🤖", layout="centered")
	st.html("""
	<style>
	.stMainBlockContainer {
	max-width: 70rem;
	}
	</style>
	""")

	# 📌 INITIALIZE SESSION STATE
	if "uploaded_image" not in st.session_state:
	st.session_state["uploaded_image"] = None
	if "example_selected" not in st.session_state:
	st.session_state["example_selected"] = False
	if "prediction_result" not in st.session_state:
	st.session_state["prediction_result"] = None

	# 📌 MAIN APP LAYOUT
	with st.container():
	st.title(
	body="🖼️ Image Classifier with CNN",
	help="An interactive application to classify images into over 1000 categories.",
	)
	st.html("<br>")

	# Use tabs for different sections of the app
	tab_app, tab_description = st.tabs(["App", "Description"])

	# 📌 APP TAB
	with tab_app:
	# Create a two-column layout for the app interface
	col_upload, col_results = st.columns(2, gap="large")

	# 📌 IMAGE UPLOAD & EXAMPLE SELECTION
	with col_upload:
	st.header("Upload an Image", divider=True)

	# File uploader widget
	uploaded_file = st.file_uploader(
	label="Drag and drop an image here or click to browse",
	type=["jpg", "jpeg", "png"],
	help="Maximum file size is 200MB",
	key="image_uploader",
	)

	st.html("<br>")
	st.subheader("Or Try an Example", divider=True)

	# Segmented control for selecting example images
	selected_example = st.segmented_control(
	label="Categories",
	options=["Animal", "Vehicle", "Object", "Building"],
	default=None,
	help="Select one of the pre-loaded examples",
	)

	st.html("<br>")

	# --- THE SINGLE CLASSIFY BUTTON ---
	classify_button = st.button(
	label="Classify Image",
	key="classify_btn",
	type="primary",
	icon="✨",
	)

	# --- LOGIC FOR IMAGE SELECTION & PREDICTION ---
	# Clear the previous prediction result if a new input is selected
	if uploaded_file or selected_example:
	st.session_state.prediction_result = None

	image_to_process = None

	if uploaded_file:
	image_to_process = Image.open(uploaded_file)

	elif selected_example:
	try:
	img_path = os.path.join(
	APP_DIR, "assets", f"{selected_example.lower()}.jpg"
	)
	image_to_process = Image.open(img_path)
	except FileNotFoundError:
	st.error(
	f"Error: The example image '{selected_example.lower()}.jpg' was not found."
	)
	st.stop()

	# 📌 PREDICTION RESULTS
	with col_results:
	st.header("Results", divider=True)

	# Display a "get started" message if no image is selected
	if not image_to_process and not st.session_state.prediction_result:
	st.info("Choose an image or an example to get a prediction.")

	# Display the image if one is selected
	if image_to_process:
	st.image(image_to_process, caption="Image to be classified")

	# If the button is clicked, run the prediction logic
	if classify_button and image_to_process:
	with st.spinner("Analyzing image..."):
	try:
	from predictor import predict_image

	predicted_label, predicted_score = predict_image(
	image_to_process
	)
	st.session_state.prediction_result = {
	"label": predicted_label.replace("_", " ").title(),
	"score": predicted_score,
	}
	except Exception as e:
	st.error(f"An error occurred during prediction: {e}")

	# Display the prediction result if available in session state
	if st.session_state.prediction_result:
	st.metric(
	label="Prediction",
	value=st.session_state.prediction_result["label"],
	delta=f"{st.session_state.prediction_result['score'] * 100:.2f}%",
	help="The predicted category and its confidence score.",
	delta_color="normal",
	)
	st.balloons()

	elif image_to_process:
	st.info("Click 'Classify Image' to see the prediction.")


	# 📌 DESCRIPTION TAB
	with tab_description:
	st.header("About This Project", divider=True)
	st.markdown(
	"""
	This project showcases a Convolutional Neural Network (CNN) model that automatically
	classifies images into over 1000 different categories.

	### Original Architecture
	The original project was built as a multi-service architecture, featuring:
	* Streamlit: For the web user interface.
	* FastAPI: As a RESTful API to handle image processing and model serving.
	* Redis: A message broker for communication between the services.

	### Portfolio Adaptation
	For a live and cost-effective demo, this application has been adapted into a single-service
	solution. The core logic of the FastAPI backend has been integrated directly into
	the Streamlit app. This demonstrates the ability to adapt a solution for
	specific deployment and resource constraints.

	### Technologies Used
	* Streamlit: For the interactive web interface.
	* TensorFlow: For loading and running the pre-trained CNN model.
	* Pre-trained Model: ResNet50 with weights trained on the ImageNet dataset.
	"""
	)