tem_loader.py (13614B)
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._clip_checkbox = QCheckBox( 224 'Hide layers below depth of interest (DOI, 2D and 3D views)' 225 ) 226 self._clip_checkbox.setChecked(True) 227 228 self._mask_checkbox = QCheckBox( 229 'Make layers below DOI partially transparent (2D view only)' 230 ) 231 self._mask_checkbox.setChecked(True) 232 233 self._opacity_spinbox = QSpinBox() 234 self._opacity_spinbox.setRange(0, 100) 235 self._opacity_spinbox.setSuffix('%') 236 self._opacity_spinbox.setValue(core.BELOW_DOI_OPACITY) 237 238 self._clip_checkbox.toggled.connect(self._update_below_doi_controls) 239 self._mask_checkbox.toggled.connect(self._update_below_doi_controls) 240 self._update_below_doi_controls() 241 242 self._dem_checkbox = QCheckBox( 243 'Adjust vertical position to digital elevation model' 244 ) 245 self._dem_checkbox.setChecked(False) 246 247 self._dem_raster_combo = QgsMapLayerComboBox() 248 self._dem_raster_combo.setProject(QgsProject.instance()) 249 self._dem_raster_combo.setFilters(_layer_filter('RasterLayer')) 250 self._dem_raster_combo.setAllowEmptyLayer(True, 'No elevation raster') 251 self._dem_raster_combo.setEnabled(False) 252 self._dem_checkbox.toggled.connect(self._dem_raster_combo.setEnabled) 253 254 opacity_form = QFormLayout() 255 opacity_form.addRow('Opacity', self._opacity_spinbox) 256 257 dem_form = QFormLayout() 258 dem_form.addRow('Elevation raster', self._dem_raster_combo) 259 260 buttons = QDialogButtonBox(_dialog_buttons()) 261 buttons.accepted.connect(self.accept) 262 buttons.rejected.connect(self.reject) 263 264 layout = QVBoxLayout() 265 layout.addWidget(self._clip_checkbox) 266 layout.addWidget(self._mask_checkbox) 267 layout.addLayout(opacity_form) 268 layout.addWidget(self._dem_checkbox) 269 layout.addLayout(dem_form) 270 layout.addWidget(buttons) 271 self.setLayout(layout) 272 273 def _update_below_doi_controls(self, *_args): 274 clipping = self._clip_checkbox.isChecked() 275 self._mask_checkbox.setEnabled(not clipping) 276 self._opacity_spinbox.setEnabled( 277 not clipping and self._mask_checkbox.isChecked() 278 ) 279 280 def options(self): 281 elevation_raster_layer = None 282 if self._dem_checkbox.isChecked(): 283 elevation_raster_layer = self._dem_raster_combo.currentLayer() 284 return { 285 'clip_below_doi': self._clip_checkbox.isChecked(), 286 'mask_below_doi': self._mask_checkbox.isChecked(), 287 'below_doi_opacity': self._opacity_spinbox.value(), 288 'elevation_raster_layer': elevation_raster_layer, 289 } 290 291 292 def _exec_dialog(dialog): 293 execute = getattr(dialog, 'exec', None) 294 if execute is None: 295 execute = dialog.exec_ 296 return execute() 297 298 299 class TEMLoaderPlugin: 300 def __init__(self, iface): 301 self.iface = iface 302 self._action = None 303 304 def initGui(self): 305 self._action = QAction('Load TEM XYZ files\u2026', self.iface.mainWindow()) 306 self._action.triggered.connect(self.run) 307 self.iface.addPluginToMenu('TEM Loader', self._action) 308 309 def unload(self): 310 self.iface.removePluginMenu('TEM Loader', self._action) 311 self._action = None 312 313 def run(self): 314 paths, _ = QFileDialog.getOpenFileNames( 315 self.iface.mainWindow(), 316 'Select TEM XYZ files', 317 '', 318 'XYZ files (*.xyz);;All files (*)', 319 ) 320 if not paths: 321 return 322 323 dialog = _ImportOptionsDialog(self.iface.mainWindow()) 324 if _exec_dialog(dialog) != _dialog_code('Accepted'): 325 return 326 327 options = dialog.options() 328 failed = [] 329 for path in paths: 330 filepath = Path(path) 331 try: 332 self._load_xyz(filepath, **options) 333 except Exception as exc: 334 failed.append(f'{filepath.name}: {exc}') 335 if failed: 336 QMessageBox.warning( 337 self.iface.mainWindow(), 338 'TEM Loader', 339 '\n'.join(failed), 340 ) 341 342 def _resolve_crs(self, filepath, project): 343 crs = None 344 source_epsg = core.detect_source_epsg(filepath) 345 if source_epsg: 346 candidate = QgsCoordinateReferenceSystem() 347 if candidate.createFromString(source_epsg) and candidate.isValid(): 348 crs = candidate 349 350 if crs is None: 351 crs = project.crs() 352 if not crs.isValid(): 353 crs = QgsCoordinateReferenceSystem() 354 crs.createFromString('EPSG:4326') 355 return crs 356 357 def _load_xyz( 358 self, 359 filepath, 360 clip_below_doi=True, 361 mask_below_doi=True, 362 below_doi_opacity=core.BELOW_DOI_OPACITY, 363 elevation_raster_layer=None, 364 ): 365 points, doi_points, layers = core.process_xyz( 366 filepath, 367 mask_below_doi=mask_below_doi, 368 below_doi_opacity=below_doi_opacity, 369 clip_below_doi=clip_below_doi, 370 ) 371 372 project = QgsProject.instance() 373 crs = self._resolve_crs(filepath, project) 374 _adjust_rows_to_dem( 375 points, 376 doi_points, 377 layers, 378 elevation_raster_layer, 379 crs, 380 project, 381 ) 382 gpkg_path = filepath.with_suffix('.gpkg') 383 transform_context = project.transformContext() 384 385 output_layers = [ 386 ('layers', layers, Qgis.WkbType.LineStringZ), 387 ] 388 if doi_points: 389 output_layers.append(('doi', doi_points, Qgis.WkbType.PointZ)) 390 output_layers.append(('points', points, Qgis.WkbType.PointZ)) 391 392 written_layer_names = [] 393 first_layer = True 394 for name, rows, wkb_type in output_layers: 395 if not rows: 396 continue 397 action = ( 398 QgsVectorFileWriter.CreateOrOverwriteFile 399 if first_layer 400 else QgsVectorFileWriter.CreateOrOverwriteLayer 401 ) 402 _write_geopackage_layer( 403 rows, 404 gpkg_path, 405 name, 406 wkb_type, 407 crs, 408 transform_context, 409 action, 410 ) 411 written_layer_names.append(name) 412 first_layer = False 413 414 loaded_layers = {} 415 for name in written_layer_names: 416 uri = _build_geopackage_uri(gpkg_path, name) 417 layer = QgsVectorLayer(uri, name, 'ogr') 418 if not layer.isValid(): 419 continue 420 421 qml = STYLES_DIR / f'{name}.qml' 422 if qml.exists(): 423 layer.loadNamedStyle(str(qml)) 424 425 project.addMapLayer(layer, False) 426 loaded_layers[name] = layer 427 428 if not loaded_layers: 429 raise ValueError('failed to load any layers') 430 431 group_name = filepath.stem 432 root = project.layerTreeRoot() 433 group = root.insertGroup(0, group_name) 434 insert_index = 0 435 for name in ('points', 'doi', 'layers'): 436 layer = loaded_layers.get(name) 437 if layer is None: 438 continue 439 group.insertLayer(insert_index, layer) 440 insert_index += 1 441 442 failed = [name for name in written_layer_names if name not in loaded_layers] 443 if failed: 444 QMessageBox.warning( 445 self.iface.mainWindow(), 446 'TEM Loader', 447 f'{filepath.name}: failed to load layers: {", ".join(failed)}', 448 )