import os
import cv2
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras.models import load_model
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D, concatenate
from tensorflow.keras.models import Model

IMG_WIDTH = 256
IMG_HEIGHT = 256
IMG_CHANNELS = 3

# Path to the model file
model_folder = 'model/'
model_file_keras = model_folder + 'model.keras'

# Function to define a U-Net model
def unet_modelorg(input_size=(IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS)):
    inputs = Input(input_size, name="input_image")
    #s = tf.keras.layers.Lambda(lambda x: x / 255)(inputs)

    #Contraction path
    c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(inputs)
    c1 = tf.keras.layers.Dropout(0.1)(c1)
    c1 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(c1)
    p1 = tf.keras.layers.MaxPooling2D((2, 2))(c1)

    c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(p1)
    c2 = tf.keras.layers.Dropout(0.1)(c2)
    c2 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(c2)
    p2 = tf.keras.layers.MaxPooling2D((2, 2))(c2)
 
    c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(p2)
    c3 = tf.keras.layers.Dropout(0.2)(c3)
    c3 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(c3)
    p3 = tf.keras.layers.MaxPooling2D((2, 2))(c3)
 
    c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(p3)
    c4 = tf.keras.layers.Dropout(0.2)(c4)
    c4 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(c4)
    p4 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(c4)
 
    c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', padding='same')(p4)
    c5 = tf.keras.layers.Dropout(0.3)(c5)
    c5 = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', padding='same')(c5)

    #Expansive path 
    u6 = tf.keras.layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c5)
    u6 = tf.keras.layers.concatenate([u6, c4])
    c6 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(u6)
    c6 = tf.keras.layers.Dropout(0.2)(c6)
    c6 = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(c6)
 
    u7 = tf.keras.layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c6)
    u7 = tf.keras.layers.concatenate([u7, c3])
    c7 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(u7)
    c7 = tf.keras.layers.Dropout(0.2)(c7)
    c7 = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(c7)
 
    u8 = tf.keras.layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c7)
    u8 = tf.keras.layers.concatenate([u8, c2])
    c8 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(u8)
    c8 = tf.keras.layers.Dropout(0.1)(c8)
    c8 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(c8)
 
    u9 = tf.keras.layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c8)
    u9 = tf.keras.layers.concatenate([u9, c1], axis=3)
    c9 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(u9)
    c9 = tf.keras.layers.Dropout(0.1)(c9)
    c9 = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(c9)
   
    outputs = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid',name="output_image_Final")(c9)
   
    model = Model(inputs=[inputs], outputs=[outputs])
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
   
    return model


def load_image(image_path, target_size=(IMG_HEIGHT,IMG_WIDTH)):
     # Load image
    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
    # Resize image
    image = cv2.resize(image, target_size)
    # Normalize image
    image = image.astype(np.float32)/255.0
    return image

def load_mask(image_path, target_size=(IMG_HEIGHT,IMG_WIDTH)):
     # Load image
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    # Resize image
    image = cv2.resize(image, target_size)
    # Normalize image
    image = image.astype(np.float32)/255.0
    # Add channel dimension
    image = np.expand_dims(image, axis=-1)
    return image

# Sample data generation for training (you should replace this with your actual data)
def generate_sample_data(num_samples=100, image_size=(IMG_HEIGHT,IMG_WIDTH)):
    TRAIN_PATHS = [
        'stage1_train/benign/Images/',
        'stage1_train/malignant/Images/',
        'stage1_train/normal/Images/'
    ]
    MASK_PATHS = [
        'stage1_train/benign/masks/',
        'stage1_train/malignant/masks/',
        'stage1_train/normal/masks/'
    ]
    
    img = []
    mask = []

    for i, train_path in enumerate(TRAIN_PATHS):
       mask_path = MASK_PATHS[i]
       
       for root, dirs, files in os.walk(train_path): 
           for file in files:
               if(file.endswith(".png")):
                   imagefile = file[:-4]; # remove extension .png
                   path = train_path + file
                   print( 'I : ' + path );
               
                   img.append(load_image(path))
                   for root, dirs, mfiles in os.walk(mask_path): 
                       for mfile in mfiles:
                           if(mfile.endswith(".png")):
                               if( -1 != mfile.find(imagefile)):
                                    pathmask = mask_path + mfile
                                    mask.append(load_mask(pathmask))
                                    print( 'M : ' + pathmask );
     
    return img, mask

def trainmodel():
    # Define and compile a new U-Net model
    model = unet_modelorg()
    print("New U-Net model created.")
    
    # Generate sample training data
    train_data, train_labels = generate_sample_data()

    train_data_Array = np.array(train_data)
    train_labels_Array = np.array(train_labels)
    # Train the model
    model.fit(train_data_Array, train_labels_Array, epochs=90, batch_size=16, validation_split=0.1)
    return model

def SaveModel(model):
    
    #Save weight file. It will be used by C++ module
    model.export(model_folder)
    #Save the trained model with weights to a file
    model.save(model_file_keras)
    #model.save_weights('model/model.weights.h5')

def predict(image, model, target_size=(IMG_HEIGHT,IMG_WIDTH)):   
    # Add batch dimension
    image = np.expand_dims(image, axis=0)
    
    # Make a prediction with the model
    prediction = model.predict(image)

    # Process the prediction as needed
    predicted_image = np.squeeze(prediction, axis=0)
    return predicted_image*255

def MakeSameSizeImage(image1, image2):
    newimage = cv2.resize(image2, (image1.shape[1], image1.shape[0]))
    return newimage
    
def ApplyOtsuThresholdAlogirhtm(image):
    rangechange = np.uint8(image)
    blurimage = cv2.GaussianBlur( rangechange, (5,5),0)
    ret, thresholdImage = cv2.threshold(blurimage, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return thresholdImage

def DrawContour(image,MaskImage):
    # Apply Otsu threshold algorithm to remove noices and make best threadshold image
    thresholdImage = ApplyOtsuThresholdAlogirhtm(MaskImage)
    contours, _ = cv2.findContours(thresholdImage, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
   
    for contour in contours:
        cv2.drawContours(image, [contour], -1, (0, 255, 0), 2)

def Initialize():
    # Check if the model file exists
    if os.path.exists(model_file_keras):
        # Load the model from the file
        model = load_model(model_file_keras)
        print("Model loaded from file.")
        return model
    else:
       #create a model and train it
       model = trainmodel()
       #Save the trained model with weights to a file
       SaveModel(model)
       return model

def ToSummary(image, title, column):
    plt.subplot(1,3,column)
    plt.title(title)
    plt.imshow(image,cmap='gray')

#----------------------------------------------------------------------#
#----------------------------------------------------------------------#
#----------------------------------------------------------------------#
# Main code
model = Initialize()

# Path to the PNG image file for prediction
image_path = input("Enter file to predict(path):")
image = cv2.imread(image_path, cv2.IMREAD_COLOR)

ToSummary(image,'original',1)

#Do AI predition upon the image
predicted_image = predict(load_image(image_path),model)

#Predicted image is in small size, resize to original image size
predicted_image = MakeSameSizeImage(image,predicted_image)

#Draw outter region of detected area
DrawContour(image,predicted_image)

# Save or display the predicted image
output_image_path = 'predicted_image.png'
cv2.imwrite(output_image_path, image)  # Rescale to original image range (0-255)

#add to summary report to show in a better way
ToSummary(predicted_image,'Predicted Mask',2)
ToSummary(image,'Predicted Area',3)

plt.show()


explain the code snippet