pyhld/qgis_scripts/create_nodes.py
2024-04-19 14:29:59 -05:00

78 lines
2.6 KiB
Python

import itertools
from qgis.PyQt.QtCore import QVariant
from qgis.core import (QgsFeature, QgsField, QgsGeometry, QgsPointXY,
QgsVectorLayer, QgsProject, QgsVectorDataProvider)
print("Creating nodes...")
# Load the EDGES layer
edges_layer = QgsProject.instance().mapLayersByName('EDGES')[0]
# Prepare the NODES layer
nodes_layer = QgsVectorLayer("Point", "NODES", "memory")
prov = nodes_layer.dataProvider()
# Add 'id' field to the NODES layer
prov.addAttributes([QgsField("id", QVariant.Int)])
nodes_layer.updateFields()
# Prepare the start_node and end_node fields for the EDGES layer
edges_layer.startEditing()
edges_prov = edges_layer.dataProvider()
edges_prov.addAttributes([QgsField("start_node", QVariant.Int),
QgsField("end_node", QVariant.Int)])
edges_layer.updateFields()
# To store the unique nodes
unique_nodes = []
id_counter = itertools.count(start=1)
node_id_map = {}
total_features = edges_layer.featureCount()
count = 0
for feature in edges_layer.getFeatures():
count += 1
print(f'Processing feature {count} of {total_features} ({(count/total_features)*100:.2f}%)')
# Get the start and end points
start_point = feature.geometry().asPolyline()[0]
end_point = feature.geometry().asPolyline()[-1]
# Check if nodes are unique and add them to NODES layer
if start_point not in unique_nodes:
unique_nodes.append(start_point)
node_id = next(id_counter)
node_feature = QgsFeature()
node_feature.setGeometry(QgsGeometry.fromPointXY(start_point))
node_feature.setAttributes([node_id])
prov.addFeature(node_feature)
node_id_map[start_point] = node_id
if end_point not in unique_nodes:
unique_nodes.append(end_point)
node_id = next(id_counter)
node_feature = QgsFeature()
node_feature.setGeometry(QgsGeometry.fromPointXY(end_point))
node_feature.setAttributes([node_id])
prov.addFeature(node_feature)
node_id_map[end_point] = node_id
# Populate the start_node and end_node fields for the EDGES layer
edges_layer.changeAttributeValue(feature.id(),
edges_layer.fields().indexOf("start_node"),
node_id_map[start_point])
edges_layer.changeAttributeValue(feature.id(),
edges_layer.fields().indexOf("end_node"),
node_id_map[end_point])
# Stop editing the EDGES layer and save changes
edges_layer.commitChanges()
# Add the layer to the Layers panel
QgsProject.instance().addMapLayer(nodes_layer)
print("Done.")