crack-detection.py

"""
🔍 Crack Detection System - Streamlit App
FINAL FIXED VERSION - Preprocessing matches training EXACTLY
"""

import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image

# ============================================
# PAGE CONFIG
# ============================================
st.set_page_config(
    page_title="Crack Detection System",
    page_icon="🔍",
    layout="wide"
)

# ============================================
# CSS
# ============================================
st.markdown("""
<style>
    .result-crack {
        background-color: #FFE5E5;
        border: 3px solid #FF4B4B;
        border-radius: 15px;
        padding: 2rem;
        text-align: center;
    }
    .result-no-crack {
        background-color: #E5FFE5;
        border: 3px solid #4CAF50;
        border-radius: 15px;
        padding: 2rem;
        text-align: center;
    }
</style>
""", unsafe_allow_html=True)

# ============================================
# BUILD MODEL - MUST MATCH TRAINING EXACTLY
# ============================================
def build_model():
    """
    Build EXACT same architecture as training.
    
    IMPORTANT: Training model included preprocess_input as a LAYER,
    so the model expects RAW pixels (0-255), not rescaled!
    
    But training used ImageDataGenerator with rescale=1./255,
    which means during training the model received pixels in [0,1] range,
    and then preprocess_input converted them further.
    
    This is the conflict we need to resolve.
    """
    from tensorflow.keras.applications import ResNet50
    from tensorflow.keras import layers, Model
    
    base_model = ResNet50(
        weights='imagenet',
        include_top=False,
        input_shape=(224, 224, 3)
    )
    
    # EXACT architecture from training
    inputs = tf.keras.Input(shape=(224, 224, 3))
    x = tf.keras.applications.resnet50.preprocess_input(inputs)
    x = base_model(x, training=False)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Dense(512, activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.3)(x)
    x = layers.Dense(256, activation='relu')(x)
    x = layers.Dropout(0.2)(x)
    outputs = layers.Dense(1, activation='sigmoid', dtype='float32')(x)
    
    model = Model(inputs, outputs)
    return model

# ============================================
# LOAD MODEL
# ============================================
@st.cache_resource
def load_model():
    """Load model weights"""
    import os
    
    # Try weight files in order
    for wf in ['best.weights.h5', 'final_weights.weights.h5']:
        if os.path.exists(wf):
            try:
                model = build_model()
                model.load_weights(wf)
                return model, wf
            except Exception as e:
                st.warning(f"Failed {wf}: {e}")
    
    # Try full model
    if os.path.exists('crack_model.h5'):
        try:
            model = tf.keras.models.load_model('crack_model.h5', compile=False)
            return model, 'crack_model.h5'
        except:
            pass
    
    return None, None

# ============================================
# PREPROCESSING - CRITICAL FIX!
# ============================================
def preprocess_image(image):
    """
    CRITICAL: Match EXACT preprocessing from training!
    
    Training pipeline:
    1. ImageDataGenerator with rescale=1./255 → pixels become [0, 1]
    2. Model has preprocess_input layer → expects [0, 255], outputs [-1, 1]
    
    This is a CONFLICT in the training code!
    
    The training fed [0,1] pixels to preprocess_input which expects [0,255].
    So preprocess_input received already-scaled values and scaled them again.
    
    To match this EXACT behavior in inference:
    - We need to rescale to [0,1] first (like ImageDataGenerator did)
    - Then feed to model (which applies preprocess_input internally)
    """
    # Convert to RGB
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Resize to 224x224
    image = image.resize((224, 224), Image.Resampling.LANCZOS)
    
    # Convert to array
    img_array = np.array(image, dtype=np.float32)
    
    # MATCH TRAINING: Apply rescale=1./255 like ImageDataGenerator did
    img_array = img_array / 255.0
    
    # Add batch dimension
    img_array = np.expand_dims(img_array, axis=0)
    
    # Model will apply preprocess_input internally
    return img_array

# ============================================
# PREDICTION
# ============================================
def predict_crack(model, image):
    """
    Make prediction.
    
    Training class indices (alphabetical from folder names):
    - 'negative' = 0 (no crack)
    - 'positive' = 1 (crack)
    
    Model output sigmoid:
    - Close to 0 → class 0 → negative → NO CRACK
    - Close to 1 → class 1 → positive → CRACK
    """
    img = preprocess_image(image)
    
    # Predict
    pred = model.predict(img, verbose=0)[0][0]
    pred = float(pred)
    
    # Class mapping
    if pred > 0.5:
        label = "🔴 CRACK DETECTED"
        confidence = pred * 100
        is_crack = True
    else:
        label = "🟢 NO CRACK"
        confidence = (1.0 - pred) * 100
        is_crack = False
    
    return label, confidence, is_crack, pred

# ============================================
# MAIN APP
# ============================================
def main():
    st.markdown("<h1 style='text-align:center;color:#FF4B4B'>🔍 Crack Detection System</h1>", unsafe_allow_html=True)
    st.markdown("<p style='text-align:center;color:#666'>Upload an image to detect cracks using AI (ResNet50)</p>", unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.header("📊 Model Info")
        st.markdown("""
        **Model:** ResNet50  
        **Dataset:** 40,000 images  
        **Accuracy:** ~98%  
        **AUC:** ~99.9%
        """)
        
        st.header("📁 Classes")
        st.markdown("""
        - **negative** → No Crack (0)
        - **positive** → Crack (1)
        """)
    
    # Load model
    model, source = load_model()
    
    if model is None:
        st.error("❌ Model not found!")
        st.stop()
    
    st.success(f"✅ Model loaded: {source}")
    
    # Show performance in sidebar
    with st.sidebar:
        st.header("📈 Performance")
        try:
            st.image("confusion_matrix.png", caption="Confusion Matrix")
        except:
            pass
        try:
            st.image("roc.png", caption="ROC Curve")
        except:
            pass
    
    # File uploader
    st.markdown("---")
    uploaded = st.file_uploader("📤 Upload Image", type=['jpg', 'jpeg', 'png', 'bmp'])
    
    if uploaded:
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("📷 Uploaded Image")
            image = Image.open(uploaded)
            st.image(image, use_column_width=True)
            st.caption(f"Size: {image.size[0]}x{image.size[1]}")
        
        with col2:
            st.subheader("🤖 Prediction")
            
            with st.spinner("Analyzing..."):
                label, conf, is_crack, raw = predict_crack(model, image)
            
            # Debug info
            st.sidebar.markdown("---")
            st.sidebar.markdown("### 🔧 Debug")
            st.sidebar.write(f"Raw output: **{raw:.6f}**")
            st.sidebar.write(f"Threshold: 0.5")
            st.sidebar.write(f"Result: {'CRACK' if raw > 0.5 else 'NO CRACK'}")
            
            # Display result
            if is_crack:
                st.markdown(f"""
                <div class="result-crack">
                    <h1 style="color:#FF4B4B;margin:0">{label}</h1>
                    <h2>Confidence: {conf:.1f}%</h2>
                    <p>⚠️ Crack detected!</p>
                </div>
                """, unsafe_allow_html=True)
            else:
                st.markdown(f"""
                <div class="result-no-crack">
                    <h1 style="color:#4CAF50;margin:0">{label}</h1>
                    <h2>Confidence: {conf:.1f}%</h2>
                    <p>✅ No crack detected!</p>
                </div>
                """, unsafe_allow_html=True)
            
            st.progress(int(min(conf, 100)))
            
            c1, c2 = st.columns(2)
            c1.metric("Result", "Crack" if is_crack else "No Crack")
            c2.metric("Confidence", f"{conf:.1f}%")
    else:
        st.info("👆 Upload an image to analyze")
        try:
            st.image("predictions.png", caption="Sample Predictions from Training")
        except:
            pass
    
    st.markdown("---")
    st.caption("🔍 Crack Detection | ResNet50 | 40k images | ~98% accuracy")

if __name__ == "__main__":
    main()