import streamlit as st
import cv2
import numpy as np
from PIL import Image

# Helper function to preprocess each image
def preprocess_image(image):
    # Convert to NumPy array and apply preprocessing
    image_np = np.array(image)
    
    # Convert to grayscale
    image_gray = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
    
    # CLAHE equalization
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    image_clahe = clahe.apply(image_gray)
    
    # Sobel X and Y gradients
    sobel_x = cv2.Sobel(image_clahe, cv2.CV_64F, 1, 0)
    sobel_y = cv2.Sobel(image_clahe, cv2.CV_64F, 0, 1)
    
    # Magnitude of gradients
    sobel_magnitude = cv2.magnitude(sobel_x, sobel_y)
    
    # Binarization
    _, binary_image = cv2.threshold(sobel_magnitude, 70, 100, cv2.THRESH_BINARY)
    
    # Erosion and dilation
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    dilated_image = cv2.dilate(binary_image, kernel, iterations=1)
    final_image = cv2.erode(dilated_image, kernel, iterations=1)
    
    # Normalize the final image to [0.0, 1.0] if float
    final_image_normalized = cv2.normalize(final_image, None, 0, 1, cv2.NORM_MINMAX)
    
    return final_image_normalized

# Streamlit App
st.title('Image Processing and Display')

# Upload image or images
uploaded_files = st.file_uploader("Choose images", type=['png', 'jpg', 'jpeg'], accept_multiple_files=True)

# Button to trigger preprocessing
if st.button('Process Images'):
    if uploaded_files:
        st.write(f"Total files uploaded: {len(uploaded_files)}")
        
        for uploaded_file in uploaded_files:
            # Load and display original image
            image = Image.open(uploaded_file)

            # Preprocess the image
            preprocessed_image = preprocess_image(image)
            
            # Display original and preprocessed images side by side
            col1, col2 = st.columns(2)

            with col1:
                st.subheader(f"Original Image: {uploaded_file.name}")
                st.image(image, caption='Original Image', use_column_width=True)

            with col2:
                st.subheader("Processed image")
                st.image(preprocessed_image, caption='Image with Blood vessels extracted', use_column_width=True, clamp=True)
    else:
        st.warning("Please upload images before processing.")