210 lines
7.0 KiB
Python
210 lines
7.0 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
掩膜转换工具模块
|
||
|
||
提供 shapefile / ndarray / dat / tif 等多种格式掩膜之间的相互转换，
|
||
以及水体掩膜的预处理逻辑。
|
||
"""
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Optional, Union
|
||
|
||
import numpy as np
|
||
|
||
try:
|
||
from osgeo import gdal, ogr
|
||
GDAL_AVAILABLE = True
|
||
except ImportError:
|
||
GDAL_AVAILABLE = False
|
||
|
||
|
||
def prepare_water_mask_for_algorithm(
|
||
water_mask: Optional[Union[str, np.ndarray]],
|
||
image_shape: Union[tuple, np.ndarray],
|
||
geotransform: tuple,
|
||
projection: str,
|
||
img_path: str,
|
||
water_mask_dir: Optional[str] = None,
|
||
callback=None
|
||
) -> Optional[np.ndarray]:
|
||
"""
|
||
准备水域掩膜供算法使用
|
||
|
||
支持格式：
|
||
- None：自动使用预先生成的 dat 格式掩膜
|
||
- numpy.ndarray：直接返回（确保是 0/1 格式）
|
||
- .dat / .tif 等栅格文件：读取并返回
|
||
- .shp 文件：先栅格化，再读取返回
|
||
|
||
Args:
|
||
water_mask: 掩膜来源
|
||
image_shape: 影像形状 (height, width) 或 (height, width, channels)
|
||
geotransform: GDAL 地理变换参数
|
||
projection: 投影信息
|
||
img_path: 影像路径（用于 shp 栅格化）
|
||
water_mask_dir: 水体掩膜目录（用于缓存栅格化的 shp 结果）
|
||
callback: 进度回调函数（可选）
|
||
|
||
Returns:
|
||
numpy数组（dtype=uint8，0=非水域，1=水域）或 None
|
||
"""
|
||
img_height, img_width = image_shape[0], image_shape[1]
|
||
|
||
if water_mask is None:
|
||
return None
|
||
|
||
# numpy 数组直接返回
|
||
if isinstance(water_mask, np.ndarray):
|
||
if water_mask.shape[:2] != (img_height, img_width):
|
||
raise ValueError(f"掩膜尺寸 {water_mask.shape[:2]} 与图像尺寸 {(img_height, img_width)} 不匹配")
|
||
return (water_mask > 0).astype(np.uint8)
|
||
|
||
# 字符串路径
|
||
if isinstance(water_mask, str):
|
||
ext = Path(water_mask).suffix.lower()
|
||
|
||
# shapefile 格式
|
||
if ext == '.shp':
|
||
return _convert_shp_to_mask(
|
||
shp_path=water_mask,
|
||
img_path=img_path,
|
||
image_shape=image_shape,
|
||
geotransform=geotransform,
|
||
projection=projection,
|
||
water_mask_dir=water_mask_dir,
|
||
callback=callback
|
||
)
|
||
|
||
# 栅格文件格式
|
||
return _load_raster_mask(water_mask, img_height, img_width)
|
||
|
||
raise ValueError(f"不支持的掩膜类型: {type(water_mask)}")
|
||
|
||
|
||
def _convert_shp_to_mask(shp_path: str, img_path: str,
|
||
image_shape: tuple,
|
||
geotransform: tuple,
|
||
projection: str,
|
||
water_mask_dir: Optional[str] = None,
|
||
callback=None) -> np.ndarray:
|
||
"""将 shapefile 栅格化为掩膜数组"""
|
||
from src.utils.extract_water_area import rasterize_shp
|
||
|
||
safe_shp_path = os.path.abspath(shp_path).replace('\\', '/')
|
||
shp_name = Path(safe_shp_path).stem
|
||
|
||
if water_mask_dir:
|
||
temp_mask_path = str(Path(water_mask_dir) / f"water_mask_{shp_name}.dat")
|
||
else:
|
||
temp_mask_path = f"/tmp/water_mask_{shp_name}.dat"
|
||
|
||
# 缓存：已栅格化则直接读取
|
||
if Path(temp_mask_path).exists():
|
||
print(f"使用已存在的栅格化掩膜: {temp_mask_path}")
|
||
return _load_raster_mask(temp_mask_path, image_shape[0], image_shape[1])
|
||
|
||
# 需要栅格化
|
||
if img_path is None:
|
||
raise ValueError("当 water_mask 为 shp 格式时，需要提供 img_path 参数用于栅格化")
|
||
|
||
print(f"正在将 SHP 栅格化: {safe_shp_path}")
|
||
rasterize_shp(safe_shp_path, temp_mask_path, img_path)
|
||
|
||
return _load_raster_mask(temp_mask_path, image_shape[0], image_shape[1])
|
||
|
||
|
||
def _load_raster_mask(mask_path: str, img_height: int, img_width: int) -> np.ndarray:
|
||
"""从栅格文件加载掩膜"""
|
||
if not GDAL_AVAILABLE:
|
||
raise ImportError("GDAL未安装，无法读取掩膜文件")
|
||
|
||
mask_dataset = gdal.Open(mask_path, gdal.GA_ReadOnly)
|
||
if mask_dataset is None:
|
||
raise ValueError(f"无法打开掩膜文件: {mask_path}")
|
||
|
||
try:
|
||
mask_array = mask_dataset.GetRasterBand(1).ReadAsArray()
|
||
finally:
|
||
mask_dataset = None
|
||
|
||
if mask_array.shape != (img_height, img_width):
|
||
raise ValueError(f"掩膜尺寸 {mask_array.shape} 与图像尺寸 {(img_height, img_width)} 不匹配")
|
||
|
||
return (mask_array > 0).astype(np.uint8)
|
||
|
||
|
||
def ensure_water_mask_dat(img_path: str,
|
||
existing_dat_path: Optional[str] = None,
|
||
output_dir: Optional[str] = None) -> str:
|
||
"""
|
||
确保存在 dat 格式的水体掩膜文件（用于步骤3/4中的算法）
|
||
|
||
如果 existing_dat_path 存在且是 .dat 文件，直接返回。
|
||
如果存在同名 .dat 文件，直接返回。
|
||
否则从 img_path 生成并保存到 output_dir。
|
||
|
||
Args:
|
||
img_path: 用于生成掩膜的遥感影像路径
|
||
existing_dat_path: 已有的 dat 格式掩膜路径（可选）
|
||
output_dir: 输出目录（可选）
|
||
|
||
Returns:
|
||
dat 格式掩膜文件路径
|
||
"""
|
||
if existing_dat_path and Path(existing_dat_path).suffix.lower() == '.dat':
|
||
if Path(existing_dat_path).exists():
|
||
return existing_dat_path
|
||
|
||
img_name = Path(img_path).stem
|
||
if output_dir is None:
|
||
output_dir = str(Path(img_path).parent)
|
||
|
||
dat_path = str(Path(output_dir) / f"{img_name}_water_mask.dat")
|
||
|
||
if Path(dat_path).exists():
|
||
return dat_path
|
||
|
||
# 如果已有其他格式的掩膜，转换为 dat
|
||
for ext in ['.tif', '.img', '.tiff']:
|
||
alt_path = str(Path(output_dir) / f"{img_name}_water_mask{ext}")
|
||
if Path(alt_path).exists():
|
||
return _convert_to_dat(alt_path, dat_path)
|
||
|
||
return dat_path # 返回目标路径，让调用方决定是否需要生成
|
||
|
||
|
||
def _convert_to_dat(src_path: str, dest_path: str) -> str:
|
||
"""将其他栅格格式转换为 ENVI dat 格式"""
|
||
if not GDAL_AVAILABLE:
|
||
raise ImportError("GDAL未安装，无法转换格式")
|
||
|
||
src_ds = gdal.Open(src_path, gdal.GA_ReadOnly)
|
||
if src_ds is None:
|
||
raise ValueError(f"无法打开源掩膜文件: {src_path}")
|
||
|
||
try:
|
||
geotransform = src_ds.GetGeoTransform()
|
||
projection = src_ds.GetProjection()
|
||
band = src_ds.GetRasterBand(1)
|
||
array = band.ReadAsArray()
|
||
|
||
driver = gdal.GetDriverByName('ENVI')
|
||
if driver is None:
|
||
driver = gdal.GetDriverByName('GTiff')
|
||
|
||
dest_ds = driver.Create(dest_path, src_ds.RasterXSize, src_ds.RasterYSize, 1, gdal.GDT_Byte)
|
||
if dest_ds is None:
|
||
raise ValueError(f"无法创建输出文件: {dest_path}")
|
||
|
||
try:
|
||
dest_ds.SetGeoTransform(geotransform)
|
||
dest_ds.SetProjection(projection)
|
||
dest_band = dest_ds.GetRasterBand(1)
|
||
dest_band.WriteArray((array > 0).astype(np.uint8))
|
||
dest_band.FlushCache()
|
||
finally:
|
||
dest_ds = None
|
||
|
||
return dest_path
|
||
finally:
|
||
src_ds = None |