基于深度神经网络的图像分割python代码

以下是使用深度学习进行图像分割的Python代码示例，其中使用的是U-Net模型：

import numpy as np
import matplotlib.pyplot as plt
import os
import cv2
from skimage import io
from keras.models import Model
from keras.layers import Input, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint
from keras import backend as K

# 设置图像大小和路径
IMG_WIDTH = 256
IMG_HEIGHT = 256
IMG_CHANNELS = 3
TRAIN_PATH = 'train/'
TEST_PATH = 'test/'

# 预处理训练数据
def preprocess_train_data():
    train_ids = next(os.walk(TRAIN_PATH))[1]
    X_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
    Y_train = np.zeros((len(train_ids), IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
    print('Preprocessing train data...')
    for i, id_ in enumerate(train_ids):
        path = TRAIN_PATH + id_
        img = io.imread(path+'/images/'+id_+'.png')[:,:,:IMG_CHANNELS]
        img = cv2.resize(img, (IMG_HEIGHT, IMG_WIDTH), interpolation=cv2.INTER_AREA)
        X_train[i] = img
        mask = np.zeros((IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool)
        for mask_file in next(os.walk(path+'/masks/'))[2]:
            mask_ = io.imread(path+'/masks/'+mask_file)
            mask_ = cv2.resize(mask_, (IMG_HEIGHT, IMG_WIDTH), interpolation=cv2.INTER_AREA)
            mask_ = np.expand_dims(mask_, axis=-1)
            mask = np.maximum(mask, mask_)
        Y_train[i] = mask
    return X_train, Y_train

# 构建U-Net模型
def unet():
    inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
    s = inputs

    c1 = Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (s)
    c1 = Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c1)
    p1 = MaxPooling2D((2, 2)) (c1)

    c2 = Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (p1)
    c2 = Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c2)
    p2 = MaxPooling2D((2, 2)) (c2)

    c3 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (p2)
    c3 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c3)
    p3 = MaxPooling2D((2, 2)) (c3)

    c4 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (p3)
    c4 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c4)
    p4 = MaxPooling2D(pool_size=(2, 2)) (c4)

    c5 = Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (p4)
    c5 = Conv2D(256, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c5)

    u6 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same') (c5)
    u6 = concatenate([u6, c4])
    c6 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (u6)
    c6 = Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c6)

    u7 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same') (c6)
    u7 = concatenate([u7, c3])
    c7 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (u7)
    c7 = Conv2D(64, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c7)

    u8 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same') (c7)
    u8 = concatenate([u8, c2])
    c8 = Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (u8)
    c8 = Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c8)

    u9 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same') (c8)
    u9 = concatenate([u9, c1], axis=3)
    c9 = Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (u9)
    c9 = Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same') (c9)

    outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)

    model = Model(inputs=[inputs], outputs=[outputs])
    model.compile(optimizer=Adam(lr=1e-4), loss='binary_crossentropy', metrics=['accuracy'])
    model.summary()
    
    return model

# 训练模型
def train():
    print('Loading train data...')
    X_train, Y_train = preprocess_train_data()
    model = unet()
    print('Fitting model...')
    earlystopper = EarlyStopping(patience=5, verbose=1)
    checkpointer = ModelCheckpoint('model.h5', verbose=1, save_best_only=True)
    results = model.fit(X_train, Y_train, validation_split=0.1, batch_size=16, epochs=50, 
                        callbacks=[earlystopper, checkpointer])
    return results

# 预测测试数据
def predict():
    print('Loading test data...')
    test_ids = next(os.walk(TEST_PATH))[1]
    X_test = np.zeros((len(test_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS), dtype=np.uint8)
    sizes_test = []
    for i, id_ in enumerate(test_ids):
        path = TEST_PATH + id_
        img = io.imread(path+'/images/'+id_+'.png')[:,:,:IMG_CHANNELS]
        sizes_test.append([img.shape[0], img.shape[1]])
        img = cv2.resize(img, (IMG_HEIGHT, IMG_WIDTH), interpolation=cv2.INTER_AREA)
        X_test[i] = img
    print('Loading model...')
    model = unet()
    print('Predicting masks on test data...')
    preds_test = model.predict(X_test, verbose=1)
    preds_test_t = (preds_test > 0.5).astype(np.uint8)
    return preds_test_t, sizes_test

# 显示图像及其掩模
def show_image_mask(image, mask):
    plt.figure(figsize=(10, 5))
    plt.subplot(121)
    plt.imshow(image)
    plt.title('Image')
    plt.subplot(122)
    plt.imshow(mask, cmap='gray')
    plt.title('Mask')

# 显示预测结果
def show_predict_results(preds_test_t, sizes_test):
    print('Resizing predicted masks to original images...')
    preds_test_upsampled = []
    for i in range(len(preds_test_t)):
        preds_test_upsampled.append(cv2.resize(np.squeeze(preds_test_t[i]), 
                                               (sizes_test[i][1], sizes_test[i][0]), 
                                               interpolation=cv2.INTER_AREA))
    print('Showing prediction results...')
    for i in range(len(preds_test_upsampled)):
        show_image_mask(X_test[i], preds_test_upsampled[i])
        
# 训练和预测
results = train()
preds_test_t, sizes_test = predict()
show_predict_results(preds_test_t, sizes_test)

这个代码示例使用U-Net模型进行图像分割，在训练过程中使用了早期停止技术和模型检查点技术，在预测过程中使用了图像插值技术。通过修改代码中的路径和参数，可以使用自己的数据集进行训练和预测。