app.py

import streamlit as tf_app
import tensorflow as tf
from PIL import Image
import numpy as np
import os

# Set up page title and layout
tf_app.set_page_config(page_title="CNN Image Classifier", layout="centered")
tf_app.title("🚀 CNN Image Classification App")
tf_app.write("Upload an image, and our trained Convolutional Neural Network will classify it!")

# Load your trained model safely
@tf_app.cache_resource
def load_my_model():
    # Force Python to look in the exact folder where app.py lives
    current_dir = os.path.dirname(os.path.abspath(__file__))
    model_path = os.path.join(current_dir, 'my_cifar10_model.keras')
    return tf.keras.models.load_model(model_path)

try:
    model = load_my_model()
    tf_app.success("Model loaded successfully!")
except Exception as e:
    tf_app.error(f"Could not find 'my_cifar10_model.keras'. Make sure it's in this exact folder! Error: {e}")

# CIFAR-10 class labels
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

# Create the drag-and-drop file uploader widget
uploaded_file = tf_app.file_uploader("Choose an image...", type=["jpg", "jpeg", "png", "webp"])

if uploaded_file is not None:
    # Display the uploaded image on the web page
    img = Image.open(uploaded_file)
    tf_app.image(img, caption='Your Uploaded Image', use_column_width=True)
    
    tf_app.write("🧠 Model is analyzing...")
    
    # Preprocess the image exactly like we did in training
    img_resized = img.resize((32, 32))
    img_array = np.array(img_resized) / 255.0
    
    # Handle grayscale or alpha channel PNG images safely
    if len(img_array.shape) == 2:
        img_array = np.stack((img_array,)*3, axis=-1)
    elif img_array.shape[2] == 4:
        img_array = img_array[:, :, :3]
        
    img_array = np.expand_dims(img_array, axis=0)
    
    # Run prediction
    predictions = model.predict(img_array)
    score = tf.nn.softmax(predictions[0])
    predicted_class = class_names[np.argmax(score)]
    confidence = 100 * np.max(score)
    
    # Output the prediction and confidence bar beautifully
    tf_app.markdown(f"### 🤖 Prediction: **{predicted_class.upper()}**")
    tf_app.progress(int(confidence))
    tf_app.write(f"📊 Confidence Score: **{confidence:.2f}%**")