基于 U-Net 的农田边界分割系统实现（前后端分离）

       本项目通过前后端分离的架构，构建了一个基于深度学习的农田边界自动提取系统。采用 Flask 搭建后端 API，前端使用 React 实现图像上传与预测结果展示，核心模型为轻量级的 U-Net，用于完成遥感图像中的农田边界分割。以下是该系统的架构目录和流程图。

# 后端（Backend）
├── model/
│   ├── unet.py               # U-Net模型定义
│   └── unet_farm.pth         # 训练好的模型权重
├── train/
│   ├── dataset.py            # 数据加载器
│   ├── train_unet.py         # 训练脚本
│   └── transforms.py         # 数据增强
├── utils/
│   └── preprocess.py         # 图像预处理
├── app.py                    # Flask主程序
├── requirements.txt          # 依赖库
├── uploads/                  # 上传文件临时存储
└── results/                  # 预测结果输出
# 前端（Frontend）
├── public/
│   └── index.html            # HTML入口
├── src/
│   ├── components/
│   │   └── UploadImage.js    # 上传组件
│   ├── App.js                # 主组件
│   ├── index.js              # 入口文件
│   └── index.css             # 全局样式
├── package.json              # 前端依赖
└── package-lock.json
# 数据（示例结构）
└── data/
    ├── images/               # 原始图像
    └── masks/                # 标注掩码

        本系统由前后端协作组成，整体可分为以下五大功能模块：

        图像上传模块：用户通过前端界面上传本地的 RGB 农田遥感图像。该模块提供图像拖拽与选择功能，并在上传后立即在网页端显示预览，提升用户交互体验。

        图像预处理模块：后端接收到图像后，利用 Pillow 和 OpenCV 等图像处理库对输入图像进行标准化预处理，包括尺寸调整、归一化处理和通道格式转换，确保图像符合模型输入要求。

        分割模型推理模块：系统采用基于 U-Net 的图像语义分割模型，对预处理后的图像进行像素级预测。该模块加载训练好的模型权重，执行前向推理，输出分割掩膜（即边界图像）。

        结果可视化与下载模块：推理完成后，系统将分割结果进行后处理（如阈值处理、边界轮廓提取），并通过前端界面显示。同时，用户可一键下载分割图像，用于后续标注或分析任务。

        后端 API 接口模块：采用 Flask 框架实现 RESTful API 服务，负责处理前端图像请求、模型加载与推理任务，并返回推理结果。支持多并发访问，具备良好的接口稳定性。

一、后端设计（Flask + U-Net）

        采用经典的 U-Net 结构，适用于小样本遥感图像分割。

import torch
import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        def conv_block(in_ch, out_ch):
            return nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 3, padding=1),
                nn.ReLU(),
                nn.Conv2d(out_ch, out_ch, 3, padding=1),
                nn.ReLU()
            )

        self.encoder1 = conv_block(in_channels, 64)
        self.pool1 = nn.MaxPool2d(2)
        self.encoder2 = conv_block(64, 128)
        self.pool2 = nn.MaxPool2d(2)

        self.bottleneck = conv_block(128, 256)

        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.decoder2 = conv_block(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.decoder1 = conv_block(128, 64)

        self.final = nn.Conv2d(64, out_channels, 1)

    def forward(self, x):
        e1 = self.encoder1(x)
        e2 = self.encoder2(self.pool1(e1))
        b = self.bottleneck(self.pool2(e2))
        d2 = self.decoder2(torch.cat([self.up2(b), e2], dim=1))
        d1 = self.decoder1(torch.cat([self.up1(d2), e1], dim=1))
        return self.final(d1)

        模型预测接口部分，这是前端与模型的桥梁，负责接收用户上传的图像 → 调用模型进行预测 → 返回分割结果。

@app.route("/predict", methods=["POST"])
def predict_route():
    if 'file' not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files['file']
    filename = str(uuid.uuid4()) + '.tif'
    filepath = os.path.join(UPLOAD_FOLDER, filename)
    file.save(filepath)

    try:
        # 加载图像
        image = Image.open(filepath).convert("RGB")
        image = image.resize((512, 512))  # 调整尺寸

        # 预处理为tensor
        image_tensor = preprocess_image(image)  # 现在传递的是PIL Image

        prediction = run_prediction(image_tensor)  # 返回numpy数组预测结果

        # 转为图像保存
        result_img = Image.fromarray((prediction * 255).astype(np.uint8)).convert("L")
        result_path = os.path.join(RESULT_FOLDER, filename.replace('.tif', '_mask.png'))
        result_img.save(result_path)

        return send_file(result_path, mimetype='image/png')

    except Exception as e:
        print(f"后端错误: {e}")
        return jsonify({'error': str(e)}), 500

二、训练过程

        我选取的数据集是看过的《A Comprehensive Deep-Learning Framework for Fine-Grained Farmland Mapping From High-Resolution Images》这一篇文献里面的，选取了部分训练，大家可以去下载也可以用本地的数据集，都没有关系。

import os
from PIL import Image
from torch.utils.data import Dataset
import torchvision.transforms as T

class FarmDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.images = os.listdir(image_dir)
        self.transform = transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_name = self.images[idx]
        img_path = os.path.join(self.image_dir, img_name)

        # 替换 "image" 为 "label"，并将后缀改为 .png
        label_name = img_name.replace("image", "label").rsplit(".", 1)[0] + ".png"
        mask_path = os.path.join(self.mask_dir, label_name)

        # 检查路径是否存在
        if not os.path.exists(mask_path):
            raise FileNotFoundError(f"标签不存在: {mask_path}")

        image = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L")

        if self.transform:
            image, mask = self.transform(image, mask)

        return image, mask

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from model.unet import UNet
from dataset import FarmDataset
from transforms import SegmentationTransform

transform = SegmentationTransform(size=(256, 256))
dataset = FarmDataset("D://农田分割系统//farm-boundary-segmentation//data//images", "D://农田分割系统//farm-boundary-segmentation//data//masks", transform=transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet(3, 1).to(device)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(10):
    model.train()
    total_loss = 0
    for imgs, masks in dataloader:
        imgs, masks = imgs.to(device), masks.to(device)
        outputs = model(imgs)
        loss = criterion(outputs, masks)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")

torch.save(model.state_dict(), "model/unet_farm.pth")

三、前端界面

        上传图像组件（UploadImage.js）：

import React, { useState } from 'react';
import axios from 'axios';
import * as UTIF from 'utif';

function UploadImage() {
  const [image, setImage] = useState(null);
  const [result, setResult] = useState(null);
  const [loading, setLoading] = useState(false);
  const [preview, setPreview] = useState(null);

  const handleChange = async (e) => {
    const file = e.target.files[0];
    if (!file) return;

    try {
      const arrayBuffer = await file.arrayBuffer();
      const ifds = UTIF.decode(arrayBuffer); // 解析TIFF
      UTIF.decodeImage(arrayBuffer, ifds[0]); // 解码第一个图像
      const rgba = UTIF.toRGBA8(ifds[0]);    // 转换为RGBA

      // 创建Canvas渲染
      const canvas = document.createElement('canvas');
      canvas.width = ifds[0].width;
      canvas.height = ifds[0].height;
      const ctx = canvas.getContext('2d');
      const imageData = new ImageData(
        new Uint8ClampedArray(rgba.buffer),
        ifds[0].width,
        ifds[0].height
      );
      ctx.putImageData(imageData, 0, 0);
      setPreview(canvas.toDataURL('image/png'));
      setImage(file);
    } catch (err) {
      console.error('TIFF解析失败:', err);
      alert('请上传有效的TIFF文件');
    }
  };

  const handleUpload = async () => {
    if (!image) return alert('请先选择图片');
    setLoading(true);

    const formData = new FormData();
    formData.append('file', image);

    try {
      const res = await axios.post('http://localhost:5000/predict', formData, {
        responseType: 'blob',
      });
      const url = URL.createObjectURL(res.data);
      setResult(url);
    } catch (err) {
      console.error(err);
      alert('预测失败，请检查后端服务');
    } finally {
      setLoading(false);
    }
  };

  return (
    <div>
      <div style={{ display: 'flex', gap: '1rem', alignItems: 'center', marginBottom: '1rem' }}>
        <input
          type="file"
          accept=".tif,.tiff"
          onChange={handleChange}
          className="input-file"
          id="file-upload"
        />
        <button
          onClick={handleUpload}
          className="btn"
          disabled={!preview || loading}
        >
          {loading ? '处理中...' : '上传并预测'}
        </button>
      </div>

      {loading && <p className="loading">正在处理图像，请稍候...</p>}

      <div className="result-container">
        {preview && (
          <div style={{ flex: 1 }}>
            <h3>原始图像</h3>
            <img
              src={preview}
              alt="原始图像"
              className="result-image"
              style={{ maxWidth: '100%', maxHeight: '500px' }}
            />
          </div>
        )}
        {result && (
          <div style={{ flex: 1 }}>
            <h3>分割结果</h3>
            <img
              src={result}
              alt="分割结果"
              className="result-image"
              style={{ maxWidth: '100%', maxHeight: '500px' }}
            />
          </div>
        )}
      </div>
    </div>
  );
}

export default UploadImage;

四、运行与部署

         后端我是在Pycharm上运行的，前端在Vscode。

cd backend
pip install -r requirements.txt
python app.py

       以下在vscode的终端输入， 如果你的vscode没有安装node.js，可以去网上找相应的教程安装，成功运行会弹出网页。

cd frontend
npm install
npm start

五、运行结果

        输入一张农田的图片，返还一个二分类的结果。

        这里用的是最基础的UNet网络来训练的，如果你希望分割的精度更高，可以参考最近的文献的模型结构，另外再增大数据集。我的系统做的比较简单，还在努力学习中，如果大家有需要，可以免费分享！