app.py

import streamlit as st
import requests
import numpy as np
from PIL import Image
from tensorflow import keras

# Cache the model download and loading to avoid reloading on every interaction
@st.cache_resource
def load_model():
    # URL of the model file on Hugging Face
    model_url = "https://huggingface.co/spaces/langatnaomi10/CNN/resolve/main/trained_model.pkl"
    
    # Download the file
    model_path = "trained_model.pkl"
    response = requests.get(model_url, stream=True)
    with open(model_path, "wb") as file:
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                file.write(chunk)
    
    # Load the saved model
    model = keras.models.load_model(model_path)
    return model

# Load the model
try:
    model = load_model()
except Exception as e:
    st.error(f"Failed to load the model: {e}")
    st.stop()

# Define a prediction function
def predict_axle_configuration(image):
    # Resize and preprocess the image
    image = image.resize((128, 128))  # Resize to match model input size
    image_array = np.array(image) / 255.0  # Normalize pixel values to [0, 1]
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    
    # Make prediction
    prediction = model.predict(image_array)
    return prediction

# Streamlit UI
st.title("Vehicle Axle Configuration Prediction")
uploaded_file = st.file_uploader("Upload a vehicle image", type=['jpg', 'jpeg', 'png'])

if uploaded_file:
    try:
        img = Image.open(uploaded_file)
        st.image(img, caption='Uploaded Image', use_column_width=True)
        st.write("Classifying...")
        
        # Get prediction
        result = predict_axle_configuration(img)
        
        # Display prediction (assuming result is a probability or class index)
        st.write(f"Predicted Axle Configuration: {result}")
    except Exception as e:
        st.error(f"An error occurred during prediction: {e}")