@@ -0,0 +1,310 @@
|
||||
import warnings
|
||||
|
||||
import geopandas as gpd
|
||||
import pandas as pd
|
||||
from shapely.geometry import LineString, MultiLineString, Point
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
EXTENDED_GATE_MATCH_EXTRA = 0.00003
|
||||
|
||||
|
||||
def _norm_text(value):
|
||||
if value is None or pd.isna(value):
|
||||
return ""
|
||||
text = str(value).strip()
|
||||
if text.lower() in {"", "nan", "none", "null"}:
|
||||
return ""
|
||||
return text
|
||||
|
||||
|
||||
def _flatten_line_geometry(geom):
|
||||
if geom is None or not geom.is_valid:
|
||||
return None
|
||||
if geom.geom_type == "LineString":
|
||||
return geom if len(geom.coords) >= 2 else None
|
||||
if geom.geom_type == "MultiLineString":
|
||||
coords = []
|
||||
for part in geom.geoms:
|
||||
if len(part.coords) >= 2:
|
||||
coords.extend(list(part.coords))
|
||||
return LineString(coords) if len(coords) >= 2 else None
|
||||
return None
|
||||
|
||||
|
||||
def filter_gate_from_pointshp(point_gdf):
|
||||
gate_gdf = point_gdf.copy()
|
||||
gate_key_cn = {"进水闸", "进排水闸", "节制闸"}
|
||||
gate_key_en = {"inlet-G", "inout-G", "main-G", "branch-G"}
|
||||
|
||||
cond1 = gate_gdf["类型"].apply(lambda x: _norm_text(x) in gate_key_cn) if "类型" in gate_gdf.columns else False
|
||||
cond2 = gate_gdf["type"].apply(lambda x: _norm_text(x) in gate_key_en) if "type" in gate_gdf.columns else False
|
||||
|
||||
if "类型" not in gate_gdf.columns and "type" not in gate_gdf.columns:
|
||||
raise ValueError("Point data is missing the type/类型 field.")
|
||||
|
||||
gate_gdf = gate_gdf[(cond1 | cond2)].reset_index(drop=True)
|
||||
print(f"筛选出 {len(gate_gdf)} 个闸门")
|
||||
return gate_gdf
|
||||
|
||||
|
||||
def gate_type_complement(row):
|
||||
type_cn = _norm_text(row.get("类型", ""))
|
||||
type_en = _norm_text(row.get("type", ""))
|
||||
channel_type = _norm_text(row.get("channel_type", ""))
|
||||
|
||||
cn2en = {
|
||||
"进水闸": "inlet-G",
|
||||
"进排水闸": "inout-G",
|
||||
"排水闸": "drain-G",
|
||||
"泵": "pump",
|
||||
"节制闸": "",
|
||||
}
|
||||
en2cn = {
|
||||
"inlet-G": "进水闸",
|
||||
"inout-G": "进排水闸",
|
||||
"drain-G": "排水闸",
|
||||
"pump": "泵",
|
||||
"main-G": "节制闸",
|
||||
"branch-G": "节制闸",
|
||||
}
|
||||
|
||||
if type_cn in cn2en:
|
||||
if type_cn == "节制闸" and channel_type:
|
||||
type_en = "main-G" if channel_type == "main-S" else "branch-G"
|
||||
else:
|
||||
type_en = cn2en[type_cn]
|
||||
elif type_en in en2cn:
|
||||
type_cn = en2cn[type_en]
|
||||
else:
|
||||
type_cn = "未知闸门"
|
||||
type_en = "unknown-G"
|
||||
|
||||
return type_cn, type_en
|
||||
|
||||
|
||||
def cal_gate_proj_on_channel(gate_point, channel_geom, tolerance=0.00010):
|
||||
flat_geom = _flatten_line_geometry(channel_geom)
|
||||
if flat_geom is None:
|
||||
return gate_point, 0.0
|
||||
|
||||
along_dist = flat_geom.project(gate_point)
|
||||
proj_point = flat_geom.interpolate(along_dist)
|
||||
return proj_point, along_dist
|
||||
|
||||
|
||||
def get_gate_belong_channel(gate_point, channel_gdf, max_match_distance=0.00010, tolerance=0.00010):
|
||||
valid_channels = channel_gdf.copy()
|
||||
if "type" in valid_channels.columns:
|
||||
valid_channels = valid_channels[valid_channels["type"].isin(["main-S", "branch-S"])].copy()
|
||||
|
||||
min_dist = float("inf")
|
||||
belong_code, belong_ctype = "", ""
|
||||
proj_point = gate_point
|
||||
along_dist = 0.0
|
||||
|
||||
for _, ch in valid_channels.iterrows():
|
||||
dist = gate_point.distance(ch.geometry)
|
||||
if dist < min_dist and dist < max_match_distance:
|
||||
min_dist = dist
|
||||
belong_code = ch.code
|
||||
belong_ctype = ch.type
|
||||
proj_point, along_dist = cal_gate_proj_on_channel(gate_point, ch.geometry, tolerance)
|
||||
|
||||
match_dist = min_dist if belong_code else None
|
||||
return belong_code, belong_ctype, proj_point, along_dist, match_dist
|
||||
|
||||
|
||||
def get_channel_start_point(channel_geom):
|
||||
flat_geom = _flatten_line_geometry(channel_geom)
|
||||
if flat_geom is None:
|
||||
return None
|
||||
return Point(flat_geom.coords[0])
|
||||
|
||||
|
||||
def get_sub_channels_sorted(chan_code, channel_gdf, tolerance=0.00010):
|
||||
parent_rows = channel_gdf[channel_gdf["code"] == chan_code]
|
||||
if len(parent_rows) == 0:
|
||||
return []
|
||||
|
||||
parent_geom = parent_rows.iloc[0].geometry
|
||||
subs = channel_gdf[channel_gdf["feed_by"] == chan_code]["code"].tolist()
|
||||
|
||||
def dist_to_parent(sub_code):
|
||||
sub_rows = channel_gdf[channel_gdf["code"] == sub_code]
|
||||
if len(sub_rows) == 0:
|
||||
return 0.0
|
||||
start_pt = get_channel_start_point(sub_rows.iloc[0].geometry)
|
||||
if start_pt is None:
|
||||
return 0.0
|
||||
_, dist = cal_gate_proj_on_channel(start_pt, parent_geom, tolerance)
|
||||
return dist
|
||||
|
||||
return sorted(subs, key=dist_to_parent, reverse=True)
|
||||
|
||||
|
||||
def dfs_traverse(chan_code, channel_gdf, gate_gdf, result_gates, tolerance=0.00010):
|
||||
gates = gate_gdf[gate_gdf["channel_code"] == chan_code].copy()
|
||||
parent_rows = channel_gdf[channel_gdf["code"] == chan_code]
|
||||
if len(parent_rows) == 0:
|
||||
return
|
||||
parent_geom = parent_rows.iloc[0].geometry
|
||||
|
||||
subs_sorted = get_sub_channels_sorted(chan_code, channel_gdf, tolerance)
|
||||
virtual_points = []
|
||||
for sub in subs_sorted:
|
||||
sub_rows = channel_gdf[channel_gdf["code"] == sub]
|
||||
if len(sub_rows) == 0:
|
||||
continue
|
||||
sub_geom = get_channel_start_point(sub_rows.iloc[0].geometry)
|
||||
if sub_geom is None:
|
||||
continue
|
||||
_, dist = cal_gate_proj_on_channel(sub_geom, parent_geom, tolerance)
|
||||
virtual_points.append({"along_dist": dist, "is_sub": True, "sub_code": sub})
|
||||
|
||||
gates_list = []
|
||||
for _, gate in gates.iterrows():
|
||||
gates_list.append({"gate": gate, "along_dist": gate.along_dist, "is_sub": False})
|
||||
for item in virtual_points:
|
||||
gates_list.append({"gate": None, "along_dist": item["along_dist"], "is_sub": True, "sub_code": item["sub_code"]})
|
||||
|
||||
gates_list = sorted(gates_list, key=lambda x: x["along_dist"], reverse=True)
|
||||
|
||||
for item in gates_list:
|
||||
if item["is_sub"]:
|
||||
dfs_traverse(item["sub_code"], channel_gdf, gate_gdf, result_gates, tolerance)
|
||||
else:
|
||||
result_gates.append(item["gate"])
|
||||
|
||||
|
||||
def sort_gates_by_dfs(gate_gdf, channel_gdf, tolerance=0.00010):
|
||||
main_channels = channel_gdf[channel_gdf["type"] == "main-S"]["code"].tolist()
|
||||
main_channels = sorted(main_channels, key=lambda code: int(str(code)[1:]))
|
||||
result_gates = []
|
||||
|
||||
for chan_code in main_channels:
|
||||
dfs_traverse(chan_code, channel_gdf, gate_gdf, result_gates, tolerance)
|
||||
|
||||
return pd.DataFrame(result_gates)
|
||||
|
||||
|
||||
def _gate_log_name(row, fallback):
|
||||
for col in ("code", "名称", "name", "编号", "id", "ID"):
|
||||
if col in row.index:
|
||||
text = _norm_text(row[col])
|
||||
if text:
|
||||
return text
|
||||
return fallback
|
||||
|
||||
|
||||
def assign_gate_code_main(point_shp_path, channel_gdf, max_match_distance=0.00010, extended_match_extra=EXTENDED_GATE_MATCH_EXTRA):
|
||||
print("1. 读取数据...")
|
||||
point_gdf = gpd.read_file(point_shp_path, encoding="utf-8")
|
||||
|
||||
gate_gdf = filter_gate_from_pointshp(point_gdf)
|
||||
if len(gate_gdf) == 0:
|
||||
print("[WARN] 未找到闸门点，返回空结果")
|
||||
empty = gpd.GeoDataFrame(columns=["geometry", "code", "类型", "type"], geometry="geometry", crs=point_gdf.crs)
|
||||
empty.attrs["detected_count"] = 0
|
||||
empty.attrs["matched_count"] = 0
|
||||
empty.attrs["extended_matched_count"] = 0
|
||||
empty.attrs["unmatched_count"] = 0
|
||||
return empty
|
||||
|
||||
gate_gdf["code"] = ""
|
||||
gate_gdf["channel_code"] = ""
|
||||
gate_gdf["channel_type"] = ""
|
||||
gate_gdf["along_dist"] = 0.0
|
||||
|
||||
print("2. 匹配闸门所属水路...")
|
||||
matched_count = 0
|
||||
extended_matched_count = 0
|
||||
unmatched_count = 0
|
||||
extended_match_distance = max_match_distance + extended_match_extra
|
||||
for idx, row in gate_gdf.iterrows():
|
||||
ch_code, ch_type, _, dist, match_dist = get_gate_belong_channel(
|
||||
row.geometry,
|
||||
channel_gdf,
|
||||
max_match_distance=max_match_distance,
|
||||
tolerance=max_match_distance,
|
||||
)
|
||||
used_extended_match = False
|
||||
if not ch_code and extended_match_extra > 0:
|
||||
ch_code, ch_type, _, dist, match_dist = get_gate_belong_channel(
|
||||
row.geometry,
|
||||
channel_gdf,
|
||||
max_match_distance=extended_match_distance,
|
||||
tolerance=extended_match_distance,
|
||||
)
|
||||
used_extended_match = bool(ch_code)
|
||||
gate_gdf.loc[idx, ["channel_code", "channel_type", "along_dist"]] = [ch_code, ch_type, dist]
|
||||
gate_name = _gate_log_name(row, f"index={idx}")
|
||||
if ch_code:
|
||||
matched_count += 1
|
||||
if used_extended_match:
|
||||
extended_matched_count += 1
|
||||
print(
|
||||
f" - 闸门 {gate_name} -> 水路 {ch_code} ({ch_type}), "
|
||||
f"along_dist={float(dist):.12f}, match_dist={float(match_dist):.12f} "
|
||||
f"[WARN] 超出常规阈值 {max_match_distance:.8f}，使用扩展阈值 {extended_match_distance:.8f} 匹配"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f" - 闸门 {gate_name} -> 水路 {ch_code} ({ch_type}), "
|
||||
f"along_dist={float(dist):.12f}, match_dist={float(match_dist):.12f}"
|
||||
)
|
||||
else:
|
||||
unmatched_count += 1
|
||||
print(f" - 闸门 {gate_name} -> 未匹配到水路")
|
||||
|
||||
print(f" 匹配完成: {matched_count}/{len(gate_gdf)}")
|
||||
if extended_matched_count:
|
||||
print(
|
||||
f" [WARN] 闸门扩展阈值匹配 {extended_matched_count} 个，"
|
||||
f"常规阈值 {max_match_distance:.8f}，扩展阈值 {extended_match_distance:.8f}"
|
||||
)
|
||||
if unmatched_count:
|
||||
print(f" [WARN] 闸门未匹配 {unmatched_count} 个，不参与编号与最终输出")
|
||||
|
||||
print("3. DFS 排序闸门...")
|
||||
gate_sorted_df = sort_gates_by_dfs(gate_gdf, channel_gdf, tolerance=max_match_distance)
|
||||
if gate_sorted_df is None or len(gate_sorted_df) == 0:
|
||||
print("[WARN] 闸门排序结果为空，返回空结果")
|
||||
empty = gpd.GeoDataFrame(columns=["geometry", "code", "类型", "type"], geometry="geometry", crs=gate_gdf.crs)
|
||||
empty.attrs["detected_count"] = int(len(gate_gdf))
|
||||
empty.attrs["matched_count"] = int(matched_count)
|
||||
empty.attrs["extended_matched_count"] = int(extended_matched_count)
|
||||
empty.attrs["unmatched_count"] = int(unmatched_count)
|
||||
return empty
|
||||
|
||||
gate_sorted_gdf = gpd.GeoDataFrame(gate_sorted_df, geometry="geometry", crs=gate_gdf.crs)
|
||||
print(" DFS 排序明细:")
|
||||
for order_idx, (_, row) in enumerate(gate_sorted_gdf.iterrows(), start=1):
|
||||
gate_name = _gate_log_name(row, f"index={order_idx}")
|
||||
print(
|
||||
f" - 顺序 {order_idx}: 闸门 {gate_name} | 水路 {row.get('channel_code', '')} "
|
||||
f"| along_dist={float(row.get('along_dist', 0.0)):.12f}"
|
||||
)
|
||||
|
||||
print("4. 赋值全局编号...")
|
||||
counter = 1
|
||||
for idx in gate_sorted_gdf.index:
|
||||
ch_code = gate_sorted_gdf.loc[idx, "channel_code"]
|
||||
gate_sorted_gdf.loc[idx, "code"] = f"{ch_code}-G{counter}"
|
||||
print(f" - 生成编号: {gate_sorted_gdf.loc[idx, 'code']}")
|
||||
counter += 1
|
||||
|
||||
print("5. 补全闸门类型...")
|
||||
for idx, row in gate_sorted_gdf.iterrows():
|
||||
cn, en = gate_type_complement(row)
|
||||
gate_sorted_gdf.loc[idx, ["类型", "type"]] = [cn, en]
|
||||
|
||||
helper_cols = ["channel_code", "channel_type", "along_dist"]
|
||||
gate_sorted_gdf = gate_sorted_gdf.drop(columns=[col for col in helper_cols if col in gate_sorted_gdf.columns])
|
||||
|
||||
print(f"完成闸门编号，共 {len(gate_sorted_gdf)} 个")
|
||||
gate_sorted_gdf.attrs["detected_count"] = int(len(gate_gdf))
|
||||
gate_sorted_gdf.attrs["matched_count"] = int(matched_count)
|
||||
gate_sorted_gdf.attrs["extended_matched_count"] = int(extended_matched_count)
|
||||
gate_sorted_gdf.attrs["unmatched_count"] = int(unmatched_count)
|
||||
return gate_sorted_gdf
|
||||
+604
@@ -0,0 +1,604 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
import threading
|
||||
import queue
|
||||
from pathlib import Path
|
||||
import tkinter as tk
|
||||
from tkinter.scrolledtext import ScrolledText
|
||||
|
||||
import geopandas as gpd
|
||||
import pandas as pd
|
||||
|
||||
from segment_code import segmen_main
|
||||
from gate_code import assign_gate_code_main
|
||||
|
||||
|
||||
CHANNEL_TOLERANCE = 0.00005
|
||||
GATE_MATCH_TOLERANCE = 0.00010
|
||||
|
||||
|
||||
def _count_filled_code(gdf):
|
||||
if gdf is None or len(gdf) == 0 or "code" not in gdf.columns:
|
||||
return 0
|
||||
return int(gdf["code"].fillna("").astype(str).str.strip().ne("").sum())
|
||||
|
||||
|
||||
# 田块赋值code
|
||||
def assign_field_code(field_shp_path, gate_gdf):
|
||||
# 读取数据
|
||||
field_gdf = gpd.read_file(field_shp_path, encoding='utf-8')
|
||||
# gate_gdf = gpd.read_file(gate_shp_path, encoding='utf-8')
|
||||
|
||||
# 确保 CRS 一致
|
||||
if field_gdf.crs != gate_gdf.crs:
|
||||
gate_gdf = gate_gdf.to_crs(field_gdf.crs)
|
||||
|
||||
# 过滤非节制闸门（田块对应的闸门）
|
||||
gate_normal = gate_gdf[~gate_gdf['type'].isin(['main-G', 'branch-G'])].copy()
|
||||
|
||||
# 按 Gxx 排序全局顺序
|
||||
def gxx_sort_key(code):
|
||||
parts = code.split('-')
|
||||
g_part = parts[1] # Gxx
|
||||
return int(g_part.replace("G",""))
|
||||
|
||||
gate_normal['G_sort'] = gate_normal['code'].apply(gxx_sort_key)
|
||||
gate_normal = gate_normal.sort_values('G_sort').reset_index(drop=True)
|
||||
|
||||
# 初始化田块 code 字段
|
||||
field_gdf['code'] = ""
|
||||
if len(gate_normal) == 0:
|
||||
print("[WARN] 未找到可用闸门，田块不赋码直接返回")
|
||||
print(f"田块编号完成，共 {len(field_gdf)} 个田块")
|
||||
return field_gdf
|
||||
|
||||
# 1️⃣ 循环田块，找到包含闸门并赋值初始 code
|
||||
for idx, field in field_gdf.iterrows():
|
||||
# intersects 包含边界
|
||||
contained_gates = gate_normal[gate_normal.geometry.intersects(field.geometry)]
|
||||
|
||||
if len(contained_gates) == 0:
|
||||
print(f"[WARN] 田块 {field['name']} 内未找到非节制闸门，可能坐标微偏或在边界,尝试更换方法查找。")
|
||||
# 尝试微调 buffer 查找
|
||||
buffer = field.geometry.buffer(0.000009)
|
||||
contained_gates = gate_normal[gate_normal.geometry.intersects(buffer)]
|
||||
if len(contained_gates) == 0:
|
||||
print(f"[SKIP] 田块 {field['name']} 仍未找到对应闸门。")
|
||||
continue
|
||||
|
||||
# 按 Gxx 排序，取第一个闸门的 Sxx-Gxx 前缀
|
||||
contained_gates = contained_gates.sort_values('G_sort')
|
||||
sg_prefix = contained_gates.iloc[0]['code']
|
||||
|
||||
# Fxx 使用全局顺序
|
||||
field_gdf.at[idx, 'code'] = sg_prefix
|
||||
field_gdf.at[idx, 'gcode'] = int(sg_prefix.split("-")[-1][1:])
|
||||
|
||||
print(f"[OK] 田块 {field['name']} 编号: {field_gdf.at[idx, 'code']}")
|
||||
|
||||
# 2️⃣ 按 gcode 升序排序田块
|
||||
field_gdf = field_gdf.sort_values('gcode').reset_index(drop=True)
|
||||
# print(field_gdf)
|
||||
|
||||
# 3️⃣ 循环赋最终 Fxx
|
||||
f_counter = 1 # 全局 Fxx
|
||||
for idx, field in field_gdf.iterrows():
|
||||
# 跳过没有闸门的田块
|
||||
if not field['code']:
|
||||
continue
|
||||
f_code = f"F{f_counter}"
|
||||
field_gdf.at[idx, 'code'] = f"{field['code']}-{f_code}"
|
||||
f_counter += 1
|
||||
print(f"[OK] 田块 {field.get('name', idx)} 最终编号: {field_gdf.at[idx, 'code']}")
|
||||
|
||||
# 删除辅助列gcode
|
||||
field_gdf = field_gdf.drop(columns=['gcode'])
|
||||
|
||||
|
||||
# 保存
|
||||
print(f"田块编号完成，共 {len(field_gdf)} 个田块")
|
||||
return field_gdf
|
||||
|
||||
# 排水口赋值，并进行合并
|
||||
def merge_points_with_drain_gdf(gate_gdf, pump_gdf, drain_gdf, field_gdf):
|
||||
"""
|
||||
合并闸门、泵、排水口为一个 GeoDataFrame。
|
||||
排水口 code = 田块 code + '-drain'
|
||||
"""
|
||||
print(f"[OK] 已读取闸门: {len(gate_gdf)} 个")
|
||||
print(f"[OK] 已读取泵: {len(pump_gdf)} 个")
|
||||
|
||||
# 筛选排水口
|
||||
drain_points = drain_gdf[
|
||||
(drain_gdf.get('type', '') == 'drain-G') |
|
||||
(drain_gdf.get('类型', '') == '排水口')
|
||||
].copy()
|
||||
print(f"[OK] 筛选排水口: {len(drain_points)} 个")
|
||||
|
||||
# 给排水口赋值 code
|
||||
drain_points['code'] = ""
|
||||
for idx, drain in drain_points.iterrows():
|
||||
contained_fields = field_gdf[field_gdf.geometry.intersects(drain.geometry)]
|
||||
if len(contained_fields) == 0:
|
||||
# 微调 1米 buffer
|
||||
buffer = drain.geometry.buffer(0.000009)
|
||||
contained_fields = field_gdf[field_gdf.geometry.intersects(buffer)]
|
||||
|
||||
if len(contained_fields) == 0:
|
||||
print(f"[WARN] 排水口 {drain.get('name', idx)} 未找到对应田块，跳过")
|
||||
continue
|
||||
|
||||
# code = 田块 code 去掉末尾 -Fxx 后，再拼接 -drain
|
||||
field_code = str(contained_fields.iloc[0]['code'])
|
||||
parts = field_code.split("-")
|
||||
if len(parts) > 0 and parts[-1].startswith("F") and parts[-1][1:].isdigit():
|
||||
base_code = "-".join(parts[:-1])
|
||||
else:
|
||||
base_code = field_code
|
||||
drain_points.at[idx, 'code'] = f"{base_code}-drain"
|
||||
print(f"[OK] 排水口 {drain.get('name', idx)} 编号: {drain_points.at[idx, 'code']}")
|
||||
|
||||
# 按字段并集合并闸门、泵和排水口，尽量保留原始点位字段
|
||||
merged_gdf = pd.concat([gate_gdf, pump_gdf, drain_points], ignore_index=True)
|
||||
merged_gdf = gpd.GeoDataFrame(merged_gdf, geometry='geometry', crs=gate_gdf.crs)
|
||||
merged_gdf = complement_point_types(merged_gdf)
|
||||
|
||||
return merged_gdf
|
||||
|
||||
|
||||
def complement_point_types(point_gdf):
|
||||
point_gdf = point_gdf.copy()
|
||||
if "类型" not in point_gdf.columns:
|
||||
point_gdf["类型"] = ""
|
||||
if "type" not in point_gdf.columns:
|
||||
point_gdf["type"] = ""
|
||||
|
||||
cn2en = {
|
||||
"进水闸": "inlet-G",
|
||||
"进排水闸": "inout-G",
|
||||
"排水闸": "drain-G",
|
||||
"排水口": "drain-G",
|
||||
"节制闸": "main-G",
|
||||
"泵": "pump",
|
||||
"泵房": "pump",
|
||||
}
|
||||
en2cn = {
|
||||
"inlet-g": "进水闸",
|
||||
"inout-g": "进排水闸",
|
||||
"drain-g": "排水闸",
|
||||
"main-g": "节制闸",
|
||||
"branch-g": "节制闸",
|
||||
"pump": "泵",
|
||||
}
|
||||
|
||||
def _norm_text(value):
|
||||
if value is None or pd.isna(value):
|
||||
return ""
|
||||
text = str(value).strip()
|
||||
if text.lower() in ("", "nan", "none", "null"):
|
||||
return ""
|
||||
return text
|
||||
|
||||
for idx, row in point_gdf.iterrows():
|
||||
cn = _norm_text(row.get("类型", ""))
|
||||
en = _norm_text(row.get("type", ""))
|
||||
en_key = en.lower()
|
||||
|
||||
if not en and cn in cn2en:
|
||||
point_gdf.at[idx, "type"] = cn2en[cn]
|
||||
en = cn2en[cn]
|
||||
en_key = en.lower()
|
||||
if not cn and en_key in en2cn:
|
||||
point_gdf.at[idx, "类型"] = en2cn[en_key]
|
||||
|
||||
return point_gdf
|
||||
|
||||
|
||||
def _extract_num(code):
|
||||
if code is None or pd.isna(code):
|
||||
return None
|
||||
s = "".join([ch for ch in str(code) if ch.isdigit()])
|
||||
return int(s) if s else None
|
||||
|
||||
|
||||
def _format_rank_text(value, empty_value=""):
|
||||
if value is None or pd.isna(value):
|
||||
return empty_value
|
||||
text = str(value).strip()
|
||||
if text.lower() in ("", "nan", "none", "null", "<na>"):
|
||||
return empty_value
|
||||
try:
|
||||
number = float(text)
|
||||
except ValueError:
|
||||
return text
|
||||
if number.is_integer():
|
||||
return str(int(number))
|
||||
return text
|
||||
|
||||
|
||||
def add_distance_rank_channels(channel_gdf):
|
||||
if "code" not in channel_gdf.columns or "feed_by" not in channel_gdf.columns:
|
||||
print("[WARN] 渠道缺少 code/feed_by 字段，跳过 dis_rank")
|
||||
return channel_gdf
|
||||
|
||||
feed_map = dict(zip(channel_gdf["code"], channel_gdf["feed_by"]))
|
||||
|
||||
def find_main_channel(code):
|
||||
current = code
|
||||
visited = set()
|
||||
while True:
|
||||
parent = feed_map.get(current)
|
||||
if parent is None or parent in visited:
|
||||
return None
|
||||
visited.add(parent)
|
||||
if "P" in str(parent):
|
||||
return current if "S" in str(current) else parent
|
||||
current = parent
|
||||
|
||||
channel_gdf = channel_gdf.copy()
|
||||
channel_gdf["_main_chan"] = channel_gdf["code"].apply(find_main_channel)
|
||||
channel_gdf["_code_num"] = channel_gdf["code"].apply(_extract_num)
|
||||
channel_gdf["dis_rank"] = (
|
||||
channel_gdf.groupby("_main_chan")["_code_num"]
|
||||
.rank(method="dense", ascending=True)
|
||||
.astype("Int64")
|
||||
)
|
||||
drain_mask = pd.Series(False, index=channel_gdf.index)
|
||||
if "ch_type" in channel_gdf.columns:
|
||||
drain_mask = drain_mask | channel_gdf["ch_type"].fillna("").astype(str).str.lower().eq("drain-s")
|
||||
if "渠型" in channel_gdf.columns:
|
||||
drain_mask = drain_mask | channel_gdf["渠型"].fillna("").astype(str).str.strip().eq("排水渠")
|
||||
channel_gdf["dis_rank"] = channel_gdf["dis_rank"].apply(_format_rank_text)
|
||||
channel_gdf.loc[drain_mask, "dis_rank"] = "0"
|
||||
channel_gdf = channel_gdf.drop(columns=["_main_chan", "_code_num"])
|
||||
print("[INFO] 渠道 dis_rank 已生成")
|
||||
return channel_gdf
|
||||
|
||||
|
||||
def add_distance_rank_fields(field_gdf):
|
||||
if "code" not in field_gdf.columns:
|
||||
print("[WARN] 田块缺少 code 字段，跳过 dis_rank")
|
||||
return field_gdf
|
||||
|
||||
def extract_s(code):
|
||||
if code is None or pd.isna(code):
|
||||
return None
|
||||
s = str(code)
|
||||
idx = s.find("S")
|
||||
if idx == -1:
|
||||
return None
|
||||
num = ""
|
||||
for ch in s[idx + 1 :]:
|
||||
if ch.isdigit():
|
||||
num += ch
|
||||
else:
|
||||
break
|
||||
return f"S{num}" if num else None
|
||||
|
||||
def extract_f_num(code):
|
||||
if code is None or pd.isna(code):
|
||||
return None
|
||||
s = str(code)
|
||||
idx = s.find("F")
|
||||
if idx == -1:
|
||||
return None
|
||||
num = ""
|
||||
for ch in s[idx + 1 :]:
|
||||
if ch.isdigit():
|
||||
num += ch
|
||||
else:
|
||||
break
|
||||
return int(num) if num else None
|
||||
|
||||
field_gdf = field_gdf.copy()
|
||||
field_gdf["_S_code"] = field_gdf["code"].apply(extract_s)
|
||||
field_gdf["_F_num"] = field_gdf["code"].apply(extract_f_num)
|
||||
field_gdf["dis_rank"] = (
|
||||
field_gdf.groupby("_S_code")["_F_num"]
|
||||
.rank(method="dense", ascending=True)
|
||||
.astype("Int64")
|
||||
)
|
||||
field_gdf["dis_rank"] = field_gdf["dis_rank"].apply(_format_rank_text)
|
||||
empty_code_mask = field_gdf["code"].fillna("").astype(str).str.strip().eq("")
|
||||
field_gdf.loc[empty_code_mask, "dis_rank"] = ""
|
||||
field_gdf = field_gdf.drop(columns=["_S_code", "_F_num"])
|
||||
print("[INFO] 田块 dis_rank 已生成")
|
||||
return field_gdf
|
||||
|
||||
|
||||
def run_pipeline(points_shp_path, field_shp_path, channel_shp_path, out_dir):
|
||||
out_dir = Path(out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
geojson_dir = out_dir.parent / f"{out_dir.name}_geojson"
|
||||
geojson_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
segment_out_path = out_dir / Path(channel_shp_path).name
|
||||
field_out_path = out_dir / Path(field_shp_path).name
|
||||
point_out_path = out_dir / Path(points_shp_path).name
|
||||
segment_geojson_path = geojson_dir / (Path(channel_shp_path).stem + ".geojson")
|
||||
field_geojson_path = geojson_dir / (Path(field_shp_path).stem + ".geojson")
|
||||
point_geojson_path = geojson_dir / (Path(points_shp_path).stem + ".geojson")
|
||||
|
||||
print(f"[INFO] points: {points_shp_path}")
|
||||
print(f"[INFO] fields: {field_shp_path}")
|
||||
print(f"[INFO] channels: {channel_shp_path}")
|
||||
print(f"[INFO] out_dir: {out_dir}")
|
||||
|
||||
segment_gdf, pump_gdf = segmen_main(
|
||||
pump_shp_path=points_shp_path,
|
||||
channel_shp_path=channel_shp_path,
|
||||
tolerance=CHANNEL_TOLERANCE,
|
||||
)
|
||||
segment_gdf = add_distance_rank_channels(segment_gdf)
|
||||
segment_gdf.to_file(segment_out_path, encoding="utf-8")
|
||||
segment_gdf.to_file(segment_geojson_path, driver="GeoJSON")
|
||||
|
||||
gate_gdf = assign_gate_code_main(
|
||||
point_shp_path=points_shp_path,
|
||||
channel_gdf=segment_gdf,
|
||||
max_match_distance=GATE_MATCH_TOLERANCE,
|
||||
)
|
||||
|
||||
field_gdf = assign_field_code(
|
||||
field_shp_path=field_shp_path,
|
||||
gate_gdf=gate_gdf,
|
||||
)
|
||||
field_gdf = add_distance_rank_fields(field_gdf)
|
||||
field_gdf.to_file(field_out_path, encoding="utf-8")
|
||||
field_gdf.to_file(field_geojson_path, driver="GeoJSON")
|
||||
|
||||
drain_gdf = gpd.read_file(points_shp_path, encoding="utf-8")
|
||||
drain_gdf = drain_gdf[
|
||||
(drain_gdf["type"] == "drain-G")
|
||||
| (drain_gdf.get("类型", "") == "排水口")
|
||||
].copy()
|
||||
|
||||
merged_gdf = merge_points_with_drain_gdf(
|
||||
gate_gdf=gate_gdf,
|
||||
pump_gdf=pump_gdf,
|
||||
drain_gdf=drain_gdf,
|
||||
field_gdf=field_gdf,
|
||||
)
|
||||
merged_gdf.to_file(point_out_path, encoding="utf-8")
|
||||
merged_gdf.to_file(point_geojson_path, driver="GeoJSON")
|
||||
|
||||
print(f"[OK] 合并完成，导出文件: {point_out_path} | 总点数: {len(merged_gdf)}")
|
||||
stats = {
|
||||
"pump_detected": int(len(pump_gdf)) if pump_gdf is not None else 0,
|
||||
"pump_completed": _count_filled_code(pump_gdf),
|
||||
"channel_detected": int(segment_gdf.attrs.get("detected_count", len(segment_gdf))) if segment_gdf is not None else 0,
|
||||
"channel_completed": _count_filled_code(segment_gdf),
|
||||
"gate_detected": int(gate_gdf.attrs.get("detected_count", len(gate_gdf))) if gate_gdf is not None else 0,
|
||||
"gate_completed": _count_filled_code(gate_gdf),
|
||||
"gate_extended_matched": int(gate_gdf.attrs.get("extended_matched_count", 0)) if gate_gdf is not None else 0,
|
||||
"gate_unmatched": int(gate_gdf.attrs.get("unmatched_count", 0)) if gate_gdf is not None else 0,
|
||||
}
|
||||
print(
|
||||
"[SUMMARY] 本农场统计 | "
|
||||
f"泵站: {stats['pump_completed']}/{stats['pump_detected']} | "
|
||||
f"水路: {stats['channel_completed']}/{stats['channel_detected']} | "
|
||||
f"闸门: {stats['gate_completed']}/{stats['gate_detected']}"
|
||||
)
|
||||
if stats["gate_extended_matched"] or stats["gate_unmatched"]:
|
||||
print(
|
||||
"[SUMMARY] 闸门匹配提醒 | "
|
||||
f"扩展阈值匹配: {stats['gate_extended_matched']} | "
|
||||
f"未匹配: {stats['gate_unmatched']}"
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
def _find_target_shps(folder):
|
||||
points = []
|
||||
fields = []
|
||||
channels = []
|
||||
for p in folder.iterdir():
|
||||
if p.is_file() and p.suffix.lower() == ".shp":
|
||||
name = p.name
|
||||
if "节制闸" in name:
|
||||
points.append(p)
|
||||
elif "田块" in name:
|
||||
fields.append(p)
|
||||
elif "水路" in name:
|
||||
channels.append(p)
|
||||
|
||||
def pick_one(items, label):
|
||||
if len(items) == 0:
|
||||
return None
|
||||
if len(items) > 1:
|
||||
print(f"[WARN] {folder.name} 多个 {label} 文件，使用第一个: {items[0].name}")
|
||||
return items[0]
|
||||
|
||||
return pick_one(points, "节制闸"), pick_one(fields, "田块"), pick_one(channels, "水路")
|
||||
|
||||
|
||||
def _is_skipped_folder(folder):
|
||||
name = folder.name
|
||||
if name in ("初始文件", "__pycache__"):
|
||||
return True
|
||||
if "result" in name.lower():
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _choose_base_dir(tool_dir):
|
||||
tool_dir = Path(tool_dir).resolve()
|
||||
|
||||
def has_target_dirs(base):
|
||||
for d in base.iterdir():
|
||||
if not d.is_dir():
|
||||
continue
|
||||
if _is_skipped_folder(d):
|
||||
continue
|
||||
points, fields, channels = _find_target_shps(d)
|
||||
if points and fields and channels:
|
||||
return True
|
||||
return False
|
||||
|
||||
if has_target_dirs(tool_dir):
|
||||
return tool_dir, None
|
||||
return tool_dir.parent, tool_dir.name
|
||||
|
||||
|
||||
class _TeeStream:
|
||||
def __init__(self, log_fp, msg_queue):
|
||||
self.log_fp = log_fp
|
||||
self.msg_queue = msg_queue
|
||||
|
||||
def write(self, msg):
|
||||
if not msg:
|
||||
return
|
||||
self.log_fp.write(msg)
|
||||
self.log_fp.flush()
|
||||
self.msg_queue.put(msg)
|
||||
|
||||
def flush(self):
|
||||
self.log_fp.flush()
|
||||
|
||||
|
||||
def _run_batch_worker(timestamp, tool_dir, base_dir, tool_folder_name):
|
||||
start_ts = time.time()
|
||||
|
||||
print(f"[INFO] tool_dir: {tool_dir}")
|
||||
print(f"[INFO] base_dir: {base_dir}")
|
||||
if tool_folder_name:
|
||||
print(f"[INFO] skip tool folder: {tool_folder_name}")
|
||||
|
||||
result_root = base_dir / f"result_{timestamp}"
|
||||
result_root.mkdir(parents=True, exist_ok=True)
|
||||
print(f"[INFO] result_root: {result_root}")
|
||||
|
||||
targets = []
|
||||
for d in base_dir.iterdir():
|
||||
if not d.is_dir():
|
||||
continue
|
||||
if _is_skipped_folder(d):
|
||||
continue
|
||||
if tool_folder_name and d.name == tool_folder_name:
|
||||
continue
|
||||
points, fields, channels = _find_target_shps(d)
|
||||
if points and fields and channels:
|
||||
targets.append((d, points, fields, channels))
|
||||
|
||||
if not targets:
|
||||
print("[ERROR] 未找到目标农场文件夹，请确定文件夹与工具在同一文件夹下。")
|
||||
else:
|
||||
print(f"[INFO] 目标农场数量: {len(targets)}")
|
||||
|
||||
ok = []
|
||||
failed = []
|
||||
failed_reasons = {}
|
||||
total_stats = {
|
||||
"pump_detected": 0,
|
||||
"pump_completed": 0,
|
||||
"channel_detected": 0,
|
||||
"channel_completed": 0,
|
||||
"gate_detected": 0,
|
||||
"gate_completed": 0,
|
||||
"gate_extended_matched": 0,
|
||||
"gate_unmatched": 0,
|
||||
}
|
||||
|
||||
for d, points, fields, channels in targets:
|
||||
name = d.name
|
||||
out_dir = result_root / f"{name}_result"
|
||||
print(f"\n[INFO] 处理农场: {name}")
|
||||
try:
|
||||
stats = run_pipeline(points, fields, channels, out_dir) or {}
|
||||
for key in total_stats:
|
||||
total_stats[key] += int(stats.get(key, 0))
|
||||
ok.append(name)
|
||||
except Exception as e:
|
||||
failed.append(name)
|
||||
failed_reasons[name] = f"{type(e).__name__}: {e}"
|
||||
print(f"[ERROR] 处理失败: {name}: {failed_reasons[name]}")
|
||||
traceback.print_exc()
|
||||
|
||||
elapsed = time.time() - start_ts
|
||||
print("\n[SUMMARY] 处理完成")
|
||||
print(f"[SUMMARY] 成功: {len(ok)}")
|
||||
print(f"[SUMMARY] 失败: {len(failed)}")
|
||||
print(f"[SUMMARY] 泵站: 检测到 {total_stats['pump_detected']} | 完成 {total_stats['pump_completed']}")
|
||||
print(f"[SUMMARY] 水路: 检测到 {total_stats['channel_detected']} | 完成 {total_stats['channel_completed']}")
|
||||
print(f"[SUMMARY] 闸门: 检测到 {total_stats['gate_detected']} | 完成 {total_stats['gate_completed']}")
|
||||
if total_stats["gate_extended_matched"] or total_stats["gate_unmatched"]:
|
||||
print(
|
||||
f"[SUMMARY] 闸门匹配提醒: 扩展阈值匹配 {total_stats['gate_extended_matched']} | "
|
||||
f"未匹配 {total_stats['gate_unmatched']}"
|
||||
)
|
||||
if failed:
|
||||
print(f"[SUMMARY] 失败列表: {', '.join(failed)}")
|
||||
for name in failed:
|
||||
print(f"[SUMMARY] 失败原因: {name}: {failed_reasons.get(name, 'unknown')}")
|
||||
print(f"[SUMMARY] 耗时: {elapsed:.2f} 秒")
|
||||
print("[INFO] 结束后请手动关闭窗口。")
|
||||
|
||||
|
||||
def run_batch_gui():
|
||||
start_ts = time.time()
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S", time.localtime(start_ts))
|
||||
if getattr(sys, "frozen", False):
|
||||
tool_dir = Path(sys.executable).resolve().parent
|
||||
try:
|
||||
os.chdir(tool_dir)
|
||||
except Exception:
|
||||
pass
|
||||
base_dir = tool_dir
|
||||
tool_folder_name = None
|
||||
else:
|
||||
tool_dir = Path(__file__).resolve().parent
|
||||
base_dir, tool_folder_name = _choose_base_dir(tool_dir)
|
||||
log_path = base_dir / f"log_{timestamp}.txt"
|
||||
|
||||
root = tk.Tk()
|
||||
root.title("Irrigation Batch Tool")
|
||||
root.geometry("900x600")
|
||||
root.configure(bg="black")
|
||||
|
||||
text = ScrolledText(
|
||||
root,
|
||||
wrap=tk.WORD,
|
||||
bg="black",
|
||||
fg="white",
|
||||
insertbackground="white",
|
||||
)
|
||||
text.pack(fill=tk.BOTH, expand=True)
|
||||
|
||||
msg_queue = queue.Queue()
|
||||
log_fp = open(log_path, "w", encoding="utf-8")
|
||||
|
||||
def poll_queue():
|
||||
while True:
|
||||
try:
|
||||
msg = msg_queue.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
text.insert(tk.END, msg)
|
||||
text.see(tk.END)
|
||||
root.after(100, poll_queue)
|
||||
|
||||
def worker():
|
||||
old_out = sys.stdout
|
||||
old_err = sys.stderr
|
||||
tee = _TeeStream(log_fp, msg_queue)
|
||||
sys.stdout = tee
|
||||
sys.stderr = tee
|
||||
try:
|
||||
_run_batch_worker(timestamp, tool_dir, base_dir, tool_folder_name)
|
||||
finally:
|
||||
sys.stdout = old_out
|
||||
sys.stderr = old_err
|
||||
log_fp.flush()
|
||||
|
||||
threading.Thread(target=worker, daemon=True).start()
|
||||
root.after(100, poll_queue)
|
||||
root.mainloop()
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_batch_gui()
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,390 @@
|
||||
import warnings
|
||||
|
||||
import geopandas as gpd
|
||||
import pandas as pd
|
||||
from shapely.geometry import LineString, MultiLineString, Point
|
||||
from shapely.ops import nearest_points
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
|
||||
def _norm_text(value):
|
||||
if value is None or pd.isna(value):
|
||||
return ""
|
||||
text = str(value).strip()
|
||||
if text.lower() in {"", "nan", "none", "null"}:
|
||||
return ""
|
||||
return text
|
||||
|
||||
|
||||
def _match_text(value, candidates):
|
||||
text = _norm_text(value)
|
||||
if not text:
|
||||
return False
|
||||
return text in candidates
|
||||
|
||||
|
||||
def _has_pump_code(value):
|
||||
return "P" in _norm_text(value).upper()
|
||||
|
||||
|
||||
def _extract_pump_number(value):
|
||||
text = _norm_text(value).upper().replace("P", "")
|
||||
return int(text) if text.isdigit() else None
|
||||
|
||||
|
||||
def _pump_sort_key(value):
|
||||
text = _norm_text(value).upper()
|
||||
number = _extract_pump_number(text)
|
||||
if number is not None:
|
||||
return (0, number, text)
|
||||
return (1, 0, text)
|
||||
|
||||
|
||||
def _next_available_pump_number(existing_codes):
|
||||
numeric_values = [num for num in (_extract_pump_number(code) for code in existing_codes) if num is not None]
|
||||
return max(numeric_values, default=0) + 1
|
||||
|
||||
|
||||
def _flatten_line_geometry(geom):
|
||||
if geom is None or not geom.is_valid:
|
||||
return None
|
||||
if geom.geom_type == "LineString":
|
||||
return geom if len(geom.coords) >= 2 else None
|
||||
if geom.geom_type == "MultiLineString":
|
||||
coords = []
|
||||
for part in geom.geoms:
|
||||
if len(part.coords) >= 2:
|
||||
coords.extend(list(part.coords))
|
||||
return LineString(coords) if len(coords) >= 2 else None
|
||||
return None
|
||||
|
||||
|
||||
def _channel_sort_key(row):
|
||||
start_point = row["start_point"]
|
||||
return (-start_point.y, -start_point.x, row.name)
|
||||
|
||||
|
||||
def _is_drain_channel(row):
|
||||
ch_type = _norm_text(row.get("ch_type", "")).lower()
|
||||
channel_type_name = _norm_text(row.get("渠型", ""))
|
||||
return ch_type == "drain-s" or channel_type_name == "排水渠"
|
||||
|
||||
|
||||
def preprocess_pump_points(pump_shp_path):
|
||||
"""Filter pump features, preserve valid P codes, and sort them stably."""
|
||||
gdf_pump = gpd.read_file(pump_shp_path)
|
||||
gdf_pump = gdf_pump.loc[:, ~gdf_pump.columns.duplicated()].copy()
|
||||
|
||||
type_mask = pd.Series(False, index=gdf_pump.index)
|
||||
if "type" in gdf_pump.columns:
|
||||
type_mask = gdf_pump["type"].astype(str).str.lower().eq("pump")
|
||||
if "类型" in gdf_pump.columns:
|
||||
type_mask = type_mask | gdf_pump["类型"].apply(lambda x: _match_text(x, {"泵", "泵站", "泵房"}))
|
||||
|
||||
if "type" not in gdf_pump.columns and "类型" not in gdf_pump.columns:
|
||||
raise ValueError("Pump point data is missing the type/类型 field.")
|
||||
|
||||
gdf_pump = gdf_pump[type_mask].copy().reset_index(drop=True)
|
||||
if "code" not in gdf_pump.columns:
|
||||
gdf_pump["code"] = ""
|
||||
|
||||
if len(gdf_pump) == 0:
|
||||
return gdf_pump
|
||||
|
||||
valid_mask = gdf_pump["code"].apply(_has_pump_code)
|
||||
existing_codes = gdf_pump.loc[valid_mask, "code"].tolist()
|
||||
existing_upper = {_norm_text(code).upper() for code in existing_codes if _norm_text(code)}
|
||||
next_number = _next_available_pump_number(existing_codes)
|
||||
|
||||
for idx in gdf_pump.index[~valid_mask]:
|
||||
while True:
|
||||
candidate = f"P{next_number}"
|
||||
next_number += 1
|
||||
if candidate.upper() not in existing_upper:
|
||||
gdf_pump.at[idx, "code"] = candidate
|
||||
existing_upper.add(candidate.upper())
|
||||
break
|
||||
|
||||
sort_parts = gdf_pump["code"].apply(_pump_sort_key)
|
||||
gdf_pump["_sort_group"] = sort_parts.apply(lambda x: x[0])
|
||||
gdf_pump["_sort_num"] = sort_parts.apply(lambda x: x[1])
|
||||
gdf_pump["_sort_text"] = sort_parts.apply(lambda x: x[2])
|
||||
gdf_pump = gdf_pump.sort_values(["_sort_group", "_sort_num", "_sort_text"]).reset_index(drop=True)
|
||||
gdf_pump = gdf_pump.drop(columns=["_sort_group", "_sort_num", "_sort_text"])
|
||||
|
||||
return gdf_pump
|
||||
|
||||
|
||||
def get_channel_start_end(channel_gdf):
|
||||
"""Extract start/end points and drop invalid line geometries."""
|
||||
channel_gdf = channel_gdf.copy()
|
||||
channel_gdf["_flat_geometry"] = channel_gdf["geometry"].apply(_flatten_line_geometry)
|
||||
channel_gdf = channel_gdf[channel_gdf["_flat_geometry"].notna()].copy()
|
||||
channel_gdf["start_point"] = channel_gdf["_flat_geometry"].apply(lambda geom: Point(geom.coords[0]))
|
||||
channel_gdf["end_point"] = channel_gdf["_flat_geometry"].apply(lambda geom: Point(geom.coords[-1]))
|
||||
return channel_gdf.drop(columns=["_flat_geometry"])
|
||||
|
||||
|
||||
def is_point_on_channel_segment(point, channel_geom, tolerance=0.00005):
|
||||
"""Check whether a point lies on the interior of a channel segment."""
|
||||
if point is None or channel_geom is None or not channel_geom.is_valid:
|
||||
return False
|
||||
|
||||
flat_geom = _flatten_line_geometry(channel_geom)
|
||||
if flat_geom is None:
|
||||
return False
|
||||
|
||||
coords = list(flat_geom.coords)
|
||||
if len(coords) < 2:
|
||||
return False
|
||||
|
||||
start_p = Point(coords[0])
|
||||
end_p = Point(coords[-1])
|
||||
if point.distance(start_p) < tolerance or point.distance(end_p) < tolerance:
|
||||
return False
|
||||
|
||||
if not point.within(flat_geom.buffer(tolerance)):
|
||||
return False
|
||||
|
||||
for i in range(len(coords) - 1):
|
||||
segment = LineString([coords[i], coords[i + 1]])
|
||||
if point.within(segment.buffer(tolerance)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def calculate_along_distance(channel_geom, point, tolerance=0.00005):
|
||||
"""Calculate the distance along a channel from its start to a projected point."""
|
||||
if point is None or channel_geom is None or not channel_geom.is_valid:
|
||||
return 0.0
|
||||
|
||||
flat_geom = _flatten_line_geometry(channel_geom)
|
||||
if flat_geom is None:
|
||||
return 0.0
|
||||
|
||||
coords = list(flat_geom.coords)
|
||||
if len(coords) < 2:
|
||||
return 0.0
|
||||
|
||||
proj_point, _ = nearest_points(flat_geom, point)
|
||||
total_dist = 0.0
|
||||
|
||||
for i in range(len(coords) - 1):
|
||||
segment = LineString([coords[i], coords[i + 1]])
|
||||
if proj_point.within(segment.buffer(tolerance)):
|
||||
total_dist += Point(coords[i]).distance(proj_point)
|
||||
break
|
||||
total_dist += segment.length
|
||||
|
||||
return total_dist
|
||||
|
||||
|
||||
def _find_main_channels(channel_gdf, pump_gdf, tolerance):
|
||||
main_by_pump = {pump_idx: [] for pump_idx in pump_gdf.index}
|
||||
main_owner = {}
|
||||
|
||||
for idx, row in channel_gdf.iterrows():
|
||||
best_match = None
|
||||
for pump_idx, pump_row in pump_gdf.iterrows():
|
||||
dist = row["start_point"].distance(pump_row["geometry"])
|
||||
if dist < tolerance:
|
||||
score = (dist, pump_idx)
|
||||
if best_match is None or score < best_match[0]:
|
||||
best_match = (score, pump_idx)
|
||||
if best_match is None:
|
||||
continue
|
||||
pump_idx = best_match[1]
|
||||
main_by_pump[pump_idx].append(idx)
|
||||
main_owner[idx] = pump_gdf.at[pump_idx, "code"]
|
||||
|
||||
for pump_idx, indices in main_by_pump.items():
|
||||
main_by_pump[pump_idx] = sorted(indices)
|
||||
|
||||
return main_by_pump, main_owner
|
||||
|
||||
|
||||
def _select_best_parent(child_idx, child_row, channel_gdf, tolerance):
|
||||
start_point = child_row["start_point"]
|
||||
candidates = []
|
||||
|
||||
for parent_idx, parent_row in channel_gdf.iterrows():
|
||||
if parent_idx == child_idx:
|
||||
continue
|
||||
if not is_point_on_channel_segment(start_point, parent_row["geometry"], tolerance):
|
||||
continue
|
||||
|
||||
dist = start_point.distance(parent_row["geometry"])
|
||||
along_dist = calculate_along_distance(parent_row["geometry"], start_point, tolerance)
|
||||
candidates.append((round(dist, 12), -along_dist, parent_idx))
|
||||
|
||||
if not candidates:
|
||||
return None
|
||||
|
||||
candidates.sort()
|
||||
return candidates[0][2]
|
||||
|
||||
|
||||
def _build_parent_map(channel_gdf, root_indices, tolerance):
|
||||
parent_map = {}
|
||||
for idx, row in channel_gdf.iterrows():
|
||||
if idx in root_indices:
|
||||
continue
|
||||
parent_idx = _select_best_parent(idx, row, channel_gdf, tolerance)
|
||||
if parent_idx is not None:
|
||||
parent_map[idx] = parent_idx
|
||||
return parent_map
|
||||
|
||||
|
||||
def _collect_reachable_indices(root_indices, parent_map):
|
||||
reachable = set(root_indices)
|
||||
unresolved = set(parent_map.keys())
|
||||
|
||||
progress = True
|
||||
while progress:
|
||||
progress = False
|
||||
for idx in list(unresolved):
|
||||
parent_idx = parent_map.get(idx)
|
||||
if parent_idx in reachable:
|
||||
reachable.add(idx)
|
||||
unresolved.remove(idx)
|
||||
progress = True
|
||||
|
||||
return reachable
|
||||
|
||||
|
||||
def _build_children_map(parent_map, reachable_indices):
|
||||
children_map = {}
|
||||
for child_idx, parent_idx in parent_map.items():
|
||||
if child_idx not in reachable_indices or parent_idx not in reachable_indices:
|
||||
continue
|
||||
children_map.setdefault(parent_idx, []).append(child_idx)
|
||||
return children_map
|
||||
|
||||
|
||||
def _sort_children(parent_idx, children_indices, channel_gdf, tolerance):
|
||||
parent_geom = channel_gdf.at[parent_idx, "geometry"]
|
||||
return sorted(
|
||||
children_indices,
|
||||
key=lambda child_idx: calculate_along_distance(parent_geom, channel_gdf.at[child_idx, "start_point"], tolerance),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
|
||||
def _assign_tree_codes(channel_gdf, root_idx, pump_code, children_map, code_counter, tolerance):
|
||||
code_counter += 1
|
||||
root_code = f"S{code_counter}"
|
||||
channel_gdf.at[root_idx, "code"] = root_code
|
||||
channel_gdf.at[root_idx, "type"] = "main-S"
|
||||
channel_gdf.at[root_idx, "feed_by"] = pump_code
|
||||
|
||||
children = _sort_children(root_idx, children_map.get(root_idx, []), channel_gdf, tolerance)
|
||||
for child_idx in children:
|
||||
code_counter = _assign_branch_codes(channel_gdf, child_idx, root_code, children_map, code_counter, tolerance)
|
||||
|
||||
return code_counter
|
||||
|
||||
|
||||
def _assign_branch_codes(channel_gdf, child_idx, parent_code, children_map, code_counter, tolerance):
|
||||
code_counter += 1
|
||||
child_code = f"S{code_counter}"
|
||||
channel_gdf.at[child_idx, "code"] = child_code
|
||||
channel_gdf.at[child_idx, "type"] = "branch-S"
|
||||
channel_gdf.at[child_idx, "feed_by"] = parent_code
|
||||
|
||||
children = _sort_children(child_idx, children_map.get(child_idx, []), channel_gdf, tolerance)
|
||||
for grandchild_idx in children:
|
||||
code_counter = _assign_branch_codes(channel_gdf, grandchild_idx, child_code, children_map, code_counter, tolerance)
|
||||
|
||||
return code_counter
|
||||
|
||||
|
||||
def _assign_drain_codes(channel_gdf, start_number):
|
||||
drain_indices = channel_gdf[channel_gdf["_is_drain"]].copy()
|
||||
if len(drain_indices) == 0:
|
||||
return start_number
|
||||
|
||||
sorted_indices = sorted(drain_indices.index, key=lambda idx: _channel_sort_key(channel_gdf.loc[idx]))
|
||||
|
||||
code_counter = start_number
|
||||
for idx in sorted_indices:
|
||||
code_counter += 1
|
||||
channel_gdf.at[idx, "code"] = f"S{code_counter}"
|
||||
channel_gdf.at[idx, "type"] = ""
|
||||
channel_gdf.at[idx, "feed_by"] = ""
|
||||
|
||||
return code_counter
|
||||
|
||||
|
||||
def assign_channel_codes(pump_gdf, channel_gdf, tolerance=0.00005):
|
||||
"""Assign main/branch channel codes first, then append drain channels."""
|
||||
channel_gdf = channel_gdf.copy()
|
||||
channel_gdf = get_channel_start_end(channel_gdf)
|
||||
print(f"过滤后有效渠道数量: {len(channel_gdf)}")
|
||||
|
||||
for col in ["code", "type", "feed_by"]:
|
||||
channel_gdf[col] = ""
|
||||
|
||||
channel_gdf["_is_drain"] = channel_gdf.apply(_is_drain_channel, axis=1)
|
||||
work_gdf = channel_gdf[~channel_gdf["_is_drain"]].copy()
|
||||
drain_count = int(channel_gdf["_is_drain"].sum())
|
||||
|
||||
print(f"非排水渠数量: {len(work_gdf)}")
|
||||
print(f"排水渠数量: {drain_count}")
|
||||
|
||||
code_counter = 0
|
||||
if len(work_gdf) > 0:
|
||||
main_by_pump, main_owner = _find_main_channels(work_gdf, pump_gdf, tolerance)
|
||||
root_indices = set(main_owner.keys())
|
||||
parent_map = _build_parent_map(work_gdf, root_indices, tolerance)
|
||||
reachable_indices = _collect_reachable_indices(root_indices, parent_map)
|
||||
children_map = _build_children_map(parent_map, reachable_indices)
|
||||
|
||||
for pump_idx, pump_row in pump_gdf.iterrows():
|
||||
root_list = [idx for idx in main_by_pump.get(pump_idx, []) if idx in reachable_indices]
|
||||
for root_idx in root_list:
|
||||
code_counter = _assign_tree_codes(work_gdf, root_idx, pump_row["code"], children_map, code_counter, tolerance)
|
||||
|
||||
for idx in work_gdf.index:
|
||||
channel_gdf.at[idx, "code"] = work_gdf.at[idx, "code"]
|
||||
channel_gdf.at[idx, "type"] = work_gdf.at[idx, "type"]
|
||||
channel_gdf.at[idx, "feed_by"] = work_gdf.at[idx, "feed_by"]
|
||||
|
||||
unassigned_count = int((work_gdf["code"] == "").sum())
|
||||
if unassigned_count:
|
||||
print(f"未编号非排水渠数量: {unassigned_count}")
|
||||
|
||||
code_counter = _assign_drain_codes(channel_gdf, code_counter)
|
||||
|
||||
temp_cols = ["start_point", "end_point", "_is_drain"]
|
||||
return channel_gdf.drop(columns=[col for col in temp_cols if col in channel_gdf.columns])
|
||||
|
||||
|
||||
def segmen_main(pump_shp_path, channel_shp_path, tolerance=0.00005):
|
||||
"""Entry point for channel coding."""
|
||||
print("=" * 60)
|
||||
pump_gdf = preprocess_pump_points(pump_shp_path)
|
||||
print(f"识别到泵点数量: {len(pump_gdf)}")
|
||||
for _, row in pump_gdf.iterrows():
|
||||
print(f" 泵 {row['code']}: 坐标({row['geometry'].x:.6f}, {row['geometry'].y:.6f})")
|
||||
|
||||
channel_gdf = gpd.read_file(channel_shp_path)
|
||||
detected_count = int(len(channel_gdf))
|
||||
print(f"原始渠道数量: {detected_count}")
|
||||
|
||||
result_gdf = assign_channel_codes(pump_gdf, channel_gdf, tolerance)
|
||||
result_gdf.attrs["detected_count"] = detected_count
|
||||
|
||||
main_num = len(result_gdf[result_gdf["type"] == "main-S"])
|
||||
branch_num = len(result_gdf[result_gdf["type"] == "branch-S"])
|
||||
drain_num = len(result_gdf[(result_gdf["code"] != "") & (result_gdf["type"] == "") & (result_gdf["feed_by"] == "")])
|
||||
unassigned_num = len(result_gdf[result_gdf["code"] == ""])
|
||||
|
||||
print("\n===== 编号结果统计 =====")
|
||||
print(f"主渠数量: {main_num}")
|
||||
print(f"支渠数量: {branch_num}")
|
||||
print(f"排水渠数量: {drain_num}")
|
||||
print(f"未编号渠道数量: {unassigned_num}")
|
||||
print("=" * 60)
|
||||
|
||||
return result_gdf, pump_gdf
|
||||
Reference in New Issue
Block a user