使用Python和常见深度学习库(如`TensorFlow`和`Keras`)实现简单医学图像分割的示例

发布于:2025-04-04 ⋅ 阅读:(78) ⋅ 点赞:(0)

以下是一个使用Python和常见深度学习库(如TensorFlowKeras)实现简单医学图像分割的示例。这里以U-Net网络为例,用于医学图像分割。

import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
import os
import cv2


# 构建U-Net模型
def build_unet(input_shape):
    inputs = layers.Input(input_shape)

    # 编码器部分
    c1 = layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(inputs)
    c1 = layers.Dropout(0.1)(c1)
    c1 = layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c1)
    p1 = layers.MaxPooling2D((2, 2))(c1)

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

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

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

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

    # 解码器部分
    u6 = layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c5)
    u6 = layers.concatenate([u6, c4])
    c6 = layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(u6)
    c6 = layers.Dropout(0.2)(c6)
    c6 = layers.Conv2D(128, (3, 3), activation='relu', kernel_initializer='he_normal', padding='same')(c6)

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

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

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

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

    model = models.Model(inputs=[inputs], outputs=[outputs])
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    return model


# 加载和预处理数据
def load_data(data_dir, mask_dir, img_size):
    images = []
    masks = []
    for img_name in os.listdir(data_dir):
        img_path = os.path.join(data_dir, img_name)
        mask_path = os.path.join(mask_dir, img_name)

        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, img_size)
        img = img / 255.0
        img = np.expand_dims(img, axis=-1)
        images.append(img)

        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        mask = cv2.resize(mask, img_size)
        mask = mask / 255.0
        mask = np.expand_dims(mask, axis=-1)
        masks.append(mask)

    images = np.array(images)
    masks = np.array(masks)
    return images, masks


# 主函数
def main():
    # 数据路径
    data_dir = 'path/to/your/images'
    mask_dir = 'path/to/your/masks'
    img_size = (256, 256)

    # 加载数据
    images, masks = load_data(data_dir, mask_dir, img_size)

    # 构建模型
    model = build_unet(input_shape=(256, 256, 1))

    # 训练模型
    model.fit(images, masks, epochs=10, batch_size=16)

    # 保存模型
    model.save('medical_image_segmentation_model.h5')


if __name__ == "__main__":
    main()
    

代码说明:

  1. U-Net模型构建build_unet函数构建了一个U-Net网络,包含编码器和解码器部分。
  2. 数据加载和预处理load_data函数用于加载医学图像和对应的分割掩码,并进行预处理,如调整大小和归一化。
  3. 主函数main函数中,你需要将data_dirmask_dir替换为实际的图像和掩码数据路径,然后加载数据、构建模型、训练模型并保存模型。

运行步骤:

  1. 确保你已经安装了TensorFlowNumPyOpenCV库。
  2. 将代码保存为medical_image_segmentation.py文件。
  3. 准备好医学图像和对应的分割掩码数据,并将其路径替换到代码中的data_dirmask_dir
  4. 在VSCode中打开该文件,在终端中运行python medical_image_segmentation.py

网站公告

今日签到

点亮在社区的每一天
去签到