app.py

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()