Spaces:
Build error
Build error
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 | |
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}") | |