上传文件至「src」

Signed-off-by: zhengzhoawen <zzw123@foxmail.com>

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

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", "<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()
+    
+
+
+    

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