import tensorflow as tf
from keras import Input, Model
from keras.layers import Conv2D, MaxPooling2D, Conv2DTranspose, concatenate, BatchNormalization, Activation
from keras.layers import ConvLSTM2D, TimeDistributed, UpSampling2D

def convlstm_unet(input_size=(None, None, None, 3), timesteps=3):
    # 输入：时间序列的图像，形状为 (batch_size, timesteps, height, width, channels)
    inputs = Input(shape=input_size)
    
    # 下采样路径
    # 时间分布的卷积 LSTM 层
    c1 = ConvLSTM2D(filters=64, kernel_size=(3, 3), padding='same', return_sequences=True)(inputs)
    c1 = BatchNormalization()(c1)
    c1 = Activation('relu')(c1)
    p1 = TimeDistributed(MaxPooling2D(pool_size=(2, 2)))(c1)

    c2 = ConvLSTM2D(filters=128, kernel_size=(3, 3), padding='same', return_sequences=True)(p1)
    c2 = BatchNormalization()(c2)
    c2 = Activation('relu')(c2)
    p2 = TimeDistributed(MaxPooling2D(pool_size=(2, 2)))(c2)

    c3 = ConvLSTM2D(filters=256, kernel_size=(3, 3), padding='same', return_sequences=False)(p2)
    c3 = BatchNormalization()(c3)
    c3 = Activation('relu')(c3)

    # 上采样路径
    u1 = UpSampling2D(size=(2, 2))(c3)
    # 将 c2 的最后一个时间步与上采样结果进行拼接
    merge1 = concatenate([u1, c2[:, -1, :, :, :]], axis=-1)
    c4 = Conv2D(filters=128, kernel_size=(3, 3), padding='same')(merge1)
    c4 = BatchNormalization()(c4)
    c4 = Activation('relu')(c4)

    u2 = UpSampling2D(size=(2, 2))(c4)
    merge2 = concatenate([u2, c1[:, -1, :, :, :]], axis=-1)
    c5 = Conv2D(filters=64, kernel_size=(3, 3), padding='same')(merge2)
    c5 = BatchNormalization()(c5)
    c5 = Activation('relu')(c5)

    outputs = Conv2D(filters=1, kernel_size=(1, 1), activation='sigmoid')(c5)
    
    model = Model(inputs=inputs, outputs=outputs)
    return model

# 创建模型
model = convlstm_unet(input_size=(3, 128, 128, 3))
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 打印模型结构
model.summary()



import os
import numpy as np
from keras_preprocessing.image import load_img, img_to_array

# 加载图像序列和对应的掩膜序列
def load_sequence_data(image_folder, mask_folder, timesteps):
    image_files = sorted(os.listdir(image_folder))
    mask_files = sorted(os.listdir(mask_folder))
    
    images = []
    masks = []

    # 确保图像和掩膜数量一致
    assert len(image_files) == len(mask_files), "图像和掩膜数量不一致"

    num_samples = len(image_files) - timesteps + 1

    for i in range(num_samples):
        img_sequence = []
        mask_sequence = []

        for t in range(timesteps):
            img_path = os.path.join(image_folder, image_files[i + t])
            mask_path = os.path.join(mask_folder, mask_files[i + t])
            
            img = load_img(img_path, target_size=(128, 128))
            img = img_to_array(img) / 255.0
            img_sequence.append(img)
            
            mask = load_img(mask_path, target_size=(128, 128), color_mode='grayscale')
            mask = img_to_array(mask) / 255.0
            mask_sequence.append(mask)
        
        images.append(np.array(img_sequence))
        # 对于掩膜，我们只需要最后一帧的掩膜
        masks.append(mask_sequence[-1])
        
    return np.array(images), np.array(masks)

# 定义时间步长
timesteps = 3

# 调用函数加载数据
train_images, train_masks = load_sequence_data(
    r"C:\Users\jack8\OneDrive\113-1_school and company\main file\data\Mix train\train\img",
    r"C:\Users\jack8\OneDrive\113-1_school and company\main file\data\Mix train\train\mask",
    timesteps
)

print("训练数据形状：", train_images.shape)
print("训练掩膜形状：", train_masks.shape)

###幫我增加過擬合機制，10epoch