2023-04-17 13:05:10 +08:00

198 lines
7.1 KiB
Python

import shutil
from rscder.gui.actions import ActionManager
from rscder.utils.icons import IconInstance
from rscder.utils.project import Project, PairLayer, ResultPointLayer
from rscder.plugins.basic import BasicPlugin
from PyQt5.QtWidgets import QDialog, QHBoxLayout, QFileDialog, QComboBox, QVBoxLayout, QPushButton, QLabel, QLineEdit, QAction
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import Qt
from osgeo import gdal, gdal_array
import numpy as np
class ExportDialog(QDialog):
TXT='*.txt'
TIF='*.tif'
def __init__(self, parent=None, out_type = TXT):
super().__init__(parent)
self.setWindowTitle('Export')
self.setWindowIcon(IconInstance().LOGO)
self.out_path = None
self.result_layer = None
result_layer_select_label = QLabel('选择结果:')
result_layer_select = QComboBox(self)
result_layer_select.addItem('---', None)
for layer in Project().layers.values():
for result_layer in layer.layers:
if isinstance(result_layer, ResultPointLayer):
result_layer_select.addItem(IconInstance().VECTOR, result_layer.name, result_layer)
def on_result_layer_select(index):
self.result_layer = result_layer_select.currentData()
result_layer_select.currentIndexChanged.connect(on_result_layer_select)
out_path_label = QLabel('输出路径:')
out_path_text = QLineEdit(self)
out_path_text.setReadOnly(True)
out_path_text.setPlaceholderText('选择输出路径')
def on_out_path_btn():
select_file = QFileDialog.getSaveFileName(self, '选择输出路径', '', out_type)
if select_file[0]:
out_path_text.setText(select_file[0])
self.out_path = select_file[0]
out_path_btn = QPushButton('...', self)
out_path_btn.clicked.connect(on_out_path_btn)
ok_btn = QPushButton('确定', self)
ok_btn.clicked.connect(self.accept)
cancel_btn = QPushButton('取消', self)
cancel_btn.clicked.connect(self.reject)
hbox1 = QHBoxLayout()
hbox1.addWidget(result_layer_select_label)
hbox1.addWidget(result_layer_select)
hbox2 = QHBoxLayout()
hbox2.addWidget(out_path_label)
hbox2.addWidget(out_path_text)
hbox2.addWidget(out_path_btn)
hbox3 = QHBoxLayout()
hbox3.addWidget(ok_btn,0,alignment=Qt.AlignHCenter)
hbox3.addWidget(cancel_btn,0,alignment=Qt.AlignHCenter)
vbox = QVBoxLayout()
vbox.addLayout(hbox1)
vbox.addLayout(hbox2)
vbox.addLayout(hbox3)
self.setLayout(vbox)
class ExportPlugin(BasicPlugin):
@staticmethod
def info():
return {
'name': 'Export',
'description': 'Export to other format',
'author': 'RSCDER',
}
def set_action(self):
self.export_txt = QAction(IconInstance().DOCUMENT, '导出为 Arcgis 兼容的TXT', self.mainwindow)
self.export_txt.triggered.connect(self.export_txt_action)
self.export_bin = QAction(IconInstance().DOCUMENT, '导出栅格二值变化检测结果')
self.export_bin.triggered.connect(self.export_bin_action)
self.export_rgb = QAction(IconInstance().DOCUMENT, '导出变化幅度渲染结果')
self.export_rgb.triggered.connect(self.export_rgb_action)
ActionManager().export_menu.addAction(self.export_txt)
ActionManager().export_menu.addAction(self.export_bin)
ActionManager().export_menu.addAction(self.export_rgb)
# self.ctx['toolbar'].addAction(self.export_txt)
def export_bin_action(self):
dialog = ExportDialog(self.mainwindow, ExportDialog.TIF)
if dialog.exec_():
result:ResultPointLayer = dialog.result_layer
out = dialog.out_path
if result:
shutil.copy(result.result_path['path'], out)
self.message_box.info('导出成功')
def export_txt_action(self):
dialog = ExportDialog(self.mainwindow, ExportDialog.TXT)
if dialog.exec_():
result = dialog.result_layer
out = dialog.out_path
if result:
shutil.copy(result.path, out)
self.message_box.info('导出成功')
def export_rgb_action(self):
dialog = ExportDialog(self.mainwindow, ExportDialog.TIF)
btn = QPushButton('选择样式文件')
style_path = ''
def select_style():
select_file = QFileDialog.getOpenFileName(self, '选择样式文件', '*.*')
if select_file[0]:
# style_path.setText(select_file[0])
style_path = select_file[0]
btn.setText(select_file[0])
btn.clicked.connect(select_style)
dialog.layout().addWidget(btn)
default_style = [
[0, 0, 255, 0],
[1, 255, 0, 0]
]
if style_path is '':
style = default_style
else:
style = np.loadtxt(style_path, comments='#', delimiter=',')
style = style.tolist()
if dialog.exec_():
result = dialog.result_layer.path
out = dialog.out_path
self.render(style, result, out)
def render(self, style:list, path, out):
data = gdal_array.LoadFile(path)
if len(data.shape) == 3:
data = data[..., 0]
def get_color(v):
first_color = []
second_color = []
first_value = 0
second_value = 1
for s in style:
if s[0] <= v:
first_value = s[0]
first_color = s[1:]
else:
second_value = s[0]
second_color = s[1:]
break
if second_value == -1:
return np.array(style[-1][1:])
first_dis = (v - first_value) / (second_value - first_value)
second_dis = (second_value - v) / (second_value - first_value)
first_color = np.array(first_color)
second_color = np.array(second_color)
color = first_color* first_dis + second_color * second_dis
return color
get_color = np.frompyfunc(get_color, nin=1, nout=1)
Y, X = data.shape
rgbs = get_color(data.reshape(-1,))
rgbs = np.stack(rgbs).reshape((Y, X, 3)).astype(np.uint8)
driver = gdal.GetDriverByName('GTiff')
# out_tif = os.path.join(Project().other_path, 'temp.tif')
out_ds = driver.Create(out, X, Y, 3, gdal.GDT_Byte)
ds = gdal.Open(path)
geo = ds.GetGeoTransform()
proj = ds.GetProjectionRef()
out_ds.SetGeoTransform(geo)
out_ds.SetProjection(proj)
for i in range(3):
out_ds.GetRasterBand(i+1).WriteArray(rgbs[..., i])
del out_ds