feat: 水质分析系统用户体验核心升级
This commit is contained in:
@ -242,20 +242,21 @@ class WaterQualityInversionPipeline:
|
||||
ndwi_threshold: float = 0.4,
|
||||
use_ndwi: bool = False,
|
||||
skip_dependency_check: bool = False,
|
||||
generate_png: bool = True) -> str:
|
||||
generate_png: bool = True,
|
||||
output_path: Optional[str] = None) -> str:
|
||||
"""
|
||||
步骤1: 生成或设置水域mask
|
||||
|
||||
|
||||
支持三种方式生成水域掩膜：
|
||||
1. 基于shp文件栅格化
|
||||
2. 使用现有的栅格格式掩膜文件
|
||||
3. 基于NDWI从影像自动生成水体掩膜
|
||||
|
||||
|
||||
当提供img_path时，会自动生成PNG预览图，基于波长选择RGB波段：
|
||||
- 红波段: ~650nm
|
||||
- 绿波段: ~550nm
|
||||
- 蓝波段: ~460nm
|
||||
|
||||
|
||||
Args:
|
||||
mask_path: 水体掩膜文件路径，支持：
|
||||
- shp格式文件（.shp）：需要提供img_path用于栅格化
|
||||
@ -265,7 +266,9 @@ class WaterQualityInversionPipeline:
|
||||
ndwi_threshold: NDWI阈值（当use_ndwi=True时使用）
|
||||
use_ndwi: 是否使用NDWI方法从影像生成水体掩膜
|
||||
generate_png: 是否生成输入影像的PNG预览图（默认True）
|
||||
|
||||
output_path: 指定输出掩膜文件的保存路径（可选）。如果提供，掩膜将保存到此路径；
|
||||
如果为None，则使用默认路径（self.water_mask_dir）
|
||||
|
||||
Returns:
|
||||
dat格式的水域掩膜文件路径
|
||||
"""
|
||||
@ -285,37 +288,44 @@ class WaterQualityInversionPipeline:
|
||||
raise ValueError("当use_ndwi=True时，必须提供img_path参数用于生成NDWI掩膜")
|
||||
if not Path(img_path).exists():
|
||||
raise ValueError(f"影像文件不存在: {img_path}")
|
||||
|
||||
|
||||
print(f"使用NDWI方法从影像生成水体掩膜，阈值={ndwi_threshold}...")
|
||||
output_path = str(self.water_mask_dir / "water_mask_from_ndwi.dat")
|
||||
|
||||
# 使用用户指定的输出路径，或使用默认路径
|
||||
if output_path:
|
||||
ndwi_output_path = output_path
|
||||
# 确保输出目录存在
|
||||
os.makedirs(Path(output_path).parent, exist_ok=True)
|
||||
else:
|
||||
ndwi_output_path = str(self.water_mask_dir / "water_mask_from_ndwi.dat")
|
||||
|
||||
# 检查文件是否已存在，避免重复生成
|
||||
if Path(output_path).exists():
|
||||
print(f"检测到已存在的NDWI掩膜文件，直接使用: {output_path}")
|
||||
self.water_mask_path = output_path
|
||||
if Path(ndwi_output_path).exists():
|
||||
print(f"检测到已存在的NDWI掩膜文件，直接使用: {ndwi_output_path}")
|
||||
self.water_mask_path = ndwi_output_path
|
||||
step_end_time = time.time()
|
||||
self._record_step_time("步骤1: 生成水域mask", step_start_time, step_end_time, status="skipped")
|
||||
print(f"水域掩膜已设置: {self.water_mask_path}")
|
||||
|
||||
|
||||
# 生成水域掩膜叠加图（如果不存在）
|
||||
overlay_path = self.water_mask_dir / "water_mask_overlay.png"
|
||||
if generate_png and img_path is not None and Path(img_path).exists() and not overlay_path.exists():
|
||||
self._generate_water_mask_overlay(img_path, self.water_mask_path)
|
||||
|
||||
|
||||
return self.water_mask_path
|
||||
|
||||
|
||||
# 执行NDWI水体提取
|
||||
from src.utils.extract_water_area import ndwi
|
||||
ndwi(img_path, ndwi_threshold, output_path)
|
||||
self.water_mask_path = output_path
|
||||
ndwi(img_path, ndwi_threshold, ndwi_output_path)
|
||||
self.water_mask_path = ndwi_output_path
|
||||
step_end_time = time.time()
|
||||
self._record_step_time("步骤1: 生成水域mask", step_start_time, step_end_time)
|
||||
print(f"已生成NDWI水体掩膜: {self.water_mask_path}")
|
||||
|
||||
|
||||
# 生成水域掩膜叠加图
|
||||
if generate_png:
|
||||
self._generate_water_mask_overlay(img_path, self.water_mask_path)
|
||||
|
||||
|
||||
return self.water_mask_path
|
||||
|
||||
elif mask_path is None:
|
||||
@ -331,37 +341,44 @@ class WaterQualityInversionPipeline:
|
||||
# 如果是shp文件，需要栅格化为dat
|
||||
if img_path is None:
|
||||
raise ValueError("当mask_path为shp格式时，必须提供img_path参数用于栅格化")
|
||||
|
||||
|
||||
print(f"检测到shp格式的水体掩膜，正在转换为dat格式...")
|
||||
output_path = str(self.water_mask_dir / "water_mask_from_shp.dat")
|
||||
|
||||
# 使用用户指定的输出路径，或使用默认路径
|
||||
if output_path:
|
||||
shp_output_path = output_path
|
||||
# 确保输出目录存在
|
||||
os.makedirs(Path(output_path).parent, exist_ok=True)
|
||||
else:
|
||||
shp_output_path = str(self.water_mask_dir / "water_mask_from_shp.dat")
|
||||
|
||||
# 检查文件是否已存在，避免重复栅格化
|
||||
if Path(output_path).exists():
|
||||
print(f"检测到已存在的栅格化掩膜文件，直接使用: {output_path}")
|
||||
self.water_mask_path = output_path
|
||||
if Path(shp_output_path).exists():
|
||||
print(f"检测到已存在的栅格化掩膜文件，直接使用: {shp_output_path}")
|
||||
self.water_mask_path = shp_output_path
|
||||
step_end_time = time.time()
|
||||
self._record_step_time("步骤1: 生成水域mask", step_start_time, step_end_time, status="skipped")
|
||||
print(f"水域掩膜已设置: {self.water_mask_path}")
|
||||
|
||||
|
||||
# 生成水域掩膜叠加图（如果不存在）
|
||||
overlay_path = self.water_mask_dir / "water_mask_overlay.png"
|
||||
if generate_png and img_path is not None and Path(img_path).exists() and not overlay_path.exists():
|
||||
self._generate_water_mask_overlay(img_path, self.water_mask_path)
|
||||
|
||||
|
||||
return self.water_mask_path
|
||||
|
||||
|
||||
# 执行栅格化
|
||||
from src.utils.extract_water_area import rasterize_shp
|
||||
rasterize_shp(mask_path, output_path, img_path)
|
||||
self.water_mask_path = output_path
|
||||
rasterize_shp(mask_path, shp_output_path, img_path)
|
||||
self.water_mask_path = shp_output_path
|
||||
step_end_time = time.time()
|
||||
self._record_step_time("步骤1: 生成水域mask", step_start_time, step_end_time)
|
||||
print(f"已生成dat格式的水域掩膜: {self.water_mask_path}")
|
||||
|
||||
|
||||
# 生成水域掩膜叠加图
|
||||
if generate_png:
|
||||
self._generate_water_mask_overlay(img_path, self.water_mask_path)
|
||||
|
||||
|
||||
return self.water_mask_path
|
||||
|
||||
else:
|
||||
@ -483,13 +500,13 @@ class WaterQualityInversionPipeline:
|
||||
fontsize=12, fontweight='bold')
|
||||
ax.axis('off')
|
||||
|
||||
# 添加比例尺信息
|
||||
# 添加比例尺信息（白色文字，黑色背景下清晰可见）
|
||||
geo_transform = dataset.GetGeoTransform()
|
||||
if geo_transform:
|
||||
pixel_size_x = abs(geo_transform[1])
|
||||
pixel_size_y = abs(geo_transform[5])
|
||||
scale_text = f"分辨率: {pixel_size_x:.2f}m x {pixel_size_y:.2f}m | 尺寸: {width} x {height}"
|
||||
fig.text(0.5, 0.02, scale_text, ha='center', fontsize=10, style='italic')
|
||||
fig.text(0.5, 0.02, scale_text, ha='center', fontsize=10, style='italic', color='white')
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(png_path, dpi=150, bbox_inches='tight', pad_inches=0.1)
|
||||
|
||||
Reference in New Issue
Block a user