UAV/grid.py

import csv
import numpy as np
import matplotlib.pyplot as plt
import math

from shapely.geometry import box, MultiPoint
from shapely.ops import unary_union
from scipy.spatial import cKDTree
from utils.gps_extractor import GPSExtractor


# ---------------------- overlap 截断为不超过 10% ----------------------
def clamp_overlap(overlap):
    if overlap < 0:
        return 0.0
    elif overlap > 0.1:
        return 0.1
    else:
        return overlap


# ====================== 1) 生成可用矩形并记录其覆盖点集 ======================

def generate_rectangles_with_point_indices(points, w, h, overlap=0.1, min_points=800):
    """
    在 bounding box 内，以 (w, h) + overlap 布置网格，生成所有矩形。
    过滤：只保留"矩形内点数 >= min_points"的矩形。

    返回:
      rect_info_list: list of (rect_polygon, covered_indices)
        - rect_polygon: Shapely Polygon
        - covered_indices: 一个 set，表示该矩形覆盖的所有点索引
    """
    overlap = clamp_overlap(overlap)
    if len(points) == 0:
        return []

    minx, miny = np.min(points, axis=0)
    maxx, maxy = np.max(points, axis=0)

    # 特殊情况：只有一个点或非常小范围 -> 很难满足 800 点
    if abs(maxx - minx) < 1e-15 and abs(maxy - miny) < 1e-15:
        return []

    # 步长
    step_x = w * (1 - overlap)
    step_y = h * (1 - overlap)

    x_coords = np.arange(minx, maxx + step_x, step_x)
    y_coords = np.arange(miny, maxy + step_y, step_y)

    # 建立 KDTree，加速查找
    tree = cKDTree(points)
    rect_info_list = []

    for x in x_coords:
        for y in y_coords:
            rect_poly = box(x, y, x + w, y + h)
            rx_min, ry_min, rx_max, ry_max = rect_poly.bounds

            cx = (rx_min + rx_max) / 2
            cy = (ry_min + ry_max) / 2
            r = math.sqrt((rx_max - rx_min) ** 2 + (ry_max - ry_min) ** 2) / 2

            candidate_ids = tree.query_ball_point([cx, cy], r)
            if not candidate_ids:
                continue

            covered_set = set()
            for idx_pt in candidate_ids:
                px, py = points[idx_pt]
                if rx_min <= px <= rx_max and ry_min <= py <= ry_max:
                    covered_set.add(idx_pt)

            # 如果覆盖的点数不足 min_points，就不保留
            if len(covered_set) < min_points:
                continue

            rect_info_list.append((rect_poly, covered_set))

    return rect_info_list


# ====================== 2) 贪心算法选取子集覆盖所有点 ======================
def cover_all_points_greedy(points, rect_info_list):
    """
    给定所有点 points 以及 "可用矩形+覆盖点集合" rect_info_list，
    要求:
       - 选出若干矩形，使得所有点都被覆盖 (每个点至少属于1个选中矩形)
       - 最终并集面积最小 (做近似贪心)

    返回:
       chosen_rects: 最终选出的矩形列表 (每个是 shapely Polygon)
    """

    n = len(points)
    all_indices = set(range(n))  # 所有点的索引
    uncovered = set(all_indices)  # 尚未被覆盖的点索引

    chosen_rects = []
    union_polygon = None  # 当前已选矩形的并集

    # 如果没有任何矩形可用，就直接失败
    if not rect_info_list:
        return []

    # 为了在贪心过程中快速评估"新矩形带来的额外并集面积"，
    # 我们每次选择矩形后更新 union_polygon，然后比较 union_polygon.union(new_rect).area - union_polygon.area
    # 但 union_polygon 初始为 None

    while uncovered:
        best_gain = 0
        best_new_area = float('inf')
        best_rect = None
        best_covered_new = set()

        for rect_poly, covered_set in rect_info_list:
            # 计算能覆盖多少"尚未覆盖"的点
            newly_covered = uncovered.intersection(covered_set)
            if not newly_covered:
                continue

            # 计算额外增加的并集面积
            if union_polygon is None:
                # 第一次选，union_polygon 为空 => new_area = rect_poly.area
                new_area = rect_poly.area
                area_increase = new_area
            else:
                # 计算 union_polygon ∪ rect_poly 的面积
                test_union = union_polygon.union(rect_poly)
                new_area = test_union.area
                area_increase = new_area - union_polygon.area

            # 贪心策略：最大化 (覆盖点数量) / (面积增量)
            # 或者 equivalently, (覆盖点数量) 多、(面积增量) 小 都是好
            ratio = len(newly_covered) / max(area_increase, 1e-12)

            # 我们要找到 ratio 最大的那个
            if ratio > best_gain:
                best_gain = ratio
                best_new_area = area_increase
                best_rect = rect_poly
                best_covered_new = newly_covered

        if best_rect is None:
            # 没有可选的矩形能覆盖任何剩余点 => 失败 (无法覆盖所有点)
            return []

        # 选中 best_rect
        chosen_rects.append(best_rect)
        uncovered -= best_covered_new

        # 更新并集
        if union_polygon is None:
            union_polygon = best_rect
        else:
            union_polygon = union_polygon.union(best_rect)

    return chosen_rects


# ====================== 3) 主流程: 离散搜索 (w,h) + 贪心覆盖 ======================
def find_optimal_rectangles_cover_all_points(
        points,
        base_w,
        base_h,
        overlap=0.1,
        steps=5,
        min_points=800
):
    """
    在 [0.5*base_w,1.5*base_w] x [0.5*base_h,1.5*base_h] 的离散区间枚举 (w,h)，
    - 生成可用矩形(≥800 点)的列表
    - 用贪心算法选出子集来覆盖所有点
    - 计算选中矩形的并集面积
    选出面积最小的方案并返回
    """
    overlap = clamp_overlap(overlap)
    n = len(points)
    if n == 0:
        return [], (base_w, base_h), 0.0  # 没有点就不用覆盖了

    w_candidates = np.linspace(0.3 * base_w, 2 * base_w, steps)
    h_candidates = np.linspace(0.3 * base_h, 2 * base_h, steps)

    best_rects = []
    best_area = float('inf')
    best_w, best_h = base_w, base_h

    for w in w_candidates:
        for h in h_candidates:
            rect_info_list = generate_rectangles_with_point_indices(points, w, h, overlap, min_points)
            if not rect_info_list:
                # 说明没有任何矩形能达到 "≥800点"
                continue

            # 用贪心覆盖所有点
            chosen_rects = cover_all_points_greedy(points, rect_info_list)
            if not chosen_rects:
                # 无法覆盖所有点
                continue

            # 计算并集面积
            union_poly = unary_union(chosen_rects)
            area_covered = union_poly.area

            if area_covered < best_area:
                best_area = area_covered
                best_rects = chosen_rects
                best_w, best_h = w, h

    return best_rects, (best_w, best_h), best_area


# ====================== 4) 读取 CSV + 可视化 ======================
def plot_image_points_cover_all_min_area(
        image_dir,  # 新参数：图片文件夹路径
        base_rect_width=0.001,
        base_rect_height=0.001,
        overlap=0.1,
        steps=5,
        min_points=800
):
    """
    从图片文件夹读取GPS坐标:
     1) 使用 GPSExtractor 从图片中提取GPS坐标
     2) 在 [0.5*base_w,1.5*base_w] x [0.5*base_h,1.5*base_h] 离散搜索 (w,h)
     3) 对每个 (w,h), 先生成所有"含≥800点"的矩形 => 再用贪心覆盖所有点 => 计算并集面积
     4) 最小并集面积的方案即近似最优解
     5) 最终用该方案可视化
    """
    overlap = clamp_overlap(overlap)

    # 使用 GPSExtractor 读取图片GPS坐标
    extractor = GPSExtractor(image_dir)
    gps_df = extractor.extract_all_gps()
    
    if gps_df.empty:
        print("未能从图片中提取到GPS坐标。")
        return
        
    points = np.column_stack((gps_df['lon'], gps_df['lat']))  # (N, 2), [x=lon, y=lat]
    n = len(points)
    if n == 0:
        print("No points extracted from images.")
        return

    # 贪心 + 离散搜索
    chosen_rects, (best_w, best_h), best_area = find_optimal_rectangles_cover_all_points(
        points,
        base_w=base_rect_width,
        base_h=base_rect_height,
        overlap=overlap,
        steps=steps,
        min_points=min_points
    )

    if not chosen_rects:
        print(f"无法找到满足 '每个矩形≥{min_points}点' 且覆盖所有点 的方案，试着调大尺寸/步数/overlap。")
        return

    # 可视化
    plt.figure(figsize=(10, 8))
    # 画点
    plt.scatter(points[:, 0], points[:, 1], c='red', s=10, label='Points')

    # 画矩形
    for i, rect in enumerate(chosen_rects):
        if rect.is_empty:
            continue
        x, y = rect.exterior.xy
        plt.fill(x, y, edgecolor='green', fill=False, alpha=0.3,
                 label='Chosen Rectangles' if i == 0 else "")

    plt.title(
        f"Cover All Points, Each Rect≥{min_points} pts, Minimal Union Area\n"
        f"base=({base_rect_width:.6f} x {base_rect_height:.6f}), overlap≤{overlap}, steps={steps}\n"
        f"best (w,h)=({best_w:.6f},{best_h:.6f}), union area={best_area:.6f}, #rect={len(chosen_rects)}"
    )
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")
    plt.legend()
    plt.grid(True)
    plt.show()


# ------------------ 测试入口 ------------------
if __name__ == "__main__":
    image_dir = r"C:\datasets\134\code\images"  # 替换为你的图片文件夹路径
    plot_image_points_cover_all_min_area(
        image_dir,
        base_rect_width=0.01,
        base_rect_height=0.01,
        overlap=0.05,  # 会被截断到 0.1
        steps=40,
        min_points=100
    )