tem_loader.py (12937B)
1 import math 2 from pathlib import Path 3 4 from qgis.PyQt.QtCore import QMetaType 5 from qgis.PyQt.QtWidgets import ( 6 QAction, 7 QCheckBox, 8 QDialog, 9 QDialogButtonBox, 10 QFileDialog, 11 QFormLayout, 12 QMessageBox, 13 QSpinBox, 14 QVBoxLayout, 15 ) 16 from qgis.core import ( 17 Qgis, 18 QgsCoordinateReferenceSystem, 19 QgsCoordinateTransform, 20 QgsCsException, 21 QgsFeature, 22 QgsField, 23 QgsFields, 24 QgsGeometry, 25 QgsPointXY, 26 QgsProject, 27 QgsVectorFileWriter, 28 QgsVectorLayer, 29 ) 30 from qgis.gui import QgsMapLayerComboBox 31 32 from . import core 33 34 35 STYLES_DIR = Path(__file__).parent / 'styles' 36 GEOMETRY_FIELD = 'Geometry' 37 STRING_FIELDS = {'Line', 'StationNo', 'Color'} 38 INTEGER_FIELDS = {'NumLayers', 'Layer', 'Opacity'} 39 40 41 def _build_geopackage_uri(gpkg_path, layer_name): 42 return f'{gpkg_path.resolve()}|layername={layer_name}' 43 44 45 def _field_type(field_name): 46 if field_name in STRING_FIELDS: 47 return QMetaType.Type.QString 48 if field_name in INTEGER_FIELDS: 49 return QMetaType.Type.Int 50 return QMetaType.Type.Double 51 52 53 def _build_fields(rows): 54 fields = QgsFields() 55 if not rows: 56 return fields 57 for field_name in rows[0]: 58 if field_name == GEOMETRY_FIELD: 59 continue 60 fields.append(QgsField(field_name, _field_type(field_name))) 61 return fields 62 63 64 def _rows_to_features(rows, fields): 65 field_names = [field.name() for field in fields] 66 features = [] 67 for row in rows: 68 feature = QgsFeature(fields) 69 feature.setGeometry(QgsGeometry.fromWkt(row[GEOMETRY_FIELD])) 70 feature.setAttributes([row.get(name) for name in field_names]) 71 features.append(feature) 72 return features 73 74 75 def _writer_no_error(): 76 writer_error = getattr(QgsVectorFileWriter, 'WriterError', QgsVectorFileWriter) 77 return writer_error.NoError 78 79 80 def _dialog_button(name): 81 standard_button = getattr( 82 QDialogButtonBox, 83 'StandardButton', 84 QDialogButtonBox, 85 ) 86 return getattr(standard_button, name) 87 88 89 def _dialog_buttons(): 90 return _dialog_button('Ok') | _dialog_button('Cancel') 91 92 93 def _dialog_code(name): 94 dialog_code = getattr(QDialog, 'DialogCode', QDialog) 95 return getattr(dialog_code, name) 96 97 98 def _layer_filter(name): 99 layer_filter = getattr(Qgis, 'LayerFilter', Qgis) 100 return getattr(layer_filter, name) 101 102 103 def _warn_dem_sample_failed(raster_layer, x, y): 104 print( 105 'TEM Loader warning: ' 106 f'point ({x}, {y}) is outside DEM raster {raster_layer.name()}' 107 ) 108 109 110 def _sample_dem_elevation(raster_layer, source_crs, project, x, y): 111 transform = QgsCoordinateTransform(source_crs, raster_layer.crs(), project) 112 try: 113 raster_point = transform.transform(QgsPointXY(x, y)) 114 except QgsCsException as exc: 115 print( 116 'TEM Loader warning: ' 117 f'could not transform point ({x}, {y}) to DEM raster ' 118 f'{raster_layer.name()}: {exc}' 119 ) 120 return None 121 122 value, ok = raster_layer.dataProvider().sample(raster_point, 1) 123 try: 124 elevation = float(value) 125 except (TypeError, ValueError): 126 elevation = math.nan 127 if not ok or math.isnan(elevation): 128 _warn_dem_sample_failed(raster_layer, x, y) 129 return None 130 return elevation 131 132 133 def _point_wkt(x, y, z): 134 return f'POINT Z ({x} {y} {z})' 135 136 137 def _line_wkt(x, y, z_top, z_bottom): 138 return f'LINESTRING Z ({x} {y} {z_top}, {x} {y} {z_bottom})' 139 140 141 def _adjust_rows_to_dem(points, doi_points, layers, raster_layer, source_crs, 142 project): 143 if raster_layer is None: 144 return 145 146 dem_z_by_xy = {} 147 for point in points: 148 x = point['X'] 149 y = point['Y'] 150 dem_z = _sample_dem_elevation(raster_layer, source_crs, project, x, y) 151 if dem_z is None: 152 continue 153 dem_z_by_xy[(x, y)] = dem_z 154 point['Z'] = dem_z 155 point['Geometry'] = _point_wkt(x, y, dem_z) 156 157 for doi_point in doi_points: 158 x = doi_point['X'] 159 y = doi_point['Y'] 160 dem_z = dem_z_by_xy.get((x, y)) 161 if dem_z is None: 162 continue 163 z_doi = dem_z - doi_point['DOI'] 164 doi_point['Z'] = z_doi 165 doi_point['ZDOI'] = z_doi 166 doi_point['Geometry'] = _point_wkt(x, y, z_doi) 167 168 for layer in layers: 169 x = layer['X'] 170 y = layer['Y'] 171 dem_z = dem_z_by_xy.get((x, y)) 172 if dem_z is None: 173 continue 174 z_top = dem_z - layer['DepthTop'] 175 z_bottom = dem_z - layer['DepthBottom'] 176 z_mid = (z_top + z_bottom) / 2 177 layer['Z'] = dem_z 178 layer['ZTop'] = z_top 179 layer['ZMid'] = z_mid 180 layer['ZBottom'] = z_bottom 181 layer['Geometry'] = _line_wkt(x, y, z_top, z_bottom) 182 183 184 def _write_geopackage_layer(rows, gpkg_path, layer_name, wkb_type, crs, 185 transform_context, action): 186 fields = _build_fields(rows) 187 options = QgsVectorFileWriter.SaveVectorOptions() 188 options.driverName = 'GPKG' 189 options.fileEncoding = 'UTF-8' 190 options.layerName = layer_name 191 options.actionOnExistingFile = action 192 193 writer = QgsVectorFileWriter.create( 194 str(gpkg_path), 195 fields, 196 wkb_type, 197 crs, 198 transform_context, 199 options, 200 ) 201 if writer is None: 202 raise ValueError(f'failed to create GeoPackage layer {layer_name}') 203 204 try: 205 if writer.hasError() != _writer_no_error(): 206 message = writer.errorMessage() 207 details = f': {message}' if message else '' 208 raise ValueError( 209 f'failed to create GeoPackage layer {layer_name}{details}' 210 ) 211 features = _rows_to_features(rows, fields) 212 if features and not writer.addFeatures(features): 213 raise ValueError(f'failed to write GeoPackage layer {layer_name}') 214 finally: 215 del writer 216 217 218 class _ImportOptionsDialog(QDialog): 219 def __init__(self, parent=None): 220 super().__init__(parent) 221 self.setWindowTitle('TEM Loader Options') 222 223 self._mask_checkbox = QCheckBox( 224 'Mask out layers below depth of interest (DOI)' 225 ) 226 self._mask_checkbox.setChecked(True) 227 228 self._opacity_spinbox = QSpinBox() 229 self._opacity_spinbox.setRange(0, 100) 230 self._opacity_spinbox.setSuffix('%') 231 self._opacity_spinbox.setValue(core.BELOW_DOI_OPACITY) 232 self._opacity_spinbox.setEnabled(self._mask_checkbox.isChecked()) 233 self._mask_checkbox.toggled.connect(self._opacity_spinbox.setEnabled) 234 235 self._dem_checkbox = QCheckBox( 236 'Adjust vertical position to digital elevation model' 237 ) 238 self._dem_checkbox.setChecked(False) 239 240 self._dem_raster_combo = QgsMapLayerComboBox() 241 self._dem_raster_combo.setProject(QgsProject.instance()) 242 self._dem_raster_combo.setFilters(_layer_filter('RasterLayer')) 243 self._dem_raster_combo.setAllowEmptyLayer(True, 'No elevation raster') 244 self._dem_raster_combo.setEnabled(False) 245 self._dem_checkbox.toggled.connect(self._dem_raster_combo.setEnabled) 246 247 opacity_form = QFormLayout() 248 opacity_form.addRow('Opacity', self._opacity_spinbox) 249 250 dem_form = QFormLayout() 251 dem_form.addRow('Elevation raster', self._dem_raster_combo) 252 253 buttons = QDialogButtonBox(_dialog_buttons()) 254 buttons.accepted.connect(self.accept) 255 buttons.rejected.connect(self.reject) 256 257 layout = QVBoxLayout() 258 layout.addWidget(self._mask_checkbox) 259 layout.addLayout(opacity_form) 260 layout.addWidget(self._dem_checkbox) 261 layout.addLayout(dem_form) 262 layout.addWidget(buttons) 263 self.setLayout(layout) 264 265 def options(self): 266 elevation_raster_layer = None 267 if self._dem_checkbox.isChecked(): 268 elevation_raster_layer = self._dem_raster_combo.currentLayer() 269 return { 270 'mask_below_doi': self._mask_checkbox.isChecked(), 271 'below_doi_opacity': self._opacity_spinbox.value(), 272 'elevation_raster_layer': elevation_raster_layer, 273 } 274 275 276 def _exec_dialog(dialog): 277 execute = getattr(dialog, 'exec', None) 278 if execute is None: 279 execute = dialog.exec_ 280 return execute() 281 282 283 class TEMLoaderPlugin: 284 def __init__(self, iface): 285 self.iface = iface 286 self._action = None 287 288 def initGui(self): 289 self._action = QAction('Load TEM XYZ files\u2026', self.iface.mainWindow()) 290 self._action.triggered.connect(self.run) 291 self.iface.addPluginToMenu('TEM Loader', self._action) 292 293 def unload(self): 294 self.iface.removePluginMenu('TEM Loader', self._action) 295 self._action = None 296 297 def run(self): 298 paths, _ = QFileDialog.getOpenFileNames( 299 self.iface.mainWindow(), 300 'Select TEM XYZ files', 301 '', 302 'XYZ files (*.xyz);;All files (*)', 303 ) 304 if not paths: 305 return 306 307 dialog = _ImportOptionsDialog(self.iface.mainWindow()) 308 if _exec_dialog(dialog) != _dialog_code('Accepted'): 309 return 310 311 options = dialog.options() 312 failed = [] 313 for path in paths: 314 filepath = Path(path) 315 try: 316 self._load_xyz(filepath, **options) 317 except Exception as exc: 318 failed.append(f'{filepath.name}: {exc}') 319 if failed: 320 QMessageBox.warning( 321 self.iface.mainWindow(), 322 'TEM Loader', 323 '\n'.join(failed), 324 ) 325 326 def _resolve_crs(self, filepath, project): 327 crs = None 328 source_epsg = core.detect_source_epsg(filepath) 329 if source_epsg: 330 candidate = QgsCoordinateReferenceSystem() 331 if candidate.createFromString(source_epsg) and candidate.isValid(): 332 crs = candidate 333 334 if crs is None: 335 crs = project.crs() 336 if not crs.isValid(): 337 crs = QgsCoordinateReferenceSystem() 338 crs.createFromString('EPSG:4326') 339 return crs 340 341 def _load_xyz( 342 self, 343 filepath, 344 mask_below_doi=True, 345 below_doi_opacity=core.BELOW_DOI_OPACITY, 346 elevation_raster_layer=None, 347 ): 348 points, doi_points, layers = core.process_xyz( 349 filepath, 350 mask_below_doi=mask_below_doi, 351 below_doi_opacity=below_doi_opacity, 352 ) 353 354 project = QgsProject.instance() 355 crs = self._resolve_crs(filepath, project) 356 _adjust_rows_to_dem( 357 points, 358 doi_points, 359 layers, 360 elevation_raster_layer, 361 crs, 362 project, 363 ) 364 gpkg_path = filepath.with_suffix('.gpkg') 365 transform_context = project.transformContext() 366 367 output_layers = [ 368 ('layers', layers, Qgis.WkbType.LineStringZ), 369 ] 370 if doi_points: 371 output_layers.append(('doi', doi_points, Qgis.WkbType.PointZ)) 372 output_layers.append(('points', points, Qgis.WkbType.PointZ)) 373 374 written_layer_names = [] 375 first_layer = True 376 for name, rows, wkb_type in output_layers: 377 if not rows: 378 continue 379 action = ( 380 QgsVectorFileWriter.CreateOrOverwriteFile 381 if first_layer 382 else QgsVectorFileWriter.CreateOrOverwriteLayer 383 ) 384 _write_geopackage_layer( 385 rows, 386 gpkg_path, 387 name, 388 wkb_type, 389 crs, 390 transform_context, 391 action, 392 ) 393 written_layer_names.append(name) 394 first_layer = False 395 396 loaded_layers = {} 397 for name in written_layer_names: 398 uri = _build_geopackage_uri(gpkg_path, name) 399 layer = QgsVectorLayer(uri, name, 'ogr') 400 if not layer.isValid(): 401 continue 402 403 qml = STYLES_DIR / f'{name}.qml' 404 if qml.exists(): 405 layer.loadNamedStyle(str(qml)) 406 407 project.addMapLayer(layer, False) 408 loaded_layers[name] = layer 409 410 if not loaded_layers: 411 raise ValueError('failed to load any layers') 412 413 group_name = filepath.stem 414 root = project.layerTreeRoot() 415 group = root.insertGroup(0, group_name) 416 insert_index = 0 417 for name in ('points', 'doi', 'layers'): 418 layer = loaded_layers.get(name) 419 if layer is None: 420 continue 421 group.insertLayer(insert_index, layer) 422 insert_index += 1 423 424 failed = [name for name in written_layer_names if name not in loaded_layers] 425 if failed: 426 QMessageBox.warning( 427 self.iface.mainWindow(), 428 'TEM Loader', 429 f'{filepath.name}: failed to load layers: {", ".join(failed)}', 430 )