trohith89's picture
Update app.py
88f40ec verified
raw
history blame
2.25 kB
import streamlit as st
import cv2
import numpy as np
from tensorflow.keras.models import load_model
import pickle
from PIL import Image
import os
# Load the model and label encoder
@st.cache_resource
def load_resources():
# Custom loading to handle compatibility
try:
model = load_model('captains_cv2_model.keras', compile=False) # Load without compiling first
except Exception as e:
st.error(f"Model loading failed: {str(e)}")
raise
with open('label_encoder.pkl', 'rb') as file:
le = pickle.load(file)
return model, le
# Preprocess the image
def preprocess_image(image_path):
img1 = cv2.imread(image_path)
img1 = cv2.resize(img1, (64, 64)) # Resize to 64x64
img1 = np.asarray(img1) # Shape: (64, 64, 3)
img1 = img1[np.newaxis, :, :, :] # Shape: (1, 64, 64, 3)
return img1
# Main app
def main():
model, le = load_resources()
st.title("Image Classification App")
st.write("Upload an image to get a prediction")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption='Uploaded Image', use_column_width=True)
file_extension = os.path.splitext(uploaded_file.name)[1].lower()
temp_filename = f"temp_image{file_extension}"
with open(temp_filename, "wb") as f:
f.write(uploaded_file.getvalue())
try:
processed_img = preprocess_image(temp_filename)
st.write(f"Processed image shape: {processed_img.shape}")
prediction = model.predict(processed_img)
predicted_class = le.inverse_transform([np.argmax(prediction)])
st.write("Prediction:", predicted_class[0])
st.write("Prediction Probabilities:")
for class_name, prob in zip(le.classes_, prediction[0]):
st.write(f"{class_name}: {prob:.4f}")
except Exception as e:
st.error(f"An error occurred: {str(e)}")
if os.path.exists(temp_filename):
os.remove(temp_filename)
if __name__ == '__main__':
main()