高光谱图像分类 - HybridSN

import numpy as np
import scipy.io as sio
import spectral
from sklearn.decomposition import PCA
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix, cohen_kappa_score
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from operator import truediv

# 3 层 三维卷积，一层二维卷积，三层全连接层
class HybridSN(nn.Module):
    def __init__(self, num_classes=16):
        super(HybridSN, self).__init__()
        # 第一层三维卷积：Conv3d_1, Input:(1, 30, 25, 25), 8 个 7 × 3 × 3 的卷积核
        # Output: (8, 24, 23, 23)
        self.conv1 = nn.Sequential(
            nn.Conv3d(1, 8, (7, 3, 3)),
            nn.ReLU()
        )

        # 第二层三维卷积：Conv3d_2，Input:(8, 24, 23, 23)，16 个 5 × 3 × 3 的卷积核
        # Output: (16, 20, 21, 21)
        self.conv2 = nn.Sequential(
            nn.Conv3d(8, 16, (5, 3, 3)),
            nn.ReLU(inplace=True)
        )

        # 第三层三维卷积：Conv3d_3, Input(16, 20, 21, 21)，32 个 3 × 3 × 3 的卷积核
        # Output: (32, 18, 19, 19)
        self.conv3 = nn.Sequential(
            nn.Conv3d(16, 32, (3, 3, 3), stride=1, padding=0),
            nn.ReLU(inplace=True)
        )

        # 第四层是二维卷积，但是要二维卷积之前就要把第三层的输出 reshape 成可以二维卷积的形状
        # 其实就是样本数和维度数堆叠：32 × 18 = 576，所以这一层二维卷积的输入为 576
        # Output: (64, 17, 17)
        self.conv4_2d = nn.Sequential(
            nn.Conv2d(574, 64, (3, 3)),
            nn.ReLU(inplace=True)
        )

        # 二维卷积层之后就是一个全连接层，而在送入这个全连接层之前，要把二维卷积之后的输出降为一维才能送入全连接层
        # 所以在 forward 里直接使用 flatten()，那么输入全连接层参数即为：64 × 17 × 17 = 18496
        # 这一全连接层有 256 个节点
        self.dense1 = nn.Linear(18496, 256)

        # 接着又是一个全连接层，有 128 个节点
        self.dense2 = nn.Linear(256, 128)

        # 最终输出层，有 num_classes 这么多节点，因为我们最后就是要分 num_classes 这么多地物类别（Indian Pines）数据集
        self.dense_out = nn.Linear(128, num_classes)

        # 论文中使用了 Dropout ，参数为 0.4
        self.drop = nn.Dropout(0.4)

    # 定义前向传播函数
    def forward(self, x):
        out = self.conv1(x) # 第一层三维卷积
        out = self.conv2(out) # 第二层三维卷积
        out = self.conv3(out) # 第三层三维卷积
        out = self.conv4_2d(out.reshape(out.shape[0], -1, 19, 19)) # 第四层二维卷积
        out = out.reshape(out.shape[0], -1)
        out = F.relu(self.dropout(self.dense1(out)))
        out = F.relu(self.dropout(self.dense2(out)))
        out - self.dense_out(out)
        return out

# 对高光谱数据集 X 应用 PCA 变换

def applyPCA(X, num_components):
    newX = np.reshape(X, (-1, X.shape[2]))
    # print(type(newX))
    # print(newX.shape) # (145*145=21025, 200)
    pca = PCA(n_components=num_components, whiten=True)
    newX = pca.fit_transform(newX)
    # print(newX.shape) # (21025, 30)
    newX = np.reshape(newX, (X.shape[0], X.shape[1], num_components))
    # print(newX.shape)
    # print("------")
    return newX


# 对单个像素周围提取 patch 时，边缘像素就无法取了，因此，给这部分像素进行 padding 操作
# 经过 padWithZeros 之后，宽、高会变成 169,169， 也就是左右各 padding 了 12
def padWithZeros(X, margin=12):
    # print(X)
    newX = np.zeros((X.shape[0] + 2 * margin, X.shape[1] + 2 * margin, X.shape[2]))
    #print(newX.shape)  # (169, 169, 30) 全是 0
    x_offset = margin
    y_offset = margin
    newX[x_offset:X.shape[0] + x_offset, y_offset:X.shape[1] + y_offset, :] = X
    return newX


#  对单个像素周围提取 patch 时，边缘像素就无法取了，因此，给这部分像素进行 padding 操作
def createImageCubes(X, y, windowSize=5, removeZeroLabels=True):
    # 给 X 做 padding
    margin = int((windowSize - 1) / 2)  # 12
    zeroPaddedX = padWithZeros(X, margin=margin)  # (169, 30, 30)
    # split patches
    patchesData = np.zeros((X.shape[0] * X.shape[1], windowSize, windowSize, X.shape[2]))
    # print(patchesData.shape) # (21025, 25, 25, 200)
    patchesLabels = np.zeros((X.shape[0] * X.shape[1]))
    # print(patchesLabels.shape) # (21025, )
    patchIndex = 0
    for r in range(margin, zeroPaddedX.shape[0] - margin):  # for r in range(12, 169-12)
        for c in range(margin, zeroPaddedX.shape[1] - margin):  # for c in range(12, 169-12)
            patch = zeroPaddedX[r - margin:r + margin + 1, c - margin:c + margin + 1]
            # print(patch.shape)
            # zeroPaddedX[r - 12 : r + 12 + 1, c -12 : c + 12 + 1]
            patchesData[patchIndex, :, :, :] = patch
            patchesLabels[patchIndex] = y[r - margin, c - margin]
            patchIndex = patchIndex + 1
    # 如果没有 removeZeroLabels
    # patchesData 的 shape 将会是 (21025, 25, 25, 30)
    # patchesLabels 的 shape 将会是 (21025, )
    # removeZeroLabels 的作用就是去掉 gt 标签集 groundtruth 中为 0 的数据，因为里面的数据值有 0~16，而刚好 1~16 刚好对应地物分类的 16 类
    if removeZeroLabels:
        patchesData = patchesData[patchesLabels > 0, :, :, :]
        patchesLabels = patchesLabels[patchesLabels > 0]
        patchesLabels -= 1
    return patchesData, patchesLabels

# 随机划分训练集和测试集
# 这里最后取 testradio=0.9 也就是随机抽 90% 为测试集，10% 为训练集
def splitTrainTestSet(X, y, testRatio, randomState=345):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=testRatio, random_state=randomState, stratify=y)
    return X_train, X_test, y_train, y_test

# 构建训练数据集
""" Training dataset"""
class TrainDS(torch.utils.data.Dataset):
    def __init__(self):
        self.len = Xtrain.shape[0]
        self.x_data = torch.FloatTensor(Xtrain)
        self.y_data = torch.LongTensor(ytrain)

    def __getitem__(self, index):
        # 根据索引返回数据和对应的标签
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        # 返回文件数据的数目
        return self.len


# 构建测试数据集
""" Testing dataset"""
class TestDS(torch.utils.data.Dataset):
    def __init__(self):
        self.len = Xtest.shape[0]
        self.x_data = torch.FloatTensor(Xtest)
        self.y_data = torch.LongTensor(ytest)

    def __getitem__(self, index):
        # 根据索引返回数据和对应的标签
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        # 返回文件数据的数目
        return self.len


# 创建 trainloader 和 testloader
trainset = TrainDS()
testset = TestDS()
train_loader = torch.utils.data.DataLoader(dataset=trainset, batch_size=128, shuffle=True, num_workers=2)
test_loader = torch.utils.data.DataLoader(dataset=testset, batch_size=128, shuffle=False, num_workers=2)

# 选择 GPU 进行训练
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 网络放到GPU上
net = HybridSN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 开始训练
total_loss = 0
# net.train()
for epoch in range(100):
    for i, (inputs, labels) in enumerate(train_loader):
        inputs = inputs.to(device)
        labels = labels.to(device)
        # 优化器梯度归零
        optimizer.zero_grad()
        # 正向传播 +　反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print('[Epoch: %d]   [loss avg: %.4f]   [current loss: %.4f]' % (epoch + 1, total_loss / (epoch + 1), loss.item()))

print('Finished Training')

# 模型测试 (添加 net.eval())

net.eval()
count = 0
for inputs, _ in test_loader:
    inputs = inputs.to(device)
    outputs = net(inputs)
    outputs = np.argmax(outputs.detach().cpu().numpy(), axis=1)
    if count == 0:
        y_pred_test = outputs
        count = 1
    else:
        y_pred_test = np.concatenate( (y_pred_test, outputs))

# 生成分类报告
classification = classification_report(ytest, y_pred_test, digits=4)
print(classification)

def AA_andEachClassAccuracy(confusion_matrix):
    counter = confusion_matrix.shape[0]
    list_diag = np.diag(confusion_matrix)
    list_raw_sum = np.sum(confusion_matrix, axis=1)
    each_acc = np.nan_to_num(truediv(list_diag, list_raw_sum))
    average_acc = np.mean(each_acc)
    return each_acc, average_acc


def reports(test_loader, y_test, name):
    count = 0
    # 模型测试
    for inputs, _ in test_loader:
        inputs = inputs.to(device)
        outputs = net(inputs)
        outputs = np.argmax(outputs.detach().cpu().numpy(), axis=1)
        if count == 0:
            y_pred = outputs
            count = 1
        else:
            y_pred = np.concatenate((y_pred, outputs))


    if name == 'IP':
        target_names = ['Alfalfa', 'Corn-notill', 'Corn-mintill', 'Corn'
            , 'Grass-pasture', 'Grass-trees', 'Grass-pasture-mowed',
                        'Hay-windrowed', 'Oats', 'Soybean-notill', 'Soybean-mintill',
                        'Soybean-clean', 'Wheat', 'Woods', 'Buildings-Grass-Trees-Drives',
                        'Stone-Steel-Towers']
    elif name == 'SA':
        target_names = ['Brocoli_green_weeds_1', 'Brocoli_green_weeds_2', 'Fallow', 'Fallow_rough_plow', 'Fallow_smooth',
                        'Stubble', 'Celery', 'Grapes_untrained', 'Soil_vinyard_develop', 'Corn_senesced_green_weeds',
                        'Lettuce_romaine_4wk', 'Lettuce_romaine_5wk', 'Lettuce_romaine_6wk', 'Lettuce_romaine_7wk',
                        'Vinyard_untrained', 'Vinyard_vertical_trellis']
    elif name == 'PU':
        target_names = ['Asphalt', 'Meadows', 'Gravel', 'Trees', 'Painted metal sheets', 'Bare Soil', 'Bitumen',
                        'Self-Blocking Bricks', 'Shadows']

    classification = classification_report(y_test, y_pred, target_names=target_names)
    oa = accuracy_score(y_test, y_pred)
    confusion = confusion_matrix(y_test, y_pred)
    each_acc, aa = AA_andEachClassAccuracy(confusion)
    kappa = cohen_kappa_score(y_test, y_pred)

    return classification, confusion, oa * 100, each_acc * 100, aa * 100, kappa * 100

classification, confusion, oa, each_acc, aa, kappa = reports(test_loader, ytest, 'IP')
classification = str(classification)
confusion = str(confusion)
file_name = "classification_report.txt"

with open(file_name, 'w') as x_file:
    x_file.write('\n')
    x_file.write('{} Kappa accuracy (%)'.format(kappa))
    x_file.write('\n')
    x_file.write('{} Overall accuracy (%)'.format(oa))
    x_file.write('\n')
    x_file.write('{} Average accuracy (%)'.format(aa))
    x_file.write('\n')
    x_file.write('\n')
    x_file.write('{}'.format(classification))
    x_file.write('\n')
    x_file.write('{}'.format(confusion))

X = sio.loadmat('Indian_pines_corrected.mat')['indian_pines_corrected']
y = sio.loadmat('Indian_pines_gt.mat')['indian_pines_gt']

height = y.shape[0]
width = y.shape[1]

X = applyPCA(X, num_components = num_components)
X = padWithZeros(X, patch_size//2)

# 逐像素预测类别
outputs = np.zeros((height,width))
for i in range(height):
    for j in range(width):
        if int(y[i,j]) == 0:
            continue
        else :
            image_patch = X[i:i+patch_size, j:j+patch_size, :]
            image_patch = image_patch.reshape(1,image_patch.shape[0],image_patch.shape[1], image_patch.shape[2], 1)
            X_test_image = torch.FloatTensor(image_patch.transpose(0, 4, 3, 1, 2)).to(device)
            prediction = net(X_test_image)
            prediction = np.argmax(prediction.detach().cpu().numpy(), axis=1)
            outputs[i][j] = prediction+1
    if i % 20 == 0:
        print('... ... row ', i, ' handling ... ...')
predict_image = spectral.imshow(classes = outputs.astype(int),figsize =(5,5))

# 随机划分训练集和测试集
# 这里最后取 testradio=0.9 也就是随机抽 90% 为测试集，10% 为训练集
def splitTrainTestSet(X, y, testRatio, randomState=345):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=testRatio, random_state=randomState, stratify=y)
    return X_train, X_test, y_train, y_test

# 3 层 三维卷积，一层二维卷积，三层全连接层
class HybridSN(nn.Module):
    def __init__(self, num_classes=16):
        super(HybridSN, self).__init__()
        # 第一层三维卷积：Conv3d_1, Input:(1, 30, 25, 25), 8 个 7 × 3 × 3 的卷积核
        # Output: (8, 24, 23, 23)
        self.conv1 = nn.Sequential(
            nn.Conv3d(1, 8, (7, 3, 3)),
            nn.BatchNorm3d(8),
            nn.ReLU()
        )

        # 第二层三维卷积：Conv3d_2，Input:(8, 24, 23, 23)，16 个 5 × 3 × 3 的卷积核
        # Output: (16, 20, 21, 21)
        self.conv2 = nn.Sequential(
            nn.Conv3d(8, 16, (5, 3, 3)),
            nn.BatchNorm3d(16),
            nn.ReLU(inplace=True)
        )

        # 第三层三维卷积：Conv3d_3, Input(16, 20, 21, 21)，32 个 3 × 3 × 3 的卷积核
        # Output: (32, 18, 19, 19)
        self.conv3 = nn.Sequential(
            nn.Conv3d(16, 32, (3, 3, 3), stride=1, padding=0),
            nn.BatchNorm3d(32),
            nn.ReLU(inplace=True)
        )

        # 第四层是二维卷积，但是要二维卷积之前就要把第三层的输出 reshape 成可以二维卷积的形状
        # 其实就是样本数和维度数堆叠：32 × 18 = 576，所以这一层二维卷积的输入为 576
        # Output: (64, 17, 17)
        self.conv4_2d = nn.Sequential(
            nn.Conv2d(574, 64, (3, 3)),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )

        # 二维卷积层之后就是一个全连接层，而在送入这个全连接层之前，要把二维卷积之后的输出降为一维才能送入全连接层
        # 所以在 forward 里直接使用 flatten()，那么输入全连接层参数即为：64 × 17 × 17 = 18496
        # 这一全连接层有 256 个节点
        self.dense1 = nn.Linear(18496, 256)

        # 接着又是一个全连接层，有 128 个节点
        self.dense2 = nn.Linear(256, 128)

        # 最终输出层，有 num_classes 这么多节点，因为我们最后就是要分 num_classes 这么多地物类别（Indian Pines）数据集
        self.dense_out = nn.Linear(128, num_classes)

        # 论文中使用了 Dropout ，参数为 0.4
        self.drop = nn.Dropout(0.4)

    # 定义前向传播函数
    def forward(self, x):
        out = self.conv1(x) # 第一层三维卷积
        out = self.conv2(out) # 第二层三维卷积
        out = self.conv3(out) # 第三层三维卷积
        out = self.conv4_2d(out.reshape(out.shape[0], -1, 19, 19)) # 第四层二维卷积
        out = out.reshape(out.shape[0], -1)
        out = F.relu(self.dropout(self.dense1(out)))
        out = F.relu(self.dropout(self.dense2(out)))
        out - self.dense_out(out)
        return out

#  对单个像素周围提取 patch 时，边缘像素就无法取了，因此，给这部分像素进行 padding 操作
def createImageCubes(X, y, windowSize=5, removeZeroLabels=True):
    # 给 X 做 padding
    margin = int((windowSize - 1) / 2)  # 12
    zeroPaddedX = padWithZeros(X, margin=margin)  # (169, 30, 30)
    # split patches
    patchesData = np.zeros((X.shape[0] * X.shape[1], windowSize, windowSize, X.shape[2]))
    # print(patchesData.shape) # (21025, 25, 25, 200)
    patchesLabels = np.zeros((X.shape[0] * X.shape[1]))
    # print(patchesLabels.shape) # (21025, )
    patchIndex = 0
    for r in range(margin, zeroPaddedX.shape[0] - margin):  # for r in range(12, 169-12)
        for c in range(margin, zeroPaddedX.shape[1] - margin):  # for c in range(12, 169-12)
            patch = zeroPaddedX[r - margin:r + margin + 1, c - margin:c + margin + 1]
            # print(patch.shape)
            # zeroPaddedX[r - 12 : r + 12 + 1, c -12 : c + 12 + 1]
            patchesData[patchIndex, :, :, :] = patch
            patchesLabels[patchIndex] = y[r - margin, c - margin]
            patchIndex = patchIndex + 1
    # 如果没有 removeZeroLabels
    # patchesData 的 shape 将会是 (21025, 25, 25, 30)
    # patchesLabels 的 shape 将会是 (21025, )
    # removeZeroLabels 的作用就是去掉 gt 标签集 groundtruth 中为 0 的数据，因为里面的数据值有 0~16，而刚好 1~16 刚好对应地物分类的 16 类
    if removeZeroLabels:
        patchesData = patchesData[patchesLabels > 0, :, :, :]
        patchesLabels = patchesLabels[patchesLabels > 0]

# 在每个像素周围提取 patch 
def createImageCubes(X, y, windowSize=5, removeZeroLabels = True):
    # 给 X 做 padding
    margin = int((windowSize - 1) / 2)
    zeroPaddedX = padWithZeros(X, margin=margin)
    # 获得 y 中的标记样本数
    count = 0
    for r in range(0, y.shape[0]):
        for c in range(0, y.shape[1]):
            if y[r, c] != 0:
                count = count+1

    # split patches
    patchesData = np.zeros([count, windowSize, windowSize, X.shape[2]])
    patchesLabels = np.zeros(count)

    count = 0
    for r in range(margin, zeroPaddedX.shape[0] - margin):
        for c in range(margin, zeroPaddedX.shape[1] - margin):
            if y[r-margin, c-margin] != 0:
                patch = zeroPaddedX[r - margin:r + margin + 1, c - margin:c + margin + 1]   
                patchesData[count, :, :, :] = patch
                patchesLabels[count] = y[r-margin, c-margin]
                count = count + 1

    return patchesData, patchesLabels