import streamlit as st
from PIL import Image
import tensorflow as tf
import numpy as np
import cv2
from skimage import img_as_float
from skimage.segmentation import chan_vese
from tensorflow.keras.applications.efficientnet import preprocess_input



model = tf.keras.models.load_model('efficientnetb2_model_v2.h5')

# Define class labels
CLASSES = ['Basal cell carcinoma', 'Benign keratosis', 'Melanoma']

def inpainted_image(image):
    gray_image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)
    median_filtered_image = cv2.medianBlur(gray_image, 5)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    bottom_hat = cv2.morphologyEx(median_filtered_image, cv2.MORPH_BLACKHAT, kernel)
    _, hair_mask = cv2.threshold(bottom_hat, 10, 255, cv2.THRESH_BINARY)
    inpainted_image = cv2.inpaint(median_filtered_image, hair_mask, 3, cv2.INPAINT_TELEA)
    return inpainted_image

def gac_segmentation(image):
    image = img_as_float(image)
    segmented_image = chan_vese(image, mu=0.25, lambda1=1, lambda2=1, tol=1e-3, max_num_iter=40, 
                                         dt=0.5, init_level_set="checkerboard", extended_output=True)
    return segmented_image[0]


def preprocess_image(image):
    image = image.resize((260, 260))
    image = preprocess_input(np.array(image))  
    return image

# Function to make predictions
def predict(image):
    image = preprocess_image(image)
    image = np.expand_dims(image, axis=0) 
    prediction = model.predict(image) 
    predicted_label_index = np.argmax(prediction)
    predicted_label = CLASSES[predicted_label_index]
    return predicted_label, prediction

def main():
    st.title('Skin Lesion Classifier')
    st.write('Upload an image of a skin lesion for classification.')

    uploaded_image = st.file_uploader('Choose an image...', type=['jpg', 'png', 'jpeg'])

    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        if st.button('Classify'):
            with st.spinner('Processing and Classifying...'):
                # Preprocess the image
                inpainted_img = inpainted_image(image)
                
                # Apply GAC segmentation
                segmented_img = gac_segmentation(inpainted_img)
                
                # Convert segmented image to a format suitable for prediction
                segmented_img = (segmented_img * 255).astype(np.uint8) 
                segmented_img_pil = Image.fromarray(segmented_img).convert("RGB").resize((260, 260))
                
                # Predict
                predicted_label, prediction = predict(segmented_img_pil)
                
                # Display results
                st.image(segmented_img, caption='Preprocessed and Segmented Image', use_column_width=True)
                st.write('Predicted Label:', predicted_label)
                st.write('Prediction Scores:', prediction)

if __name__ == '__main__':
    main()