131 lines
4.1 KiB
Python
131 lines
4.1 KiB
Python
|
"""
|
|||
|
|
掩膜tif文件的ROI区域并保存
|
|||
|
|
|
|||
|
|
输入: tif文件路径, roi文件路径(shp/geojson等)
|
|||
|
|
输出: 掩膜后的tif文件
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import numpy as np
|
|||
|
|
import rasterio
|
|||
|
|
from rasterio.mask import mask
|
|||
|
|
from rasterio.features import shapes
|
|||
|
|
import geopandas as gpd
|
|||
|
|
# from shapely.geometry import shape
|
|||
|
|
import logging
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def mask_tif_by_roi(tif_path: str, roi_path: str, output_path: str = None,
|
|||
|
|
nodata_value: float = None):
|
|||
|
|
"""
|
|||
|
|
使用ROI文件掩膜TIF文件
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
tif_path: 输入TIF文件路径
|
|||
|
|
roi_path: ROI文件路径（shp/geojson等）
|
|||
|
|
output_path: 输出文件路径，如果为None则自动生成
|
|||
|
|
nodata_value: nodata值，如果为None则使用原始文件的nodata值
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
str: 输出文件路径
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
tif_path = Path(tif_path)
|
|||
|
|
roi_path = Path(roi_path)
|
|||
|
|
|
|||
|
|
# 检查输入文件存在性
|
|||
|
|
if not tif_path.exists():
|
|||
|
|
raise FileNotFoundError(f"TIF文件不存在: {tif_path}")
|
|||
|
|
if not roi_path.exists():
|
|||
|
|
raise FileNotFoundError(f"ROI文件不存在: {roi_path}")
|
|||
|
|
|
|||
|
|
# 自动生成输出路径
|
|||
|
|
if output_path is None:
|
|||
|
|
output_path = tif_path.parent / f"{tif_path.stem}_masked{tif_path.suffix}"
|
|||
|
|
else:
|
|||
|
|
output_path = Path(output_path)
|
|||
|
|
|
|||
|
|
logger.info(f"开始处理: {tif_path.name}")
|
|||
|
|
logger.info(f"ROI文件: {roi_path.name}")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 读取ROI文件
|
|||
|
|
logger.info("读取ROI文件...")
|
|||
|
|
gdf = gpd.read_file(roi_path)
|
|||
|
|
|
|||
|
|
if gdf.empty:
|
|||
|
|
logger.warning("ROI文件为空，直接复制原文件")
|
|||
|
|
# 直接复制原文件
|
|||
|
|
import shutil
|
|||
|
|
shutil.copy2(tif_path, output_path)
|
|||
|
|
return str(output_path)
|
|||
|
|
|
|||
|
|
# 转换为GeoJSON格式的几何对象
|
|||
|
|
geometries = gdf.geometry.tolist()
|
|||
|
|
logger.info(f"找到 {len(geometries)} 个ROI几何对象")
|
|||
|
|
|
|||
|
|
# 读取并掩膜TIF
|
|||
|
|
logger.info("读取并掩膜TIF文件...")
|
|||
|
|
with rasterio.open(tif_path) as src:
|
|||
|
|
# 使用ROI掩膜 - 保留ROI以外的区域（ROI区域设为nodata）
|
|||
|
|
masked_data, masked_transform = mask(
|
|||
|
|
src,
|
|||
|
|
geometries,
|
|||
|
|
crop=False, # 不裁剪，保持原始尺寸
|
|||
|
|
invert=True, # 反转：ROI区域设为nodata，ROI以外保留
|
|||
|
|
nodata=nodata_value if nodata_value is not None else src.nodata
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 更新元数据
|
|||
|
|
out_meta = src.meta.copy()
|
|||
|
|
out_meta.update({
|
|||
|
|
"driver": "GTiff",
|
|||
|
|
"height": masked_data.shape[1],
|
|||
|
|
"width": masked_data.shape[2],
|
|||
|
|
"transform": masked_transform,
|
|||
|
|
"nodata": nodata_value if nodata_value is not None else src.nodata
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 保存结果
|
|||
|
|
logger.info(f"保存掩膜结果到: {output_path.name}")
|
|||
|
|
with rasterio.open(output_path, "w", **out_meta) as dest:
|
|||
|
|
dest.write(masked_data)
|
|||
|
|
|
|||
|
|
logger.info("处理完成！")
|
|||
|
|
return str(output_path)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"处理失败: {str(e)}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
parser = argparse.ArgumentParser(description="使用ROI文件掩膜TIF文件")
|
|||
|
|
parser.add_argument("tif_path", help="输入TIF文件路径")
|
|||
|
|
parser.add_argument("roi_path", help="ROI文件路径 (shp/geojson等)")
|
|||
|
|
parser.add_argument("-o", "--output", help="输出文件路径 (可选，默认自动生成)")
|
|||
|
|
parser.add_argument("-n", "--nodata", type=float, help="nodata值 (可选，默认使用原文件nodata)")
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
output_path = mask_tif_by_roi(
|
|||
|
|
args.tif_path,
|
|||
|
|
args.roi_path,
|
|||
|
|
args.output,
|
|||
|
|
args.nodata
|
|||
|
|
)
|
|||
|
|
print(f"成功完成！输出文件: {output_path}")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"错误: {e}")
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
exit(main())
|