使用PyTorch框架来训练语义分割模型——航空影像输电线路与输电塔分割数据集，使用U-Net架构进行实现。前端界面使用Flask，UI界面使用PyQt5。电缆，铁塔木塔混凝土数据集-CSDN博客

本文链接：https://blog.csdn.net/2401_88440984/article/details/144440375

在这里插入图片描述
航空影像的输电线路与输电塔分割数据集。

该数据集包含1234张图像，超高分辨率（约3840×2160）标注了11447个对象，
属于5个不同的类别，包括电缆（cable）、铁塔（tower_lattice）、木塔（tower_wooden）以及其他两种类型：混凝土或钢结构混合塔（tower_tucohy）和空白肉眼不可见线缆（void），
共3GB 在这里插入图片描述
好的，我们将基于提供的航空影像数据集构建一个输电线路与输电塔分割系统。我们将使用PyTorch框架来训练一个语义分割模型，并使用U-Net架构进行实现。前端界面将使用Flask，而UI界面将使用PyQt5。
PyQt5 UI 应用入口

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QVBoxLayout, QWidget, QFileDialog, QMessageBox
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import Qt
import requests
import json
import base64

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        self.setWindowTitle("Aerial Imagery Segmentation System")
        self.setGeometry(100, 100, 800, 600)

        self.image_label = QLabel(self)
        self.image_label.setAlignment(Qt.AlignCenter)

        self.upload_button = QPushButton("Upload Image", self)
        self.upload_button.clicked.connect(self.upload_image)

        self.result_label = QLabel("", self)
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setStyleSheet("font-size: 16px; font-weight: bold;")

        self.segmentation_label = QLabel(self)
        self.segmentation_label.setAlignment(Qt.AlignCenter)

        layout = QVBoxLayout()
        layout.addWidget(self.image_label)
        layout.addWidget(self.upload_button)
        layout.addWidget(self.result_label)
        layout.addWidget(self.segmentation_label)

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def upload_image(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(self, "QFileDialog.getOpenFileName()", "", "Images (*.png *.xpm *.jpg);;All Files (*)", options=options)
        if file_path:
            pixmap = QPixmap(file_path)
            self.image_label.setPixmap(pixmap.scaled(400, 400, Qt.KeepAspectRatio))

            files = {'file': open(file_path, 'rb')}
            response = requests.post('http://127.0.0.1:5000/predict', files=files)
            result = response.json()

            pred_base64 = result['prediction']
            pred_bytes = base64.b64decode(pred_base64)
            pred_pixmap = QPixmap()
            pred_pixmap.loadFromData(pred_bytes)

            self.segmentation_label.setPixmap(pred_pixmap.scaled(400, 400, Qt.KeepAspectRatio))

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())

项目结构

aerial_imagery_segmentation/
├── data/
│   ├─`在这里插入代码片`─ images/
│       ├── img_00001.png
│       ├── img_00002.png
│       └── ...
│   └── masks/
│       ├── img_00001_mask.png
│       ├── img_00002_mask.png
│       └── ...
├── models/
│   ├── unet.pth
├── app.py          # Flask应用入口
├── ui_app.py       # PyQt5 UI应用入口
├── train_unet.py   # U-Net模型训练脚本
├── utils.py        # 工具函数
├── requirements.txt  # 依赖项
└── README.md       # 项目说明文档

安装依赖项

首先，安装所需的Python库：

pip install torch torchvision flask pyqt5 matplotlib scikit-learn pandas seaborn albumentations opencv-python-headless

数据集准备

假设你已经有一个航空影像分割数据集，目录结构如上所示。如果没有，可以从公开的数据集中下载或自行收集。

训练模型

U-Net 模型训练

前端界面 (Flask)

`app.py` 文件

`templates/index.html` 文件

创建一个简单的HTML模板用于上传图片并显示预测结果：

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Aerial Imagery Segmentation</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            display: flex;
            justify-content: center;
            align-items: center;
            height: 100vh;
            background-color: #f4f4f4;
        }
        .container {
            text-align: center;
            background-color: white;
            padding: 20px;
            border-radius: 8px;
            box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
        }
        input[type="file"] {
            margin-bottom: 20px;
        }
        button {
            padding: 10px 20px;
            background-color: #007bff;
            color: white;
            border: none;
            border-radius: 5px;
            cursor: pointer;
        }
        button:hover {
            background-color: #0056b3;
        }
        .result {
            margin-top: 20px;
            font-size: 1.2em;
        }
        .segmentation-result {
            margin-top: 20px;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>Upload Aerial Image</h1>
        <form id="upload-form" enctype="multipart/form-data">
            <input type="file" id="file-input" name="file" accept="image/*" required>
            <br><br>
            <button type="submit">Predict</button>
        </form>
        <div class="segmentation-result" id="segmentation-result"></div>
    </div>

    <script>
        document.getElementById('upload-form').addEventListener('submit', function(event) {
            event.preventDefault();
            const formData = new FormData(this);

            fetch('/predict', {
                method: 'POST',
                body: formData
            })
            .then(response => response.json())
            .then(data => {
                const imgElement = document.createElement('img');
                imgElement.src = 'data:image/png;base64,' + data.prediction;
                document.getElementById('segmentation-result').innerHTML = '';
                document.getElementById('segmentation-result').appendChild(imgElement);
            })
            .catch(error => console.error('Error:', error));
        });
    </script>
</body>
</html>