Spaces:
Sleeping
Sleeping
File size: 5,016 Bytes
29bcdf2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
from PIL import Image
import io
# Set page config
st.set_page_config(
page_title="Stone Classification",
page_icon="🪨",
layout="wide"
)
# Custom CSS to improve the appearance
st.markdown("""
<style>
.main {
padding: 2rem;
}
.stButton>button {
width: 100%;
margin-top: 1rem;
}
.upload-text {
text-align: center;
padding: 2rem;
}
</style>
""", unsafe_allow_html=True)
@st.cache_resource
def load_model():
"""Load the trained model"""
return tf.keras.models.load_model('custom_model.h5')
def preprocess_image(image):
"""Preprocess the uploaded image"""
# Convert to RGB if needed
if image.mode != 'RGB':
image = image.convert('RGB')
# Convert to numpy array
img_array = np.array(image)
# Convert to RGB if needed
if len(img_array.shape) == 2: # Grayscale
img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
elif img_array.shape[2] == 4: # RGBA
img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
# Preprocess image similar to training
img_hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV)
img_hsv[:, :, 2] = cv2.equalizeHist(img_hsv[:, :, 2])
img_array = cv2.cvtColor(img_hsv, cv2.COLOR_HSV2RGB)
# Adjust brightness
target_brightness = 150
current_brightness = np.mean(img_array)
alpha = target_brightness / (current_brightness + 1e-5)
img_array = cv2.convertScaleAbs(img_array, alpha=alpha, beta=0)
# Apply Gaussian blur
img_array = cv2.GaussianBlur(img_array, (5, 5), 0)
# Resize
img_array = cv2.resize(img_array, (256, 256))
# Normalize
img_array = img_array.astype('float32') / 255.0
return img_array
def main():
# Title
st.title("🪨 Stone Classification")
st.write("Upload an image of a stone to classify its type")
# Initialize session state for prediction if not exists
if 'prediction' not in st.session_state:
st.session_state.prediction = None
if 'confidence' not in st.session_state:
st.session_state.confidence = None
# Create two columns
col1, col2 = st.columns(2)
with col1:
st.subheader("Upload Image")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Display uploaded image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_column_width=True)
# Add predict button
if st.button("Predict"):
try:
# Load model
model = load_model()
# Preprocess image
processed_image = preprocess_image(image)
# Make prediction
prediction = model.predict(np.expand_dims(processed_image, axis=0))
class_names = ['Artificial', 'Nature'] # Replace with your actual class names
# Get prediction and confidence
predicted_class = class_names[np.argmax(prediction)]
confidence = float(np.max(prediction)) * 100
# Store in session state
st.session_state.prediction = predicted_class
st.session_state.confidence = confidence
except Exception as e:
st.error(f"Error during prediction: {str(e)}")
with col2:
st.subheader("Prediction Results")
if st.session_state.prediction is not None:
# Create a card-like container for results
results_container = st.container()
with results_container:
st.markdown("""
<style>
.prediction-card {
padding: 2rem;
border-radius: 0.5rem;
background-color: #f0f2f6;
margin: 1rem 0;
}
</style>
""", unsafe_allow_html=True)
st.markdown("<div class='prediction-card'>", unsafe_allow_html=True)
st.markdown(f"### Predicted Class: {st.session_state.prediction}")
st.markdown(f"### Confidence: {st.session_state.confidence:.2f}%")
st.markdown("</div>", unsafe_allow_html=True)
# Add confidence bar
st.progress(st.session_state.confidence / 100)
else:
st.info("Upload an image and click 'Predict' to see the results")
# Footer
st.markdown("---")
st.markdown("Made with ❤️ using Streamlit")
if __name__ == "__main__":
main() |