WQ_GUI/src/postprocessing/point_map.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
采样点地图生成模块 - 在高光谱假彩色影像上标注采样点
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Optional, Tuple, List, Dict, Union
import warnings
from matplotlib.patches import FancyArrowPatch
import matplotlib.patheffects as path_effects

# 性能优化配置
plt.rcParams['agg.path.chunksize'] = 10000
plt.rcParams['path.simplify'] = True
plt.rcParams['path.simplify_threshold'] = 0.1

try:
    from osgeo import gdal, osr

    GDAL_AVAILABLE = True
except ImportError:
    GDAL_AVAILABLE = False
    print("警告: GDAL未安装，地理坐标转换功能可能无法正常工作")


class SamplingPointMap:
    def __init__(self, output_dir: str = "./point_maps", fast_mode: bool = False):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.fast_mode = fast_mode

        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans', 'Arial Unicode MS']
        plt.rcParams['axes.unicode_minus'] = False
        plt.rcParams['font.size'] = 12

        if fast_mode:
            plt.rcParams['figure.dpi'] = 150
            plt.rcParams['savefig.dpi'] = 150
            warnings.filterwarnings('ignore', category=UserWarning)
        else:
            plt.rcParams['figure.dpi'] = 300
            plt.rcParams['savefig.dpi'] = 300
        warnings.filterwarnings('ignore')

    def create_sampling_point_map(self, hyperspectral_path: str, csv_path: str,
                                  output_filename: Optional[str] = None, rgb_bands: Optional[List[int]] = None,
                                  point_color: str = 'red', point_size: int = 80, point_alpha: float = 0.8,
                                  show_north_arrow: bool = True, show_scale_bar: bool = True,
                                  show_legend: bool = True, dpi: int = None, downsample: bool = False) -> str:
        if not GDAL_AVAILABLE:
            raise ImportError("GDAL未安装，无法处理地理坐标转换")

        print(f"正在生成采样点地图...{' (快速模式)' if self.fast_mode else ''}")

        hyperspectral_img, geotransform, projection, width, height, sample_factor = self._read_hyperspectral(
            hyperspectral_path, rgb_bands, downsample)

        sampling_points = self._read_sampling_points(csv_path)
        rgb_image = self._create_false_color_image(hyperspectral_img)
        pixel_coords = self._geo_to_pixel(sampling_points, geotransform, width, height, projection, sample_factor)

        if output_filename is None:
            csv_name = Path(csv_path).stem
            hs_name = Path(hyperspectral_path).stem
            output_filename = f"{hs_name}_{csv_name}_sampling_map.png"

        output_path = self.output_dir / output_filename
        if dpi is None:
            dpi = 150 if self.fast_mode else 200

        self._create_map_visualization(
            rgb_image, pixel_coords, sampling_points, str(output_path), point_color, point_size, point_alpha,
            show_north_arrow, show_scale_bar, show_legend, dpi, geotransform, width, height, downsample, projection,
            sample_factor
        )

        print(f"采样点地图已保存: {output_path}")
        return str(output_path)

    def _read_hyperspectral(self, hyperspectral_path: str, rgb_bands: Optional[List[int]] = None,
                            downsample: bool = False) -> Tuple[np.ndarray, tuple, str, int, int]:
        dataset = gdal.Open(hyperspectral_path)
        if dataset is None:
            raise ValueError(f"无法打开高光谱影像: {hyperspectral_path}")

        width = dataset.RasterXSize
        height = dataset.RasterYSize
        band_count = dataset.RasterCount

        if rgb_bands is None:
            if band_count >= 3:
                try:
                    from src.utils.util import find_band_number
                    rgb_bands = [
                        find_band_number(650.0, hyperspectral_path),
                        find_band_number(550.0, hyperspectral_path),
                        find_band_number(460.0, hyperspectral_path)
                    ]
                except Exception:
                    rgb_bands = [min(band_count - 1, int(band_count * 0.25)),
                                 min(band_count - 1, int(band_count * 0.15)),
                                 min(band_count - 1, int(band_count * 0.05))]
            else:
                rgb_bands = [0, 0, 0]

        if downsample and (width > 2000 or height > 2000):
            print(f"  ⚠ 下采样暂被禁用，使用原始分辨率: {width}x{height}")
            sample_factor = 1
        else:
            sample_factor = 1

        rgb_data = []
        for band_idx in rgb_bands:
            band = dataset.GetRasterBand(band_idx + 1)
            band_data = band.ReadAsArray().astype(np.float32)
            rgb_data.append(band_data)

        if len(rgb_data) == 3:
            image_array = np.stack(rgb_data, axis=2)
        else:
            image_array = np.stack([rgb_data[0]] * 3, axis=2)

        geotransform = dataset.GetGeoTransform()
        projection = dataset.GetProjection()
        dataset = None

        return image_array, geotransform, projection, width, height, sample_factor

    def _read_sampling_points(self, csv_path: str) -> pd.DataFrame:
        """智能读取采样点，自动识别模糊列名，允许UTM坐标，自动修复颠倒坐标"""
        df = pd.read_csv(csv_path)
        if len(df.columns) < 2:
            raise ValueError("CSV文件至少需要两列（经度、纬度 或 X、Y）")

        # 智能子串匹配
        lat_aliases = ['lat', 'y', '纬']
        lon_aliases = ['lon', 'lng', 'x', '经']

        lat_col = None
        lon_col = None
        cols_lower = {c: str(c).strip().lower() for c in df.columns}

        for c, lc in cols_lower.items():
            if lat_col is None and any(a in lc for a in lat_aliases):
                lat_col = c
            elif lon_col is None and any(a in lc for a in lon_aliases):
                lon_col = c

        # 兜底：取前两列，默认列0=X(lon), 列1=Y(lat)
        if lat_col is None or lon_col is None:
            c0, c1 = df.columns[0], df.columns[1]
            lon_col, lat_col = c0, c1

        df = df.rename(columns={lat_col: 'latitude', lon_col: 'longitude'})
        df['latitude'] = pd.to_numeric(df['latitude'], errors='coerce')
        df['longitude'] = pd.to_numeric(df['longitude'], errors='coerce')
        n_nan = int(df[['latitude', 'longitude']].isna().any(axis=1).sum())
        df = df.dropna(subset=['latitude', 'longitude']).reset_index(drop=True)

        if len(df) > 0:
            lat_max = df['latitude'].abs().max()
            lon_max = df['longitude'].abs().max()

            # 智能对调：如果纬度 > 90，且经度 <= 90，说明用户把经纬度两列搞反了
            if lat_max > 90 and lon_max <= 90 and lat_max <= 180:
                print("  ⚠ 检测到经纬度数值颠倒 (纬度>90, 经度<=90)，系统已自动对调坐标列")
                df['latitude'], df['longitude'] = df['longitude'], df['latitude']
            # UTM 投影坐标判定：只要数值远大于180，就是米级别的投影系统
            elif lat_max > 180 or lon_max > 180:
                print(f"  ℹ 检测到坐标值远超180 (X:{lon_max:.1f}, Y:{lat_max:.1f})，判定为投影坐标(UTM)")

        print(f"  CSV 列匹配: lat_col='{lat_col}', lon_col='{lon_col}'")
        if n_nan:
            print(f"  剔除 {n_nan} 个无效(NaN)行")
        print(f"  读取到 {len(df)} 个有效采样点 (不再拦截越界拦截)")
        return df

    def _create_false_color_image(self, image_array: np.ndarray, rgb_bands: Optional[List[int]] = None) -> np.ndarray:
        if image_array.shape[2] != 3:
            if len(image_array.shape) == 2 or image_array.shape[2] == 1:
                image_array = np.stack([image_array] * 3, axis=2) if len(image_array.shape) == 2 else np.repeat(
                    image_array, 3, axis=2)

        def simple_linear_stretch(data, min_percent=1, max_percent=99):
            valid_data = data[np.isfinite(data)]
            if len(valid_data) == 0: return np.zeros_like(data, dtype=np.float32)
            p_low = np.percentile(valid_data, min_percent)
            p_high = np.percentile(valid_data, max_percent)
            if p_high - p_low < 1e-8:
                d_min, d_max = valid_data.min(), valid_data.max()
                return (data - d_min) / (d_max - d_min) if d_max > d_min else np.zeros_like(data, dtype=np.float32)
            stretched = (data - p_low) / (p_high - p_low)
            return np.clip(stretched, 0.0, 1.0)

        r_stretched = simple_linear_stretch(image_array[:, :, 0])
        g_stretched = simple_linear_stretch(image_array[:, :, 1])
        b_stretched = simple_linear_stretch(image_array[:, :, 2])
        rgb_image = np.nan_to_num(np.stack([r_stretched, g_stretched, b_stretched], axis=2), nan=0.0)
        rgb_image = np.clip(rgb_image, 0.0, 1.0)
        return (rgb_image * 255).astype(np.uint8)

    def _geo_to_pixel(self, sampling_points: pd.DataFrame, geotransform: tuple, width: int, height: int,
                      projection: str = "", sample_factor: int = 1) -> List[Tuple[float, float]]:
        if geotransform is None or len(sampling_points) == 0:
            return [(width / 2, height / 2) for _ in range(len(sampling_points))]

        pixel_coords = []
        gt = geotransform
        needs_transform = projection and ("PROJCS" in projection or "GEOGCS" in projection)

        # 智能判定是否为 WGS84
        sample_lon = float(sampling_points['longitude'].iloc[0])
        sample_lat = float(sampling_points['latitude'].iloc[0])
        is_wgs84 = (abs(sample_lon) <= 180) and (abs(sample_lat) <= 90)

        transform = None
        if needs_transform and is_wgs84 and GDAL_AVAILABLE:
            try:
                src_srs = osr.SpatialReference()
                src_srs.ImportFromEPSG(4326)
                dst_srs = osr.SpatialReference()
                dst_srs.ImportFromWkt(projection)
                transform = osr.CoordinateTransformation(src_srs, dst_srs)
            except Exception as e:
                transform = None
        elif not is_wgs84:
            print("  ℹ 采样点为投影坐标(UTM)，跳过WGS84投影转换，直接使用放射变换映射")

        for _, row in sampling_points.iterrows():
            lon, lat = float(row['longitude']), float(row['latitude'])

            if transform is not None:
                try:
                    proj_x, proj_y, _ = transform.TransformPoint(lon, lat)
                    x, y = (proj_x - gt[0]) / gt[1], (proj_y - gt[3]) / gt[5]
                except Exception:
                    x, y = width / 2, height / 2
            else:
                x, y = (lon - gt[0]) / gt[1], (lat - gt[3]) / gt[5]

            if sample_factor > 1:
                x, y = x / sample_factor, y / sample_factor

            pixel_coords.append((max(0, min(x, width - 1)), max(0, min(y, height - 1))))

        return pixel_coords

    def _create_map_visualization(self, rgb_image: np.ndarray, pixel_coords: List[Tuple[float, float]],
                                  sampling_points: pd.DataFrame, output_path: str, point_color: str, point_size: int,
                                  point_alpha: float, show_north_arrow: bool, show_scale_bar: bool, show_legend: bool,
                                  dpi: int, geotransform: tuple, width: int, height: int, downsample: bool = False,
                                  projection: str = "", sample_factor: int = 1):
        figsize = (10, 8) if self.fast_mode or downsample else (12, 10)
        fig, ax = plt.subplots(figsize=figsize, dpi=100 if self.fast_mode else 150)
        ax.imshow(rgb_image, interpolation='nearest' if self.fast_mode else 'bilinear')

        if pixel_coords:
            x_coords, y_coords = [p[0] for p in pixel_coords], [p[1] for p in pixel_coords]
            ax.scatter(x_coords, y_coords, c=point_color, s=point_size, alpha=point_alpha, edgecolors='white',
                       linewidth=1.5)

        if show_north_arrow: self._add_north_arrow(ax, width, height, position='bottom-left', direction='down')
        if show_scale_bar and geotransform is not None: self._add_scale_bar(ax, geotransform, width, height)

        if show_legend:
            ax.plot([], [], 'o', color=point_color, markersize=8, label=f'采样点 (n={len(sampling_points)})')
            ax.legend(loc='lower right', frameon=True, facecolor='white', edgecolor='gray')

        ax.set_title('高光谱影像采样点分布图', fontsize=16, fontweight='bold', pad=20)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.grid(True, alpha=0.2, linestyle='--')
        plt.tight_layout()

        save_kwargs = {'dpi': min(dpi, 180) if self.fast_mode else dpi, 'bbox_inches': 'tight', 'pad_inches': 0.05,
                       'facecolor': 'white'}
        plt.savefig(output_path, **save_kwargs)
        plt.close(fig)

    def _add_north_arrow(self, ax, width: int, height: int, position='top-right', direction='down', size=0.08,
                         color='white', n_color='white', outline_color='black'):
        pos_map = {'top-left': (0.08, 0.88), 'top-right': (0.92, 0.88), 'bottom-left': (0.08, 0.12),
                   'bottom-right': (0.92, 0.12)}
        arrow_x, arrow_y = width * pos_map.get(position, (0.92, 0.88))[0], height * pos_map.get(position, (0.92, 0.88))[
            1]
        dx, dy = {'up': (0, size), 'down': (0, -size), 'left': (-size, 0), 'right': (size, 0)}.get(direction,
                                                                                                   (0, -size))

        arrow = FancyArrowPatch((arrow_x, arrow_y), (arrow_x + dx * width, arrow_y + dy * height), color=color,
                                linewidth=3, arrowstyle='->', mutation_scale=20)
        ax.add_patch(arrow)
        text_y = arrow_y - height * 0.02 if direction == 'up' else arrow_y + height * 0.02
        ax.text(arrow_x, text_y, 'N', fontsize=14, fontweight='bold', color=n_color, ha='center', va='center',
                path_effects=[path_effects.withStroke(linewidth=3, foreground=outline_color)])

    def _add_scale_bar(self, ax, geotransform: tuple, width: int, height: int):
        if geotransform is None: return
        pixel_size_x = abs(geotransform[1])
        scale_length_m = (width * pixel_size_x) / 4
        scale_meters = next((s for s in [1000, 500, 200, 100, 50, 20, 10, 5, 2, 1] if s <= scale_length_m), 1)
        scale_pixels = int(scale_meters / pixel_size_x)
        bar_x, bar_y = width * 0.08, height * 0.92

        ax.plot([bar_x, bar_x + scale_pixels], [bar_y, bar_y], color='white', linewidth=4)
        ax.plot([bar_x, bar_x], [bar_y, bar_y + 8], color='white', linewidth=2)
        ax.plot([bar_x + scale_pixels, bar_x + scale_pixels], [bar_y, bar_y + 8], color='white', linewidth=2)
        ax.text(bar_x + scale_pixels / 2, bar_y, f'{scale_meters} m', fontsize=11, ha='center', va='bottom',
                fontweight='bold', bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=1))