74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
import os
|
|
import shutil
|
|
from osgeo import gdal
|
|
|
|
def fix_manual_masks(raw_dir, work_dir):
|
|
manual_files = [f for f in os.listdir(raw_dir) if f.endswith(".png")]
|
|
|
|
for manual_file in manual_files:
|
|
manual_path = os.path.join(raw_dir, manual_file)
|
|
|
|
parts = manual_file.replace("_viz.png", "").split("_")
|
|
if len(parts) >= 4:
|
|
tile_prefix = f"{parts[0]}_{parts[1]}_{parts[2]}_{parts[3]}"
|
|
work_file = f"{tile_prefix}_1_rp_mask.png"
|
|
work_path = os.path.join(work_dir, work_file)
|
|
|
|
if not os.path.exists(work_path):
|
|
print(f"Skipping {manual_file}: Work file not found.")
|
|
continue
|
|
|
|
print(f"Fixing {manual_file} using {work_file}...")
|
|
|
|
ds_work = gdal.Open(work_path)
|
|
gt_work = ds_work.GetGeoTransform()
|
|
proj_work = ds_work.GetProjection()
|
|
target_w = ds_work.RasterXSize
|
|
target_h = ds_work.RasterYSize
|
|
|
|
minx = gt_work[0]
|
|
maxy = gt_work[3]
|
|
maxx = minx + (gt_work[1] * target_w)
|
|
miny = maxy + (gt_work[5] * target_h)
|
|
|
|
# Open Source as ReadOnly, then CreateCopy to MEM to edit
|
|
ds_raw = gdal.Open(manual_path)
|
|
mem_driver = gdal.GetDriverByName("MEM")
|
|
ds_manual = mem_driver.CreateCopy("", ds_raw)
|
|
ds_raw = None # Close file
|
|
|
|
src_w = ds_manual.RasterXSize
|
|
src_h = ds_manual.RasterYSize
|
|
|
|
src_res_x = (maxx - minx) / src_w
|
|
src_res_y = (miny - maxy) / src_h
|
|
|
|
src_gt = (minx, src_res_x, 0, maxy, 0, src_res_y)
|
|
ds_manual.SetGeoTransform(src_gt)
|
|
ds_manual.SetProjection(proj_work)
|
|
|
|
warp_options = gdal.WarpOptions(
|
|
format="PNG",
|
|
width=target_w,
|
|
height=target_h,
|
|
outputBounds=(minx, miny, maxx, maxy),
|
|
resampleAlg=gdal.GRA_NearestNeighbour
|
|
)
|
|
|
|
temp_path = manual_path + ".tmp.png"
|
|
ds_fixed = gdal.Warp(temp_path, ds_manual, options=warp_options)
|
|
ds_fixed = None
|
|
ds_manual = None
|
|
ds_work = None
|
|
|
|
shutil.move(temp_path, manual_path)
|
|
|
|
work_wld = work_path.replace(".png", ".wld")
|
|
if os.path.exists(work_wld):
|
|
shutil.copy(work_wld, manual_path.replace(".png", ".wld"))
|
|
|
|
print(f" -> Fixed.")
|
|
|
|
if __name__ == "__main__":
|
|
fix_manual_masks("raw/water_masks", "work/river_masks")
|