python - 使用 UNet 分割多种材料

Question

我正在尝试在 nifti 文件 (.nii) 中分割多种材料，但每当我运行我的 UNet 时，它只会分割一种材料并返回 0 和 1（黑色和白色）的标签。我想返回 0、1 和 2 的标签。我n_labels = 3在定义模型时尝试包括在内，但模型仍然返回 0 和 1 的标签。我的训练标签用未标记、内部和外部分段。

def get_model(n_labels = 3):
    inputs = Input((sample_width, sample_height, sample_depth, 1))
    conv1 = Conv3D(n_labels, (3, 3, 3), activation='relu', padding='same')(inputs)
    conv1 = Conv3D(n_labels, (3, 3, 3), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling3D(pool_size=(2, 2, 2))(conv1)
    drop1 = Dropout(0.5)(pool1)

    conv2 = Conv3D(64, (3, 3, 3), activation='relu', padding='same')(drop1)
    conv2 = Conv3D(64, (3, 3, 3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling3D(pool_size=(2, 2, 2))(conv2)
    drop2 = Dropout(0.5)(pool2)

    conv3 = Conv3D(128, (3, 3, 3), activation='relu', padding='same')(drop2)
    conv3 = Conv3D(128, (3, 3, 3), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling3D(pool_size=(2, 2, 2))(conv3)
    drop3 = Dropout(0.3)(pool3)

    conv4 = Conv3D(256, (3, 3, 3), activation='relu', padding='same')(drop3)
    conv4 = Conv3D(256, (3, 3, 3), activation='relu', padding='same')(conv4)
    pool4 = MaxPooling3D(pool_size=(2, 2, 2))(conv4)
    drop4 = Dropout(0.3)(pool4)

    conv5 = Conv3D(512, (3, 3, 3), activation='relu', padding='same')(drop4)
    conv5 = Conv3D(512, (3, 3, 3), activation='relu', padding='same')(conv5)

    up6 = concatenate([Conv3DTranspose(256, (2, 2, 2), strides=(2, 2, 2), padding='same')(conv5), conv4], axis=4)
    conv6 = Conv3D(256, (3, 3, 3), activation='relu', padding='same')(up6)
    conv6 = Conv3D(256, (3, 3, 3), activation='relu', padding='same')(conv6)

    up7 = concatenate([Conv3DTranspose(128, (2, 2, 2), strides=(2, 2, 2), padding='same')(conv6), conv3], axis=4)
    conv7 = Conv3D(128, (3, 3, 3), activation='relu', padding='same')(up7)
    conv7 = Conv3D(128, (3, 3, 3), activation='relu', padding='same')(conv7)

    up8 = concatenate([Conv3DTranspose(64, (2, 2, 2), strides=(2, 2, 2), padding='same')(conv7), conv2], axis=4)
    conv8 = Conv3D(64, (3, 3, 3), activation='relu', padding='same')(up8)
    conv8 = Conv3D(64, (3, 3, 3), activation='relu', padding='same')(conv8)

    up9 = concatenate([Conv3DTranspose(n_labels, (2, 2, 2), strides=(2, 2, 2), padding='same')(conv8), conv1], axis=4)
    conv9 = Conv3D(n_labels, (3, 3, 3), activation='relu', padding='same')(up9)
    conv9 = Conv3D(n_labels, (3, 3, 3), activation='relu', padding='same')(conv9)

    conv10 = Conv3D(1, (1, 1, 1), activation='sigmoid')(conv9)

    model = Model(inputs=[inputs], outputs=[conv10])
    model.compile(optimizer=Adam(learning_rate = 1e-4), loss=dice_coef_loss, metrics=[dice_coef])
    return model

smooth = 1.
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return 1-dice_coef(y_true, y_pred)

n_labels = 3
model = get_model(n_labels)

observe_var = 'dice_coef'
strategy = 'max'

model_dir = '//models//'
model_checkpoint = ModelCheckpoint(model_dir + '{epoch:04}.h5', monitor=observe_var, mode='auto', 
                save_weights_only=False, save_best_only=False, period = 5)
model.fit(train_x, train_y, batch_size = 2, epochs= 5000, verbose=1, shuffle=True, validation_split=.15, 
                callbacks=[model_checkpoint])
model.save(model_dir + 'final_3d.h5')

标签： pythontensorflowimage-segmentationunity3d-unetnifti