geospatial-scripts/draw_rail_lines.py at master · yan-fab/geospatial-scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import geopandas as gpd
from shapely.geometry import LineString
import networkx as nx
from scipy.spatial.distance import pdist, squareform
import os

# Corrige warning do PROJ e GDAL no QGIS Python
os.environ['PROJ_DATA'] = r'C:\Program Files\QGIS 3.44.7\share\proj'
os.environ['GDAL_DATA'] = r'C:\Program Files\QGIS 3.44.7\share\gdal'

def process_rail_lines(input_shp, output_gpkg, layer_name):
    print(f"Processando {layer_name}...")
    try:
        gdf_pts = gpd.read_file(input_shp)
    except Exception as e:
        print(f"Erro ao ler {input_shp}: {e}")
        return

    linha_col = 'nm_linha_m'
    if linha_col not in gdf_pts.columns:
        print(f"Coluna {linha_col} não encontrada em {input_shp}.")
        return

    gdf_pts = gdf_pts.dropna(subset=['geometry'])
    lines_features = []

    for linha_nome, group in gdf_pts.groupby(linha_col):
        pts = list(group.geometry)
        if len(pts) < 2:
            continue

        # Extrair coordenadas
        coords = [(pt.x, pt.y) for pt in pts]

        # Matriz de distancias e criacao do grafo
        dist_matrix = squareform(pdist(coords))
        G = nx.Graph()
        for i in range(len(pts)):
            for j in range(i + 1, len(pts)):
                G.add_edge(i, j, weight=dist_matrix[i][j])

        # Arvore Geradora Minima (Kruskal/Prim)
        mst = nx.minimum_spanning_tree(G)

        # Desenhar os segmentos
        for u, v, data in mst.edges(data=True):
            line = LineString([pts[u], pts[v]])
            props = {
                'nm_linha_m': linha_nome,
                'geometry': line
            }
            if 'nm_empresa' in group.columns:
                props['nm_empresa'] = group.iloc[0]['nm_empresa']

            lines_features.append(props)

    if len(lines_features) > 0:
        gdf_lines = gpd.GeoDataFrame(lines_features, crs=gdf_pts.crs)
        try:
            gdf_lines.to_file(output_gpkg, driver="GPKG", layer=layer_name)
            print(f"Salvo {len(gdf_lines)} segmentos de linha em {output_gpkg}")
        except Exception as e:
            print(f"Erro ao salvar GeoPackage: {e}")
    else:
        print(f"Nenhum segmento gerado para {layer_name}")

metro_shp = r"C:\Users\yanju\Downloads\QGIS_SP_Layers\estacao_metro\estacao_metro.shp"
trem_shp = r"C:\Users\yanju\Downloads\QGIS_SP_Layers\estacao_trem\estacao_trem.shp"

metro_out = r"C:\Users\yanju\.gemini\antigravity\scratch\geo_workspace\camada_linhas_metro.gpkg"
trem_out = r"C:\Users\yanju\.gemini\antigravity\scratch\geo_workspace\camada_linhas_trem.gpkg"

process_rail_lines(metro_shp, metro_out, "linhas_metro")
process_rail_lines(trem_shp, trem_out, "linhas_trem")
print("Processamento total finalizado!")