简单使用PaddleSeg+PySide2将建筑提取模型做成GUI版
一. 模型训练
首先可以在本地使用(或者AI Studio上)pip安装PaddlePaddle-GPU版本和PaddleSeg,使用PaddleSeg提供的UNet网络,准备好(这里使用的是)Aerial imagery dataset数据集,可以很方便的进行训练(AI Studio上很多类似的项目,就不多说)。最后得到一个.pdparams
的模型参数文件即可。下面是训练的Loss和Acc等。
二、界面设计
- 首先需要在本地
pip install pyside2
,并且本地有paddle-gpu和paddleseg(没有也可以通过pip来安装)。win10下使用pyside2如果运行没有窗体产生并报错,可以参考这个。完成后到Anaconda3\Lib\site-packages\PySide2
下找到designer.exe
运行,此为QT的设计器,可以像VS一样拉出一个简单的界面。这里主要有4个控件。两个按钮分别用于加载模型和打开图片,两个标签分别用于显示图片和显示建筑提取的结果。完成界面设计后保存为.ui
文件,以便后面加载界面。
- 打开(我这里是)VS code,import界面相关的库,使用下面的代码定义一个类,为主窗体,读取设计好的ui文件,完成窗体的搭建。
from PySide2.QtWidgets import QApplication, QFileDialog, QMessageBox
from PySide2.QtGui import QImage, QPixmap # 用于显示图像的
from PySide2.QtUiTools import QUiLoader
from PySide2.QtCore import QFile
class MainForm():
def __init__(self):
# 读取加载ui文件
qtfile = QFile('ui/form.ui')
qtfile.open(QFile.ReadOnly)
qtfile.close()
self.ui = QUiLoader().load(qtfile) # 加载窗体
- 在import那里导入与paddle和paddleseg相关的库,完善属性和功能的定义,即要能加载模型参数和能打开图片进行建筑提取。
from PySide2.QtWidgets import QApplication, QFileDialog, QMessageBox
from PySide2.QtGui import QImage, QPixmap
from PySide2.QtUiTools import QUiLoader
from PySide2.QtCore import QFile
import os
import cv2
import numpy as np
import paddle
from paddleseg.models import UNet
import paddleseg.transforms as T
from paddleseg.core import infer
class MainForm():
def __init__(self):
# 读取加载ui文件
qtfile = QFile('ui/form.ui')
qtfile.open(QFile.ReadOnly)
qtfile.close()
self.ui = QUiLoader().load(qtfile) # 加载窗体
# 加载模型和图像预处理方法
self.model = UNet(num_classes=2)
self.in_params = False # 是否加载参数
self.transforms = T.Compose([T.Resize(target_size=(256, 256)), T.Normalize()])
# 点击事件
self.ui.btn_model.clicked.connect(self.load_model) # 加载模型
self.ui.btn_image.clicked.connect(self.open_image) # 加载图像
# 加载模型
def load_model(self):
# 加载图像
def open_image(self):
- 接下来分别完成两个函数,加载模型和加载图像。这里主要是使用
QFileDialog.getOpenFileName
获取文件路径,然后基操,显示图片需要用到QImage
和QPixmap
。
# 加载模型
def load_model(self):
model_path, _ = QFileDialog.getOpenFileName(
self.ui,
"选择模型参数",
os.getcwd(),
"Pdparams Files (*.pdparams)"
)
para_state_dict = paddle.load(model_path)
self.model.set_dict(para_state_dict)
self.in_params = True
QMessageBox.about(self.ui, '信息', '模型参数加载成功!')
print('Load Model Params Successful!')
# 加载图像
def open_image(self):
if self.in_params == False:
QMessageBox.warning(self.ui, '警告', '未加载模型参数,请先加载模型参数!')
return
image_path, _ = QFileDialog.getOpenFileName(
self.ui,
"选择加载图像",
os.getcwd(),
"JPG Files (*.jpg)"
)
img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
# 显示在ui上
q_img = QImage(img, img.shape[1], img.shape[0], img.strides[0], QImage.Format_RGB888)
self.ui.lab_image.setPixmap(QPixmap.fromImage(q_img))
# 预测
inf_img, _ = self.transforms(image_path)
inf_img = paddle.to_tensor(inf_img[np.newaxis, :])
pre = infer.inference(self.model, inf_img)
# 预测结果转换为二值图像
pred = paddle.argmax(pre, axis=1).numpy().reshape((256, 256)).astype('uint8') * 255
pred = cv2.resize(pred, (img.shape[0], img.shape[1]), interpolation=cv2.INTER_NEAREST) # 标签缩放
q_pred = QImage(pred, pred.shape[1], pred.shape[0], pred.strides[0], QImage.Format_Indexed8)
self.ui.lab_forecast.setPixmap(QPixmap.fromImage(q_pred))
print('Infer Successful!')
- 最后写上运行的话,显示窗体,application循环等待我们的操作。
if __name__ == '__main__':
app = QApplication([])
main_form = MainForm()
main_form.ui.show()
app.exec_()
- 到此简单的操作就结束了。
三、演示
转载:https://blog.csdn.net/qq_39798423/article/details/114695315
查看评论