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()) |