Create app.py

app.py

@@ -0,0 +1,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()