langatnaomi10's picture
Update app.py (#2)
937575c verified
import os
import streamlit as st
import requests
import numpy as np
from PIL import Image
import pickle # Using pickle since the model is saved as a .pkl file
# Define the absolute path for the model
MODEL_PATH = "/app/trained_model.pkl"
# Ensure the model has the correct read permissions
if os.path.exists(MODEL_PATH):
os.chmod(MODEL_PATH, 0o644)
# Cache model loading to avoid repeated downloads
@st.cache_resource
def load_model():
# Load the trained model from the saved .pkl file
with open(MODEL_PATH, "rb") as file:
model = pickle.load(file)
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 the 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}")