Skip to content

Commit 7e6652a

Browse files
committed
fix: adding io module
1 parent 393323e commit 7e6652a

5 files changed

Lines changed: 201 additions & 0 deletions

File tree

loopstructural/main/io/__init__.py

Whitespace-only changes.

loopstructural/main/io/pyvista2qgis.py

Whitespace-only changes.
Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import pandas as pd
2+
import geopandas as gpd
3+
from qgis.core import QgsRaster, QgsWkbTypes
4+
5+
6+
def qgsLayerToGeoDataFrame(layer) -> gpd.GeoDataFrame:
7+
if layer is None:
8+
return None
9+
features = layer.getFeatures()
10+
fields = layer.fields()
11+
data = {'geometry': []}
12+
for f in fields:
13+
data[f.name()] = []
14+
for feature in features:
15+
geom = feature.geometry()
16+
if geom.isEmpty():
17+
continue
18+
data['geometry'].append(geom)
19+
for f in fields:
20+
data[f.name()].append(feature[f.name()])
21+
return gpd.GeoDataFrame(data, crs=layer.crs().authid())
22+
23+
24+
def qgsLayerToDataFrame(layer, dtm) -> pd.DataFrame:
25+
"""Convert a vector layer to a pandas DataFrame
26+
samples the geometry using either points or the vertices of the lines
27+
28+
:param layer: _description_
29+
:type layer: _type_
30+
:param dtm: Digital Terrain Model to evaluate Z values
31+
:type dtm: _type_ or None
32+
:return: the dataframe object
33+
:rtype: pd.DataFrame
34+
"""
35+
if layer is None:
36+
return None
37+
fields = layer.fields()
38+
data = {}
39+
data['X'] = []
40+
data['Y'] = []
41+
data['Z'] = []
42+
43+
for field in fields:
44+
data[field.name()] = []
45+
for feature in layer.getFeatures():
46+
geom = feature.geometry()
47+
points = []
48+
if geom.isMultipart():
49+
if geom.type() == QgsWkbTypes.PointGeometry:
50+
points = geom.asMultiPoint()
51+
elif geom.type() == QgsWkbTypes.LineGeometry:
52+
for line in geom.asMultiPolyline():
53+
points.extend(line)
54+
# points = geom.asMultiPolyline()[0]
55+
else:
56+
if geom.type() == QgsWkbTypes.PointGeometry:
57+
points = [geom.asPoint()]
58+
elif geom.type() == QgsWkbTypes.LineGeometry:
59+
points = geom.asPolyline()
60+
61+
for p in points:
62+
coords = list(p.coords())
63+
64+
data['X'].append(coords[0])
65+
data['Y'].append(coords[1])
66+
if dtm is not None:
67+
# if a DTM is provided, use it to get the Z value
68+
# Extract the value at the point
69+
z_value = dtm.dataProvider().identify(p, QgsRaster.IdentifyFormatValue)
70+
if z_value.isValid():
71+
z_value = z_value.results()[1]
72+
else:
73+
z_value = -9999
74+
data['Z'].append(z_value)
75+
elif len(coords) > 2:
76+
# if the geometry has Z values, use them
77+
data['Z'].append(coords[2])
78+
else:
79+
#otherwise set to 0
80+
data['Z'].append(0)
81+
for field in fields:
82+
data[field.name()].append(feature[field.name()])
83+
return pd.DataFrame(data)
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
def polyline2line(geodataframe):
2+
pass
3+
4+
def gqsPoint2VtkPoint(geodataframe):
5+
pass
6+

loopstructural/main/io/samplers.py

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
import pandas as pd
2+
import geopandas as gpd
3+
from typing import Callable
4+
from abc import ABC, abstractmethod
5+
class BaseSampler(ABC):
6+
"""Base class for samplers."""
7+
8+
@abstractmethod
9+
def __call__(self, line: gpd.GeoDataFrame, dem: Callable, use_z: bool) -> pd.DataFrame:
10+
"""Sample the line and return a DataFrame with X, Y, Z coordinates and attributes."""
11+
pass
12+
class RegularSpacingSampler(BaseSampler):
13+
def __init__(self, spacing: float):
14+
self.spacing = spacing
15+
def __call__(self, line: gpd.GeoDataFrame, dem: Callable, use_z: bool) -> pd.DataFrame:
16+
"""Sample the line at regular intervals and return a DataFrame with X, Y, Z coordinates and attributes."""
17+
points = []
18+
feature_id = 0
19+
if line is None:
20+
return pd.DataFrame(points, columns=['X', 'Y', 'Z', 'feature_id'])
21+
for geom in line.geometry:
22+
attributes = line.iloc[feature_id].to_dict()
23+
attributes.pop('geometry', None) # Remove geometry from attributes
24+
if geom.geom_type == 'LineString':
25+
length = geom.length
26+
num_points = max(int(length // self.spacing), 1)
27+
for i in range(num_points + 1):
28+
point = geom.interpolate(i * self.spacing)
29+
x, y = point.x, point.y
30+
# Use Z from geometry if available, otherwise use DEM
31+
if use_z and hasattr(point, 'z'):
32+
z = point.z
33+
else:
34+
z = dem(x, y)
35+
points.append({'X': x, 'Y': y, 'Z': z, 'feature_id': feature_id, **attributes})
36+
elif geom.geom_type == 'MultiLineString':
37+
for l in geom.geoms:
38+
length = l.length
39+
num_points = max(int(length // self.spacing), 1)
40+
for i in range(num_points + 1):
41+
point = l.interpolate(i * self.spacing)
42+
x, y = point.x, point.y
43+
# Use Z from geometry if available, otherwise use DEM
44+
if use_z and hasattr(point, 'z'):
45+
z = point.z
46+
else:
47+
z = dem(x, y)
48+
points.append({'X': x, 'Y': y, 'Z': z, 'feature_id': feature_id, **attributes})
49+
elif geom.geom_type == 'Point':
50+
coords = list(geom.coords[0])
51+
# Use Z from geometry if available, otherwise use DEM
52+
if use_z and len(coords) > 2:
53+
z = coords[2]
54+
elif dem is not None:
55+
z = dem(coords[0], coords[1])
56+
else:
57+
z = 0
58+
points.append({'X': coords[0], 'Y': coords[1], 'Z': z, 'feature_id': feature_id, **attributes})
59+
feature_id += 1
60+
df = pd.DataFrame(points)
61+
return df
62+
class AllSampler(BaseSampler):
63+
"""This is a simple sampler that just returns all the points, or all of the vertices
64+
of a line. It will also copy the elevation from the DEM or the elevation set in the data manager.
65+
"""
66+
67+
def __call__(self, line: gpd.GeoDataFrame, dem: Callable, use_z: bool) -> pd.DataFrame:
68+
"""Sample the line and return a DataFrame with X, Y, Z coordinates and attributes."""
69+
points = []
70+
feature_id = 0
71+
if line is None:
72+
return pd.DataFrame(points, columns=['X', 'Y', 'Z', 'feature_id'])
73+
for geom in line.geometry:
74+
attributes = line.iloc[feature_id].to_dict()
75+
attributes.pop('geometry', None) # Remove geometry from attributes
76+
if geom.geom_type == 'LineString':
77+
coords = list(geom.coords)
78+
for coord in coords:
79+
x, y = coord[0], coord[1]
80+
# Use Z from geometry if available, otherwise use DEM
81+
if use_z and len(coord) > 2:
82+
z = coord[2]
83+
else:
84+
z = dem(x, y)
85+
points.append({'X': x, 'Y': y, 'Z': z, 'feature_id': feature_id, **attributes})
86+
elif geom.geom_type == 'MultiLineString':
87+
for l in geom.geoms:
88+
coords = list(l.coords)
89+
for coord in coords:
90+
x, y = coord[0], coord[1]
91+
# Use Z from geometry if available, otherwise use DEM
92+
if use_z and len(coord) > 2:
93+
z = coord[2]
94+
else:
95+
z = dem(x, y)
96+
points.append(
97+
{'X': x, 'Y': y, 'Z': z, 'feature_id': feature_id, **attributes}
98+
)
99+
elif geom.geom_type == 'Point':
100+
101+
coords = list(geom.coords[0])
102+
# Use Z from geometry if available, otherwise use DEM
103+
if use_z and len(coords) > 2:
104+
z = coords[2]
105+
elif dem is not None:
106+
z = dem(coords[0], coords[1])
107+
else:
108+
z = 0
109+
points.append({'X': coords[0], 'Y': coords[1], 'Z': z, 'feature_id': feature_id, **attributes})
110+
feature_id += 1
111+
df = pd.DataFrame(points)
112+
return df

0 commit comments

Comments
 (0)