diff --git a/src/gate_code.py b/src/gate_code.py new file mode 100644 index 0000000..93cf2ea --- /dev/null +++ b/src/gate_code.py @@ -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 diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..109015a --- /dev/null +++ b/src/main.py @@ -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", ""): + 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() + + + + diff --git a/src/segment_code.py b/src/segment_code.py new file mode 100644 index 0000000..9eb8a8e --- /dev/null +++ b/src/segment_code.py @@ -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