app.py

import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
from PIL import Image
from tensorflow.keras import layers, models
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.applications.efficientnet import preprocess_input
from sklearn.preprocessing import StandardScaler
import io

# Set page config
st.set_page_config(
    page_title="Stone Classification",
    page_icon="🪨",
    layout="wide"
)

# Custom CSS with improved styling
st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .stButton>button {
        width: 100%;
        margin-top: 1rem;
    }
    .prediction-card {
        padding: 2rem;
        border-radius: 0.5rem;
        background-color: #f0f2f6;
        margin: 1rem 0;
    }
    .top-predictions {
        margin-top: 2rem;
        padding: 1rem;
        background-color: white;
        border-radius: 0.5rem;
        box-shadow: 0 1px 3px rgba(0,0,0,0.12);
    }
    </style>
    """, unsafe_allow_html=True)

@st.cache_resource
def load_model():
    """Load the trained model"""
    try:
        return tf.keras.models.load_model('mlp_model.h5')
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None

def color_histogram(image, bins=16):
    """Calculate color histogram features"""
    hist_r = cv2.calcHist([image], [0], None, [bins], [0, 256]).flatten()
    hist_g = cv2.calcHist([image], [1], None, [bins], [0, 256]).flatten()
    hist_b = cv2.calcHist([image], [2], None, [bins], [0, 256]).flatten()
    
    # Normalize histograms
    hist_r = hist_r / (np.sum(hist_r) + 1e-7)
    hist_g = hist_g / (np.sum(hist_g) + 1e-7)
    hist_b = hist_b / (np.sum(hist_b) + 1e-7)
    
    return np.concatenate([hist_r, hist_g, hist_b])

def color_moments(image):
    """Calculate color moments features"""
    img = image.astype(np.float32) / 255.0
    moments = []
    
    for i in range(3):
        channel = img[:,:,i]
        mean = np.mean(channel)
        std = np.std(channel) + 1e-7  # Avoid division by zero
        skewness = np.mean(((channel - mean) / std) ** 3) if std != 0 else 0
        moments.extend([mean, std, skewness])
    
    return np.array(moments)

def dominant_color_descriptor(image, k=3):
    """Calculate dominant color descriptor"""
    pixels = image.reshape(-1, 3).astype(np.float32)
    
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
    flags = cv2.KMEANS_RANDOM_CENTERS
    
    try:
        _, labels, centers = cv2.kmeans(pixels, k, None, criteria, 10, flags)
        unique, counts = np.unique(labels, return_counts=True)
        percentages = counts / len(labels)
        return np.concatenate([centers.flatten(), percentages])
    except Exception:
        return np.zeros(k * 4)  # Return zero vector if clustering fails

def color_coherence_vector(image, k=3):
    """Calculate color coherence vector"""
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    gray = np.uint8(gray)
    
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    num_labels, labels = cv2.connectedComponents(binary)
    
    ccv = []
    for i in range(1, min(k+1, num_labels)):
        region_mask = (labels == i)
        total_pixels = np.sum(region_mask)
        ccv.extend([total_pixels, total_pixels])
    
    # Pad with zeros if needed
    ccv.extend([0] * (2 * k - len(ccv)))
    return np.array(ccv[:2*k])

@st.cache_resource
def create_feature_extractor():
    """Create and cache the feature extractor model"""
    input_shape = (256, 256, 3)
    inputs = layers.Input(shape=input_shape)
    x = layers.Lambda(preprocess_input)(inputs)
    
    base_model = EfficientNetB0(
        include_top=False,
        weights='imagenet',
        input_tensor=x
    )
    
    x = layers.GlobalAveragePooling2D()(base_model.output)
    return models.Model(inputs=inputs, outputs=x)

def extract_features(image):
    """Extract all features from an image"""
    # Convert image to uint8 for OpenCV operations
    image_uint8 = (image * 255).astype(np.uint8)
    
    # Extract traditional features
    hist_features = color_histogram(image_uint8)
    moment_features = color_moments(image_uint8)
    dominant_features = dominant_color_descriptor(image_uint8)
    ccv_features = color_coherence_vector(image_uint8)
    
    return np.concatenate([
        hist_features,
        moment_features,
        dominant_features,
        ccv_features
    ])

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 and resize
    img_array = np.array(image)
    img_array = cv2.resize(img_array, (256, 256))
    img_array = img_array.astype('float32') / 255.0
    
    # Extract traditional features
    traditional_features = extract_features(img_array)
    
    # Extract deep features
    feature_extractor = create_feature_extractor()
    deep_features = feature_extractor.predict(
        np.expand_dims(img_array, axis=0),
        verbose=0
    )
    
    # Combine features
    combined_features = np.concatenate([
        traditional_features.reshape(1, -1),
        deep_features.reshape(1, -1)
    ], axis=1)
    
    # Scale features
    scaler = StandardScaler()
    return scaler.fit_transform(combined_features)

def get_top_predictions(prediction, class_names, top_k=5):
    """Get top k predictions with their probabilities"""
    top_indices = prediction.argsort()[0][-top_k:][::-1]
    return [
        (class_names[i], float(prediction[0][i]) * 100)
        for i in top_indices
    ]

def main():
    st.title("🪨 Stone Classification")
    st.write("Upload an image of a stone to classify its type")
    
    # Initialize session state
    if 'predictions' not in st.session_state:
        st.session_state.predictions = None
    
    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:
            try:
                image = Image.open(uploaded_file)
                st.image(image, caption="Uploaded Image", use_column_width=True)
                
                with st.spinner('Analyzing image...'):
                    model = load_model()
                    if model is None:
                        st.error("Failed to load model")
                        return
                    
                    processed_image = preprocess_image(image)
                    prediction = model.predict(processed_image, verbose=0)
                    
                    class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
                    st.session_state.predictions = get_top_predictions(prediction, class_names)
                    
            except Exception as e:
                st.error(f"Error processing image: {str(e)}")
    
    with col2:
        st.subheader("Prediction Results")
        if st.session_state.predictions:
            # Display main prediction
            top_class, top_confidence = st.session_state.predictions[0]
            st.markdown(
                f"""
                <div class='prediction-card'>
                    <h3>Primary Prediction: Grade {top_class}</h3>
                    <h3>Confidence: {top_confidence:.2f}%</h3>
                </div>
                """,
                unsafe_allow_html=True
            )
            
            # Display confidence bar
            st.progress(top_confidence / 100)
            
            # Display top 5 predictions
            st.markdown("### Top 5 Predictions")
            st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)
            
            for class_name, confidence in st.session_state.predictions:
                cols = st.columns([2, 6, 2])
                with cols[0]:
                    st.write(f"Grade {class_name}")
                with cols[1]:
                    st.progress(confidence / 100)
                with cols[2]:
                    st.write(f"{confidence:.2f}%")
            
            st.markdown("</div>", unsafe_allow_html=True)
        else:
            st.info("Upload an image to see the predictions")
    
    st.markdown("---")
    st.markdown("Made with ❤️ using Streamlit")

if __name__ == "__main__":
    main()