实战| 使用深度学习分割和计算水体和农田面积【Pytorch附源码】

《博主简介》

小伙伴们好，我是阿旭。专注于人工智能、AIGC、python、计算机视觉相关分享研究。
👍感谢小伙伴们点赞、关注！

《------往期经典推荐------》

一、AI应用软件开发实战专栏【链接】

二、机器学习实战专栏【链接】，已更新31期，欢迎关注，持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
五、YOLOv8改进专栏【链接】，持续更新中~~
六、YOLO性能对比专栏【链接】，持续更新中~

《------正文------》

引言

在本文中，我们将介绍如何使用深度学习从卫星图像中分离水体（如湖泊和河流）和农田（植物生长的地方）。我们还将估计这些区域的面积，比如水体或农田覆盖了多少英亩。本文将包括PyTorch代码，并将以简单的步骤为初学者进行解释说明。

步骤1：收集图像（GeoTIFF格式）

我们需要收集卫星图像，这些图像通常以GeoTIFF格式存储。这些图像不仅仅是普通的照片-它们还存储了地理信息，这有助于我们以后测量面积。

您可以从以下位置获取GeoTIFF图像：

• Google Earth：用于查看地图。
• Sentinel-2或Landsat：用于高质量的卫星图像。

步骤2：将GeoTIFF转换为常规图像

由于我们的深度学习模型不需要存储在GeoTIFF中的地理信息，我们首先将GeoTIFF文件转换为常规图像格式，如PNG或JPEG。

Python代码：将GeoTIFF转换为PNG

import rasterio
from matplotlib import pyplot as plt
import numpy as np
# Open GeoTIFF file
with rasterio.open('path_to_geotiff.tif') as dataset:
    # Read the image data as an array (RGB bands)
    image = dataset.read([1, 2, 3])
    image = np.transpose(image, (1, 2, 0))  # Rearrange to (H, W, C) format
    # Save the image as PNG
    plt.imsave('output_image.png', image)

现在，我们有了一个可以在深度学习模型中使用的常规图像。

步骤3：图像预处理

在我们使用图像进行训练之前，我们需要调整图像的大小并将其标准化。这个过程称为预处理。

预处理代码

import cv2
import torch
# Function to load and resize images
def preprocess_image(image_path, size=(256, 256)):
    image = cv2.imread(image_path)
    resized_image = cv2.resize(image, size)
    normalized_image = resized_image / 255.0
    return normalized_image
# Example usage
processed_image = preprocess_image('output_image.png')
tensor_image = torch.tensor(processed_image).permute(2, 0, 1).unsqueeze(0)

步骤4：选择模型（U-Net）

对于这个任务，我们使用了U-Net模型，它非常适合分离图像的不同部分，比如水体和农田。

U-Net模型构建

下面是一个简化版本的U-Net模型构建代码：

import torch
import torch.nn as nn
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 1, kernel_size=1)
        )
    def forward(self, x):
        enc = self.encoder(x)
        dec = self.decoder(enc)
        return dec

步骤5：训练模型

接下来，我们将使用标记的图像（我们绘制线条以显示水体和农田位置的图像）来训练模型。

训练代码

from torch.optim import Adam
# Initialize model and optimizer
model = UNet()
optimizer = Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
# Training loop
def train_model(model, train_loader, num_epochs):
    model.train()
    for epoch in range(num_epochs):
        epoch_loss = 0
        for images, masks in train_loader:
            optimizer.zero_grad()
            outputs = model(images)
            loss = loss_fn(outputs, masks)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}')

步骤6：分割水体和农田

在训练模型之后，我们使用它来预测新图像上的水体和农田，进行图像分割。

预测代码

def predict(model, image):
    model.eval()
    with torch.no_grad():
        output = model(image)
        pred = torch.argmax(output, dim=1)
        return pred

步骤7：计算面积

要计算真实世界的面积（以英亩为单位），我们需要返回到原始GeoTIFF文件以查找每个像素的大小（以米为单位）。这将帮助我们计算分割区域的面积，如水体和农田。

面积计算代码

def calculate_area(segmented_image, pixel_size=10):
    num_water_pixels = torch.sum(segmented_image == 1).item()  # Class 1 for water
    num_crop_pixels = torch.sum(segmented_image == 2).item()   # Class 2 for crops
    
    # Calculate area in square meters
    water_area = num_water_pixels * (pixel_size ** 2)
    crop_area = num_crop_pixels * (pixel_size ** 2)
    
    # Convert to acres (1 acre = 4046.86 square meters)
    water_area_acres = water_area / 4046.86
    crop_area_acres = crop_area / 4046.86
    
    return water_area_acres, crop_area_acres
# Open the GeoTIFF file to get the pixel size
with rasterio.open('path_to_geotiff.tif') as dataset:
    pixel_size = dataset.res[0]
# Example usage
water_acres, crop_acres = calculate_area(segmented_image, pixel_size)
print(f'Water Area: {water_acres:.2f} acres, Crop Area: {crop_acres:.2f} acres')

结论

在本文中，我们学习了如何以GeoTIFF格式拍摄卫星图像，将其转换为常规图像，并使用深度学习来分离水体和农田。分割后，我们可以使用GeoTIFF中的地理信息计算这些区域的面积。通过遵循这些简单的步骤，即使是初学者也可以使用卫星图像和深度学习模型。

好了，这篇文章就介绍到这里，喜欢的小伙伴感谢给点个赞和关注，更多精彩内容持续更新~~
关于本篇文章大家有任何建议或意见，欢迎在评论区留言交流！