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}")