从几条简单的线段生成路网

输入：随手画几条线，格式为[LineString(1...),LineString(2...).....]

输出：link、node 两个（geo）dataframe

factor = 5 #分辨率，每m几个像素 

from shapely.geometry import LineString, Point, MultiPoint, GeometryCollection, Polygon, MultiLineString
from shapely.ops import snap, split
import os
import numpy as np
from collections import Counter
from sklearn.cluster import DBSCAN 
from shapely.ops import linemerge


class UnionFindSet(object): 
    """并查集"""
    def __init__(self, data_list):
        """初始化两个字典，一个保存节点的父节点，另外一个保存父节点的大小
        初始化的时候，将节点的父节点设为自身，size设为1"""
        self.father_dict = {}
        self.size_dict = {}

        for node in data_list:
            self.father_dict[node] = node
            self.size_dict[node] = 1 

    def find(self, node):
        """使用递归的方式来查找父节点

        在查找父节点的时候，顺便把当前节点移动到父节点上面
        这个操作算是一个优化
        """
        father = self.father_dict[node]
        if(node != father):
            if father != self.father_dict[father]:    # 在降低树高优化时，确保父节点大小字典正确
                self.size_dict[father] -= 1
            father = self.find(father)
        self.father_dict[node] = father
        return father

    def is_same_set(self, node_a, node_b):
        """查看两个节点是不是在一个集合里面"""
        return self.find(node_a) == self.find(node_b)

    def union(self, node_a, node_b):
        """将两个集合合并在一起"""
        if node_a is None or node_b is None:
            return

        a_head = self.find(node_a)
        b_head = self.find(node_b)

        if(a_head != b_head):
            a_set_size = self.size_dict[a_head]
            b_set_size = self.size_dict[b_head]
            if(a_set_size >= b_set_size):
                self.father_dict[b_head] = a_head
                self.size_dict[a_head] = a_set_size + b_set_size
            else:
                self.father_dict[a_head] = b_head
                self.size_dict[b_head] = a_set_size + b_set_size

def link2link_and_node(linkgpd): 
    link_gpd = linkgpd.copy() 
    link_gpd["snode"] = link_gpd["geometry"].apply(lambda x:Point(x.coords[0])) 
    link_gpd["enode"] = link_gpd["geometry"].apply(lambda x:Point(x.coords[-1])) 
    node_set = set(link_gpd["snode"].tolist() + link_gpd["enode"].tolist()) 
    
    node_gpd = gpd.GeoDataFrame() 
    node_geom_list = list(node_set)
    node_id_list = range(len(node_set)) 
    node_gpd["geometry"] = node_geom_list
    node_gpd["nodeid"] = node_id_list
    node_gpd["node_type"] = None
    
    node2id_dict = dict(zip(node_geom_list,node_id_list)) 
    link_gpd["snodeid"] = link_gpd["snode"].apply(lambda x:node2id_dict[x]) 
    link_gpd["enodeid"] = link_gpd["enode"].apply(lambda x:node2id_dict[x]) 
    link_gpd["linkid"] = range(len(link_gpd))
    link_gpd = link_gpd[["linkid","geometry","snodeid","enodeid"]]
    return link_gpd,node_gpd
    
def check_link_and_node(link,node):
    nodeid_set = set(node["nodeid"].tolist())
    link2nodeid_set = set(link["snodeid"].tolist() + link["enodeid"].tolist())
    assert len([i for i in link2nodeid_set if i not in nodeid_set]) + len([i for i in nodeid_set if i not in link2nodeid_set]) == 0
    for _,link_row in link.iterrows():
        snodeid = link_row["snodeid"]
        enodeid = link_row["enodeid"]
        snode_geom = node[node["nodeid"] == snodeid].iloc[0]["geometry"]
        enode_geom = node[node["nodeid"] == enodeid].iloc[0]["geometry"]
        assert snode_geom.distance(Point(link_row["geometry"].coords[0])) < factor * 0.1 
        assert enode_geom.distance(Point(link_row["geometry"].coords[-1])) < factor * 0.1 
    return 

def step1_split_link(linkgpd): 
    #将交叉的link打断
    lines = linkgpd["geometry"].tolist() 
    # 查找交点并添加节点
    nodes = []
    for line in lines:
        if not line:
            continue
        for other_line in lines:
            if not other_line:
                continue
            if line != other_line:
                intersection = line.intersection(other_line.buffer(factor*0.1))
                if not intersection.is_empty:
                    if isinstance(intersection, Point):
                        nodes.append(intersection) 
                    elif isinstance(intersection, LineString):
                        nodes.append(intersection.centroid) 
    # 移除重复的节点
    unique_nodes = []
    for node in nodes:
        if node not in unique_nodes: 
            unique_nodes.append(node) 
            
    new_lines = []

    # 添加边 
    for line in lines: 
        if not line: 
            continue
        #将线几何向端点几何靠近，保证点在线上 
        snapped_line = snap(line, MultiPoint(unique_nodes), tolerance=factor*0.2)
        segments = split(snapped_line, MultiPoint(unique_nodes))

        # 检查 GeometryCollection 是否可迭代
        if isinstance(segments, GeometryCollection):
            for segment in segments.geoms:
                if isinstance(segment, LineString):
                    new_lines.append(segment) 
        elif isinstance(segments, LineString): 
            new_lines.append(segments)  
    return gpd.GeoDataFrame(geometry=new_lines) 

def step2_special_node(node_gpd,special_point_list): 
    node_gpd2 = node_gpd.copy()
    #大门附node属性 
    change_node_type_id_list = [] 
    node_goem2id_dict = dict(zip(node_gpd["geometry"],node_gpd["nodeid"])) 
    for special_point in special_point_list: 
        tmp_node = min(list(node_goem2id_dict.keys()),key = lambda x:x.distance(special_point))
        tmp_nodeid = node_goem2id_dict[tmp_node]
        if tmp_node.distance(special_point) < factor * 0.1:
            change_node_type_id_list.append(tmp_nodeid) 
    change_node_type_id_set = set(change_node_type_id_list)
    node_gpd2["node_type"] = node_gpd2["nodeid"].apply(lambda x :"door" if x in change_node_type_id_set else None) 
    return node_gpd2

def step3_merge_Knode_dbscan(link_gpd,node_gpd,merge_distance = factor * 1):
    node_gpd2 = node_gpd.copy()
    link_gpd2 = link_gpd.copy() 
    node_geom_list = node_gpd["geometry"].tolist() 
    node_id_list = node_gpd["nodeid"].tolist() 
    node_type_list = node_gpd["node_type"].tolist() 
    node_coord_list = [[i.x,i.y] for i in node_geom_list] 
    estimator = DBSCAN(eps= merge_distance,min_samples=1) # 构造聚类器
    estimator.fit(np.array(node_coord_list) ) # 聚类
    label_pred = estimator.labels_ # 获取聚类标签
    label2idx_dict = dict()
    for i in range(len(label_pred)):
        key = label_pred[i]
        if key in label2idx_dict:
            label2idx_dict[key].append(i)
        else:
            label2idx_dict[key] = [i] 
#     print(label2idx_dict) 
    #对每个待合并的点，调整node表
    old2new_nodeid = dict() 
    add_nodeid_global = max(node_gpd2["nodeid"].tolist()) 
    for va in label2idx_dict.values(): 
        if len(va) < 2: 
            continue 
        #寻找是否有待合并的特殊点 
        node_type_flag = sum([1 for i in va if node_type_list[i]=="door"]) 
        if node_type_flag == 0: 
            new_point = np.mean(np.array([node_coord_list[i] for i in va]),axis=0) 
            del_nodeid_list = [node_id_list[i] for i in va] 
            add_nodeid_global = add_nodeid_global + 1 
            old2new_nodeid.update(dict(zip(del_nodeid_list,[add_nodeid_global]*len(del_nodeid_list)))) 
            add_node_row = {
                "geometry" : Point(new_point), 
                "node_type" : None, 
                "nodeid" : add_nodeid_global, 
            } 
            node_gpd2 = node_gpd2.append(add_node_row.copy(),ignore_index=True) 
        else: 
            special_geom2id_list = [[node_geom_list[i],node_id_list[i]] for i in va if node_type_list[i]=="door"] 
            del_nodeid_list = []
            for i in va:
                if node_type_list[i]!="door": 
                    del_nodeid_list.append(node_id_list[i]) 
                    old2new_nodeid.update({node_id_list[i]:min(special_geom2id_list,key=lambda x:x[0].distance(node_geom_list[i]))[1]}) 
        for nodeid in del_nodeid_list: 
            node_gpd2 = node_gpd2[node_gpd2["nodeid"]!=nodeid] 
    nodeid2geom = dict(zip(node_gpd2["nodeid"],node_gpd2["geometry"]))
    #调整link表 
    for index,linkrow in link_gpd2.iterrows():
        snodeid = linkrow["snodeid"] 
        enodeid = linkrow["enodeid"] 
        old_geom = linkrow["geometry"] 
        new_geom,new_snode,new_enode = None, None, None
        if snodeid in old2new_nodeid: 
            link_gpd2.at[index, 'snodeid'] = old2new_nodeid[snodeid]
            new_snode = nodeid2geom[old2new_nodeid[snodeid]]
        if enodeid in old2new_nodeid:
            link_gpd2.at[index, 'enodeid'] = old2new_nodeid[enodeid]
            new_enode = nodeid2geom[old2new_nodeid[enodeid]]
        if new_snode and new_enode:
            new_geom = [[new_snode.x,new_snode.y]] + old_geom.coords[1:-1] + [[new_enode.x,new_enode.y]]
        elif new_snode:
            new_geom = [[new_snode.x,new_snode.y]] + old_geom.coords[1:]
        elif new_enode:
            new_geom = old_geom.coords[:-1] + [[new_enode.x,new_enode.y]]
        if new_geom:
            link_gpd2.at[index, 'geometry'] = LineString(new_geom) 
    link_gpd2 = link_gpd2[link_gpd2["geometry"].apply(lambda x:x.length > factor*0.1)]
    return link_gpd2,node_gpd2  

def step4_merge_2degree_node(link_gpd,node_gpd): 
    #合并非特殊点的二度node  
    nodeid2cntlinknum = Counter(link_gpd["snodeid"].tolist() + link_gpd["enodeid"].tolist()) 
    nodeid2type = dict(zip(node_gpd["nodeid"],node_gpd["node_type"])) 
    nodeid2geom = dict(zip(node_gpd["nodeid"],node_gpd["geometry"])) 
    del_nodeid_list = [] 
    for nodeid,linknum in nodeid2cntlinknum.items(): 
        if linknum==2 and nodeid2type[nodeid]!="door": 
            del_nodeid_list.append(nodeid) 
    if not len(del_nodeid_list):
        return link_gpd,node_gpd 
    link_gpd2,node_gpd2 = link_gpd.copy(),node_gpd.copy() 
    #调整node表 
    node_gpd2 = node_gpd2[node_gpd2["nodeid"].apply(lambda x:x not in set(del_nodeid_list))]
    #调整link表 
    nodeid2linkid = dict() 
    for nodeid,linkid in zip(link_gpd["enodeid"].tolist()+link_gpd["snodeid"].tolist(),link_gpd["linkid"].tolist()+link_gpd["linkid"].tolist()):
        if nodeid in nodeid2linkid:
            nodeid2linkid[nodeid].append(linkid)
        else:
            nodeid2linkid[nodeid] = [linkid] 
    wait_change_linkid_couple_list = [nodeid2linkid[i] for i in del_nodeid_list] 
    wait_change_linkid_list = list(set(sum(wait_change_linkid_couple_list, []))) 
    linkid2cntnum = Counter(list(sum(wait_change_linkid_couple_list, []))) 

#     print(wait_change_linkid_couple_list) 
    #并查集 
    UF = UnionFindSet(wait_change_linkid_list) 
    for k,v in wait_change_linkid_couple_list: 
        UF.union(k,v) 
    
    linkid2snodeid = dict(zip(link_gpd["linkid"],link_gpd["snodeid"]))
    linkid2enodeid = dict(zip(link_gpd["linkid"],link_gpd["enodeid"]))
    linkid2geom = dict(zip(link_gpd["linkid"],link_gpd["geometry"]))
    union_all = dict((i,UF.find(i)) for i in wait_change_linkid_list)
    father_list = list(set(union_all.values())) 
    add_linkid_global = max(link_gpd["linkid"].tolist()) 
    for father in father_list: 
        long_linkid_list = []
        new_node_list = [] 
        for tmp_linkid in union_all: 
            if union_all[tmp_linkid] == father: 
                long_linkid_list.append(tmp_linkid) 
                if linkid2cntnum[tmp_linkid] == 1:
                    if linkid2snodeid[tmp_linkid] not in set(del_nodeid_list):
                        new_node_list.append(linkid2snodeid[tmp_linkid])
                    if linkid2enodeid[tmp_linkid] not in set(del_nodeid_list):
                        new_node_list.append(linkid2enodeid[tmp_linkid])                    
        assert len(new_node_list) == 2 
        new_snodeid,new_enodeid = new_node_list[0], new_node_list[1] 
        add_linkid_global = add_linkid_global+1 
        # 使用 linemerge 合并 LineString
        merged_line = linemerge(MultiLineString([linkid2geom[i] for i in long_linkid_list]))
        assert isinstance(merged_line,LineString)
        if Point(merged_line.coords[0]).distance(nodeid2geom[new_snodeid]) > Point(merged_line.coords[0]).distance(nodeid2geom[new_enodeid]):
            merged_line = merged_line.reverse()
        add_link_row = { 
                "geometry" : merged_line, 
                "snodeid" :new_snodeid, 
                "enodeid" :new_enodeid, 
                "linkid" :add_linkid_global,
            } 
        link_gpd2 = link_gpd2.append(add_link_row.copy(),ignore_index=True) 
    link_gpd2 = link_gpd2[link_gpd2["linkid"].apply(lambda x:x not in set(wait_change_linkid_list))]
    return link_gpd2,node_gpd2 

def step5_drop_short_1degree_link(link_gpd,node_gpd,drop_distance = factor * 1): 
    #剔除非特殊点挂接的末端一度短小link 
    nodeid2cntlinknum = Counter(link_gpd["snodeid"].tolist() + link_gpd["enodeid"].tolist()) 
    nodeid2type = dict(zip(node_gpd["nodeid"],node_gpd["node_type"])) 
    nodeid2geom = dict(zip(node_gpd["nodeid"],node_gpd["geometry"])) 
    linkid2geom = dict(zip(link_gpd["linkid"],link_gpd["geometry"]))
    nodeid2linkid = dict() 
    for nodeid,linkid in zip(link_gpd["enodeid"].tolist()+link_gpd["snodeid"].tolist(),link_gpd["linkid"].tolist()+link_gpd["linkid"].tolist()):
        if nodeid in nodeid2linkid:
            nodeid2linkid[nodeid].append(linkid)
        else:
            nodeid2linkid[nodeid] = [linkid] 
    del_nodeid_list = [] 
    del_linkid_list = []
    for nodeid,linknum in nodeid2cntlinknum.items(): 
        if linknum == 1 and nodeid2type[nodeid] != "door" and linkid2geom[nodeid2linkid[nodeid][0]].length <= drop_distance: 
            del_nodeid_list.append(nodeid) 
            del_linkid_list.append(nodeid2linkid[nodeid][0]) 
    if not len(del_nodeid_list): 
        return link_gpd,node_gpd 
    link_gpd2,node_gpd2 = link_gpd.copy(),node_gpd.copy()  
    link_gpd2 = link_gpd2[link_gpd2["linkid"].apply(lambda x:x not in set(del_linkid_list))]
    node_gpd2 = node_gpd2[node_gpd2["nodeid"].apply(lambda x:x not in set(del_nodeid_list))]
    return link_gpd2,node_gpd2 
    
def pc(name): 
    tmp_ske = gpd.GeoDataFrame([LineString(1...),LineString(2...),.....])
    link = step1_split_link(tmp_ske) 
    link,node = link2link_and_node(link) 
    
    door_path = f"../../dataset/door/{name}.npy" 
    if os.path.exists(door_path): 
        door = np.load(door_path,allow_pickle=True)
        door = [Point(i) for i in door] 
    else: 
        door = [] 
    node = step2_special_node(node,door) 
    link,node = step3_merge_Knode_dbscan(link,node,merge_distance = factor * 1)  
    link,node = step4_merge_2degree_node(link,node) 
    link,node = step5_drop_short_1degree_link(link,node,drop_distance = factor * 3) 
    link,node = step4_merge_2degree_node(link,node) 
    check_link_and_node(link,node) 
    
    return link 

pc(1226899548807409719).to_file("tmp.geojson",driver='GeoJSON')