不规则裁剪与统计

# -*- coding:utf-8 -*-
import json
import rasterio
from uuid import uuid1
from rasterio.crs import CRS
from rasterio.transform import from_origin
from rasterio.features import shapes
from rasterio.mask import mask
from typing import Union, List, Dict
import shapely.ops
from shapely.geometry import Polygon, MultiPoint, box
import geopandas as gpd
import random
import numpy as np


class PolygonValPixelCounter:
    """一个用于统计多边形范围内栅格图像象元个数的类"""
    def __init__(self, raster_path: str):
        """
        初始化类
        
        Args:
            raster_path (str): 栅格文件路径
        """
        self.raster_path = raster_path
        self.raster = None
        self.nodata = 255  # 明确指定nodata值为255
        self._load_raster()
    
    def _load_raster(self) -> None:
        """加载栅格数据"""
        try:
            self.raster = rasterio.open(self.raster_path)
        except Exception as e:
            raise ValueError(f"无法加载栅格文件: {str(e)}")
    
    def count_pixels(self, 
                    polygons: Union[Dict, List, gpd.GeoDataFrame],
                    band: int = 1) -> int:
        """
        统计多边形范围内的有效象元个数
        Args:
            polygon: 多边形几何
            band (int): 要统计的波段，默认为1
        Returns:
            int: 有效象元个数
        """
        if self.raster is None:
            raise ValueError("栅格数据未加载")
            
        try:
            # 裁剪多边形范围内的栅格数据，指定nodata为255
            
            res = []
            for ig in polygons:
                out_image, _ = mask(
                    self.raster,
                    [ig],
                    crop=True,
                    nodata=self.nodata  # 使用指定的nodata值255
                )
                
                # 获取指定波段数据并统计有效象元个数
                data = out_image[band-1]
                valid_pixels = data[data != self.nodata]  # 使用255作为nodata值
                res.append(int(valid_pixels.size))
                
            return res
            
        except Exception as e:
            raise ValueError(f"统计失败: {str(e)}")
    
    def close(self) -> None:
        """关闭栅格文件"""
        if self.raster is not None:
            self.raster.close()
            self.raster = None
    
    def __del__(self):
        """析构函数，确保文件被关闭"""
        self.close()


def split_polygon(polygon, num_parts):
    """
    将一个 polygon 随机拆分为指定数量的小块。
    参数：
        polygon (shapely.geometry.Polygon): 要拆分的多边形。
        num_parts (int): 拆分的小块数量。
    返回：
        list: 包含拆分后的多边形的列表。
    """
    # 获取多边形的边界框
    min_x, min_y, max_x, max_y = polygon.bounds

    # 生成随机点
    points = []
    for _ in range(num_parts):
        while True:
            x = random.uniform(min_x, max_x)
            y = random.uniform(min_y, max_y)
            point = shapely.geometry.Point(x, y)
            if polygon.contains(point):
                points.append(point)
                break

    # 添加多边形的中心点
    points.append(polygon.centroid)

    # 使用 Voronoi 图将多边形拆分为小块
    voronoi_diagram = shapely.ops.voronoi_diagram(shapely.geometry.MultiPoint(points))
    polygons = [polygon.intersection(cell) for cell in voronoi_diagram.geoms]

    return polygons


def raster_to_polygon(raster_file):
    """
    提取栅格图像的范围
    参数：raster_file (str): 栅格文件路径。
    返回：list: 包含多边形特征的列表。
    """
    with rasterio.open(raster_file) as src:
        result = box(src.bounds.left, src.bounds.bottom, src.bounds.right, src.bounds.top)

    return result


def create_geotiff(data, output_file, crs_epsg, origin_x, origin_y, pixel_size, nodata=None):
    """
    创建一个带空间参考的 GeoTIFF 文件。
    参数:
        data (numpy.ndarray): 栅格数据，二维数组。
        output_file (str): 输出的 GeoTIFF 文件路径。
        crs_epsg (int): 坐标系的 EPSG 代码。
        origin_x (float): 左上角的 X 坐标。
        origin_y (float): 左上角的 Y 坐标。
        pixel_size (float): 像元大小。
        nodata (float, optional): 无数据值，默认为 None。
    """
    # 设置空间参考信息
    crs = CRS.from_epsg(crs_epsg)

    # 设置仿射变换
    transform = from_origin(origin_x, origin_y, pixel_size, pixel_size)

    # 设置其他元数据
    metadata = {
        'driver': 'GTiff',  # 驱动程序，指定为 GeoTIFF
        'dtype': data.dtype,  # 数据类型
        'nodata': nodata,  # 无数据值
        'width': data.shape[1],  # 宽度（列数）
        'height': data.shape[0],  # 高度（行数）
        'count': 1,  # 波段数
        'crs': crs,  # 坐标参考系统
        'transform': transform  # 仿射变换
    }

    # 保存到 GeoTIFF 文件
    with rasterio.open(output_file, 'w', **metadata) as dst:
        dst.write(data, 1)  # 写入数据到第一个波段

    print(f'GeoTIFF 文件已保存到 {output_file}')


# 示例调用
if __name__ == "__main__":
    # 示例数据：一个简单的二维数组
    data = np.random.randint(1, 100, (100,100))  # 随机生成一个 100x100 的数组，值范围为 0 到 100
    out_file = F'tmpfile/example_{str(uuid1())[:8]}.tif'
    # 调用函数创建 GeoTIFF 文件
    create_geotiff(
        data=data,
        output_file=out_file,
        crs_epsg=4326,  # WGS84
        origin_x=113.475,  # 左上角 X 坐标
        origin_y=34.85,  # 左上角 Y 坐标
        pixel_size=0.001,  # 像元大小
        nodata=0  # 无数据值
    )

    gg = raster_to_polygon(out_file)
    gs = split_polygon(gg, num_parts=5)
    # gs = gpd.GeoSeries(gs, crs=4326).to_file(out_file.replace('.tif', '.geojson'))

    pvpc = PolygonValPixelCounter(out_file)
    res = pvpc.count_pixels(gs)
    print(res)