无人机航拍图像分割: 使用 DeepLabv3+ 和 U-Net 在 Dronet 数据集上进行图像分割

目录

1. 准备数据集

2. 构建模型

2.1 DeepLabv3+ 模型

2.2 U-Net 模型

3. 训练模型

4. 评估模型

5. 可视化预测结果

6. 结论


在本博客中,我们将探讨如何使用DeepLabv3+或U-Net模型在Dronet数据集上进行无人机航拍图像分割任务。我们将分步介绍整个实验流程,包括数据集准备、模型训练和评估。

1. 准备数据集

首先,我们需要获取Dronet数据集。Dronet数据集包含大量的无人机航拍图像和对应的分割标签。假设我们已经下载了数据集并将其解压缩到了data/dronet文件夹中。

接下来,我们将使用Python和相关库处理数据集。我们需要安装以下库:

  • numpy
  • pandas
  • opencv
  • matplotlib
  • tensorflow
  • keras

安装库后,我们可以开始处理数据集。首先,我们将读取图像和标签文件,并将它们划分为训练集和测试集:

import os
import numpy as np
import cv2
from sklearn.model_selection import train_test_split

DATA_DIR = "data/dronet"

# 获取所有图像文件名和标签文件名
images = [f for f in os.listdir(os.path.join(DATA_DIR, "images")) if f.endswith(".png")]
labels = [f for f in os.listdir(os.path.join(DATA_DIR, "labels")) if f.endswith(".png")]

# 创建训练集和测试集
train_images, test_images, train_labels, test_labels = train_test_split(images, labels, test_size=0.2, random_state=42)

# 加载图像和标签数据
def load_data(image_files, label_files):
    image_data = []
    label_data = []

    for img_file, lbl_file in zip(image_files, label_files):
        img_path = os.path.join(DATA_DIR, "images", img_file)
        lbl_path = os.path.join(DATA_DIR, "labels", lbl_file)

        img = cv2.imread(img_path)
        img = cv2.resize(img, (256, 256))
        lbl = cv2.imread(lbl_path, cv2.IMREAD_GRAYSCALE)
        lbl = cv2.resize(lbl, (256, 256))

        image_data.append(img)
        label_data.append(lbl)

    return np.array(image_data), np.array(label_data)

X_train, y_train = load_data(train_images, train_labels)
X_test, y_test = load_data(test_images, test_labels)

2. 构建模型

在这一步中,我们将实现DeepLabv3+和U-Net模型,并根据需要选择其中一个进行训练。为了简化,我们将使用Keras库来实现这些模型。

2.1 DeepLabv3+ 模型

现在我们将使用TensorFlow和Keras实现DeepLabv3+模型。我们将使用预训练的Xception网络作为骨架网络,并将其与DeepLabv3+的解码器部分结合。

import tensorflow as tf
from tensorflow.keras import layers, Model
from tensorflow.keras.applications import Xception

def create_deeplabv3_plus(input_shape=(256, 256, 3), num_classes=1):
    # 使用预训练的Xception作为骨架网络
    base_model = Xception(input_shape=input_shape, weights='imagenet', include_top=False)

    # 获取骨架网络的输出特征
    x = base_model.output

    # DeepLabv3+解码器部分
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(256, activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.5)(x)

    x = layers.Dense(256, activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dropout(0.5)(x)

    # 最后的分类层
    x = layers.Dense(num_classes, activation='sigmoid')(x)

    # 创建模型
    model = Model(inputs=base_model.input, outputs=x)

    return model

deeplabv3_plus_model = create_deeplabv3_plus()

在这个实现中,我们使用预训练的Xception网络作为骨架网络,其输出特征被送入DeepLabv3+解码器部分。解码器部分包括全局平均池化、全连接层、批量归一化和Dropout层。最后的分类层使用sigmoid激活函数输出每个类别的分数。

2.2 U-Net 模型

除了DeepLabv3+模型,我们还可以实现U-Net模型。U-Net模型由编码器和解码器部分组成,可以有效地处理语义分割问题。

def create_unet(input_shape=(256, 256, 3), num_classes=1):
    # 编码器部分
    inputs = layers.Input(shape=input_shape)

    c1 = layers.Conv2D(64, 3, activation='relu', padding='same')(inputs)
    c1 = layers.Conv2D(64, 3, activation='relu', padding='same')(c1)
    p1 = layers.MaxPooling2D(pool_size=(2, 2))(c1)

    c2 = layers.Conv2D(128, 3, activation='relu', padding='same')(p1)
    c2 = layers.Conv2D(128, 3, activation='relu', padding='same')(c2)
    p2 = layers.MaxPooling2D(pool_size=(2, 2))(c2)

    c3 = layers.Conv2D(256, 3, activation='relu', padding='same')(p2)
    c3 = layers.Conv2D(256, 3, activation='relu', padding='same')(c3)
    p3 = layers.MaxPooling2D(pool_size=(2, 2))(c3)

    c4 = layers.Conv2D(512, 3, activation='relu', padding='same')(p3)
    c4 = layers.Conv2D(512, 3, activation='relu', padding='same')(c4)
    p4 = layers.MaxPooling2D(pool_size=(2, 2))(c4)

    c5 = layers.Conv2D(1024, 3, activation='relu', padding='same')(p4)
    c5 = layers.Conv2D(1024, 3, activation='relu', padding='same')(c5)

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

    u7 = layers.Conv2DTranspose(256, 2, strides=(2, 2), padding='same')(c6)
    u7 = layers.concatenate([u7, c3])
    c7 = layers.Conv2D(256, 3, activation='relu', padding='same')(u7)
    c7 = layers.Conv2D(256, 3, activation='relu', padding='same')(c7)

    u8 = layers.Conv2DTranspose(128, 2, strides=(2, 2), padding='same')(c7)
    u8 = layers.concatenate([u8, c2])
    c8 = layers.Conv2D(128, 3, activation='relu', padding='same')(u8)
    c8 = layers.Conv2D(128, 3, activation='relu', padding='same')(c8)

    u9 = layers.Conv2DTranspose(64, 2, strides=(2, 2), padding='same')(c8)
    u9 = layers.concatenate([u9, c1], axis=3)
    c9 = layers.Conv2D(64, 3, activation='relu', padding='same')(u9)
    c9 = layers.Conv2D(64, 3, activation='relu', padding='same')(c9)

# 最后的分类层
outputs = layers.Conv2D(num_classes, 1, activation='sigmoid')(c9)

# 创建模型
model = Model(inputs=inputs, outputs=outputs)

return model

在这个实现中,我们使用Conv2D层和MaxPooling2D层构建编码器部分,并使用Conv2DTranspose层和concatenate操作构建解码器部分。每个解码器层使用与其对应的编码器层的特征进行连接。最后的分类层使用sigmoid激活函数输出每个类别的分数。

3. 训练模型

现在我们已经定义了DeepLabv3+和U-Net模型,我们可以开始训练模型。在这里,我们将使用Adam优化器和二元交叉熵损失函数。我们还将监控模型的精度和IoU(交并比)评估指标。
 

from tensorflow.keras import optimizers

# 编译DeepLabv3+模型
deeplabv3_plus_model.compile(optimizer=optimizers.Adam(lr=1e-4),
                             loss='binary_crossentropy',
                             metrics=['accuracy', tf.keras.metrics.MeanIoU(num_classes=2)])

# 训练DeepLabv3+模型
deeplabv3_plus_history = deeplabv3_plus_model.fit(X_train, y_train,
                                                  batch_size=32,
                                                  epochs=20,
                                                  validation_data=(X_test, y_test))
                                                  
# 编译U-Net模型
unet_model.compile(optimizer=optimizers.Adam(lr=1e-4),
                   loss='binary_crossentropy',
                   metrics=['accuracy', tf.keras.metrics.MeanIoU(num_classes=2)])

# 训练U-Net模型
unet_history = unet_model.fit(X_train, y_train,
                              batch_size=32,
                              epochs=20,
                              validation_data=(X_test, y_test))

在训练期间,我们还可以通过TensorBoard可视化训练进度:

 
 
# 可视化训练进度
from tensorflow.keras.callbacks import TensorBoard

tb_callback = TensorBoard(log_dir='./logs', update_freq='batch')

deeplabv3_plus_history = deeplabv3_plus_model.fit(X_train, y_train,
                                                  batch_size=32,
                                                  epochs=20,
                                                  validation_data=(X_test, y_test),
                                                  callbacks=[tb_callback])

4. 评估模型

在训练完成后,我们可以使用测试集评估DeepLabv3+和U-Net模型的性能。在这里,我们将使用sklearn.metrics中的分类报告和混淆矩阵来评估模型。

from sklearn.metrics import classification_report, confusion_matrix

# 预测图像标签
y_pred_deeplabv3_plus = deeplabv3_plus_model.predict(X_test)
y_pred_deeplabv3_plus = np.round(y_pred_deeplabv3_plus).astype(int)

y_pred_unet = unet_model.predict(X_test)
y_pred_unet = np.round(y_pred_unet).astype(int)

# 分类报告和混淆矩阵(DeepLabv3+模型)
print("DeepLabv3+模型:")
print(classification_report(y_test.flatten(), y_pred_deeplabv3_plus.flatten()))
print(confusion_matrix(y_test.flatten(), y_pred_deeplabv3_plus.flatten()))

# 分类报告和混淆矩阵(U-Net模型)
print("U-Net模型:")
print(classification_report(y_test.flatten(), y_pred_unet.flatten()))
print(confusion_matrix(y_test.flatten(), y_pred_unet.flatten()))

通过分类报告和混淆矩阵,我们可以了解模型在每个类别上的表现,以及模型在不同类别之间的混淆情况。

5. 可视化预测结果

最后,我们可以使用模型对一些无人机航拍图像进行预测,并可视化预测结果。

import matplotlib.pyplot as plt

# 随机选择几张图像进行预测
random_indexes = np.random.choice(range(len(X_test)), size=5, replace=False)

for i in random_indexes:
    # 预测
    y_pred_deeplabv3_plus = deeplabv3_plus_model.predict(X_test[i:i+1])
    y_pred_unet = unet_model.predict(X_test[i:i+1])

    # 可视化
    fig, axs = plt.subplots(1, 3, figsize=(10, 5))
    axs[0].imshow(X_test[i])
    axs[0].set_title("原图")
    axs[1].imshow(y_pred_deeplabv3_plus[0, :, :, 0], cmap='gray')
    axs[1].set_title("DeepLabv3+ 预测结果")
    axs[2].imshow(y_pred_unet[0, :, :, 0], cmap='gray')
    axs[2].set_title("U-Net 预测结果")
    plt.show()

在这个代码片段中,我们选择了五张测试集中的随机图像,并对它们进行预测和可视化。每个图像的左侧是原始图像,中间是DeepLabv3+模型的预测结果,右侧是U-Net模型的预测结果。我们可以通过这种方式直观地比较两个模型的预测效果。

6. 结论

在这篇博客中,我们介绍了如何使用DeepLabv3+和U-Net模型进行无人机航拍图像分割。我们使用Dronet数据集进行训练和测试,并实现了模型的定义、编译、训练、评估和可视化预测结果。我们还通过分类报告和混淆矩阵比较了两个模型的性能。

在实际应用中,我们可以根据具体问题选择使用DeepLabv3+或U-Net模型进行无人机航拍图像分割任务。这些模型可以用于农业、城市规划、环境监测等领域,可以帮助我们更好地理解和利用航拍图像中的信息。

猜你喜欢

转载自blog.csdn.net/m0_68036862/article/details/130452948