3d-qcnet_train

import os
import numpy as np
import nibabel as nib
import cv2
from keras.models import Model
from keras.layers import Dense
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint
from densenet import DenseNet3D  # 从 densenet.py 导入 DenseNet3D 模型
from sklearn.model_selection import train_test_split
from keras.utils import to_categorical

# Function to load and preprocess data
def load_and_preprocess_data(data_path):
    X_data = []
    y_data = []

    # Assuming the data is structured in folders: 'good' and 'bad'
    for label in ['good', 'bad']:
        label_path = os.path.join(data_path, label)
        for file in os.listdir(label_path):
            if file.endswith('.nii') or file.endswith('.nii.gz'):  # 只加载 .nii 文件
                img = nib.load(os.path.join(label_path, file)).get_fdata()  # Load the 3D volume
                img_resized = np.zeros((96, 96, 70))  # Target size (96, 96, 70)

                # Clip or pad the volume to make it 70 slices deep
                depth = img.shape[-1]
                if depth > 70:
                    img = img[:, :, :70]  # Clip to 70 slices if the depth is greater than 70
                elif depth < 70:
                    pad_width = [(0, 0), (0, 0), (0, 70 - depth)]  # Pad to 70 slices
                    img = np.pad(img, pad_width, mode='constant', constant_values=0)

                # Resize each slice to 96x96
                for i in range(70):
                    slice_img = cv2.resize(img[:, :, i], (96, 96))  # Resize each 2D slice to 96x96
                    img_resized[:, :, i] = slice_img

                X_data.append(img_resized)  # Add resized volume to list
                y_data.append(0 if label == 'bad' else 1)  # 0 for 'bad', 1 for 'good'

    # Convert to numpy arrays and reshape for input into the model
    X_data = np.array(X_data).reshape(-1, 96, 96, 70, 1)  # Add channel dimension
    y_data = np.array(y_data)

    return X_data, y_data

# Function to train the model
def train_model(data_path, model_save_path, epochs=20, batch_size=5, learning_rate=1e-4):
    # Load and preprocess the data
    X_data, y_data = load_and_preprocess_data(data_path)
    
    # 将标签转换为 one-hot 编码
    y_data = to_categorical(y_data, num_classes=2)

    # Split data into training and validation sets
    X_train, X_val, y_train, y_val = train_test_split(X_data, y_data, test_size=0.2, random_state=42)

    # Create the model using the DenseNet3D architecture
    model = DenseNet3D(input_shape=(96, 96, 70, 1), include_top=True)

    # Add final classification layer (for 2 classes)
    output_layer = Dense(2, activation='softmax')(model.output)
    model = Model(inputs=model.input, outputs=output_layer)

    # Compile the model with Adam optimizer and binary crossentropy loss
    model.compile(optimizer=Adam(lr=learning_rate), loss='binary_crossentropy', metrics=['accuracy'])

    # Create a checkpoint callback to save the best model during training
    checkpoint = ModelCheckpoint(model_save_path, monitor='val_loss', save_best_only=True, save_weights_only=True)

    # Train the model with the specified batch size, epochs, and the checkpoint callback
    model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, validation_data=(X_val, y_val),
              callbacks=[checkpoint])

    print("Model training complete and saved to:", model_save_path)

# Main entry point of the script
if __name__ == "__main__":
    # 直接在这里定义需要的路径和参数，而不是使用命令行参数
    data_path = "/path/to/your/data"  # 替换为实际数据路径
    model_save_path = "/path/to/save/model.h5"  # 替换为保存模型的路径
    epochs = 20  # 训练轮数
    batch_size = 5  # 批量大小
    learning_rate = 1e-4  # 学习率

    # 调用训练模型函数
    train_model(data_path, model_save_path, epochs=epochs, batch_size=batch_size, learning_rate=learning_rate)