app.py

import streamlit as st
import cv2
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image
import matplotlib.pyplot as plt

# Load the trained U-Net model (replace with the correct path to your model file)
@st.cache_resource
def load_unet_model():
    model = load_model('unet_crack_detection.h5')
    return model

# Function to process the image and make predictions
def process_and_predict(image, model):
    # Convert the image to grayscale
    img_gray = np.array(image.convert("L"))

    # Resize the image to the target size used in training (e.g., 256x256)
    target_size = (256, 256)
    img_resized = cv2.resize(img_gray, target_size)

    # Normalize the image
    img_resized = img_resized / 255.0
    img_resized = np.expand_dims(img_resized, axis=-1)  # Shape: (256, 256, 1)
    img_resized = np.expand_dims(img_resized, axis=0)   # Shape: (1, 256, 256, 1)

    # Predict the segmentation mask
    predicted_mask = model.predict(img_resized)
    predicted_mask = predicted_mask[0, :, :, 0]  # Remove batch and channel dimensions

    # Create overlay for the mask
    img_original_resized = cv2.resize(np.array(image), target_size)
    overlay = img_original_resized.copy()
    overlay[predicted_mask > 0.5] = [255, 0, 255]  # Color mask areas (e.g., red for mask regions)

    # Superimpose the overlay onto the original image with transparency
    alpha = 1
    superimposed_img = cv2.addWeighted(img_original_resized, 1 - alpha, overlay, alpha, 0)

    return img_original_resized, predicted_mask, superimposed_img

# Set up the Streamlit interface
st.title("Concrete Crack Detection App")
st.write("Upload an image, and the app will use a CNN model to predict and overlay the segmentation mask.")

# Load the model
model = load_unet_model()

# File uploader for the user to upload an image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Load the uploaded image
    image = Image.open(uploaded_file)
    
    # Display the original image
    st.subheader("Original Image")
    st.image(image)

    # Process the image and generate predictions
    img_original, predicted_mask, superimposed_img = process_and_predict(image, model)

    # Display the predicted mask
    #st.subheader("Predicted Mask")
    #fig, ax = plt.subplots()
    #ax.imshow(predicted_mask, cmap='gray')
    #ax.axis('off')
    #st.pyplot(fig)

    # Display the superimposed image
    st.subheader("Superimposed Image with Mask")
    superimposed_img_rgb = cv2.cvtColor(superimposed_img, cv2.COLOR_BGR2RGB)
    st.image(superimposed_img_rgb)