diff --git a/pyproject.toml b/pyproject.toml index cc4cdebb..2ca933af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,6 +90,7 @@ ignore = [ "PTH123", # `open()` should be replaced by `Path.open()` "PTH207", # "Replace `glob` with `Path.glob` or `Path.rglob` "ISC001", # The following rules may cause conflicts when used with the formatter + "TRY003", # Avoid specifying long messages outside the exception ] exclude = ["scripts"] diff --git a/src/dolphin/_overviews.py b/src/dolphin/_overviews.py new file mode 100644 index 00000000..a53b4589 --- /dev/null +++ b/src/dolphin/_overviews.py @@ -0,0 +1,175 @@ +import argparse +import logging +from enum import Enum +from os import fspath +from pathlib import Path +from typing import Sequence + +from osgeo import gdal +from tqdm.contrib.concurrent import thread_map + +from dolphin._types import PathOrStr + +gdal.UseExceptions() + +logger = logging.getLogger(__name__) + +DEFAULT_LEVELS = [4, 8, 16, 32, 64] + + +class Resampling(Enum): + """GDAL resampling algorithm.""" + + NEAREST = "nearest" + AVERAGE = "average" + + +class ImageType(Enum): + """Types of images produced by dolphin.""" + + UNWRAPPED = "unwrapped" + INTERFEROGRAM = "interferogram" + CORRELATION = "correlation" + CONNCOMP = "conncomp" + PS = "ps" + + +IMAGE_TYPE_TO_RESAMPLING = { + ImageType.UNWRAPPED: Resampling.AVERAGE, + ImageType.INTERFEROGRAM: Resampling.AVERAGE, + ImageType.CORRELATION: Resampling.AVERAGE, + ImageType.CONNCOMP: Resampling.NEAREST, + # No max in resampling, yet, which would be best + # https://github.com/OSGeo/gdal/issues/3683 + ImageType.PS: Resampling.AVERAGE, +} + + +def create_image_overviews( + file_path: Path | str, + levels: Sequence[int] = DEFAULT_LEVELS, + image_type: ImageType | None = None, + resampling: Resampling | None = None, + external: bool = False, + compression: str = "LZW", +): + """Add GDAL compressed overviews to an existing file. + + Parameters + ---------- + file_path : Path + Path to the file to process. + levels : Sequence[int] + List of overview levels to add. + Default = [4, 8, 16, 32, 64] + image_type : ImageType, optional + If provided, looks up the default resampling algorithm + most appropriate for this type of raster. + resampling : str or Resampling + GDAL resampling algorithm for overviews. Required + if not specifying `image_type`. + external : bool, default = False + Use external overviews (.ovr files). + compression: str, default = "LZW" + Compression algorithm to use for overviews. + See https://gdal.org/programs/gdaladdo.html for options. + + """ + if image_type is None and resampling is None: + raise ValueError("Must provide `image_type` or `resampling`") + if image_type is not None: + resampling = IMAGE_TYPE_TO_RESAMPLING[ImageType(image_type)] + else: + resampling = Resampling(resampling) + + flags = gdal.GA_Update if not external else gdal.GA_ReadOnly + ds = gdal.Open(fspath(file_path), flags) + if ds.GetRasterBand(1).GetOverviewCount() > 0: + logger.info("%s already has overviews. Skipping.") + return + + gdal.SetConfigOption("COMPRESS_OVERVIEW", compression) + gdal.SetConfigOption("GDAL_NUM_THREADS", "2") + ds.BuildOverviews(resampling.value, levels) + + +def create_overviews( + file_paths: Sequence[PathOrStr], + levels: Sequence[int] = DEFAULT_LEVELS, + image_type: ImageType | None = None, + resampling: Resampling = Resampling.AVERAGE, + max_workers: int = 5, +) -> None: + """Process many files to add GDAL overviews and compression. + + Parameters + ---------- + file_paths : Sequence[PathOrStr] + Sequence of file paths to process. + levels : Sequence[int] + Sequence of overview levels to add. + Default = [4, 8, 16, 32, 64] + image_type : ImageType, optional + If provided, looks up the default resampling algorithm + resampling : str or Resampling + GDAL resampling algorithm for overviews. Required + if not specifying `image_type`. + max_workers : int, default = 5 + Number of parallel threads to run. + + """ + thread_map( + lambda file_path: create_image_overviews( + Path(file_path), + levels=list(levels), + image_type=image_type, + resampling=resampling, + ), + file_paths, + max_workers=max_workers, + ) + + +def run(): + """Add compressed GDAL overviews to files.""" + parser = argparse.ArgumentParser( + description="Add compressed GDAL overviews to files.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("file_paths", nargs="+", type=str, help="Path to files") + parser.add_argument( + "--levels", + "-l", + nargs="*", + default=[4, 8, 16, 32, 64], + type=int, + help="Overview levels to add.", + ) + parser.add_argument( + "--resampling", + "-r", + default=Resampling("nearest"), + choices=[r.value for r in Resampling], + type=Resampling, + help="Resampling algorithm to use when building overviews", + ) + parser.add_argument( + "--max-workers", + "-n", + default=5, + type=int, + help="Number of parallel files to process", + ) + + args = parser.parse_args() + + # Convert resampling argument from string to Resampling Enum + resampling_enum = Resampling(args.resampling) + + create_overviews( + file_paths=args.file_paths, + levels=args.levels, + resampling=resampling_enum, + max_workers=args.max_workers, + overwrite=args.overwrite, + ) diff --git a/src/dolphin/io/_readers.py b/src/dolphin/io/_readers.py index d5923a74..96bb7540 100644 --- a/src/dolphin/io/_readers.py +++ b/src/dolphin/io/_readers.py @@ -1051,7 +1051,7 @@ def iter_blocks( self.queue_read(rows, cols) queued_slices.append((rows, cols)) - logger.info(f"Processing {self._block_shape} sized blocks... {tqdm_kwargs}") + logger.info(f"Processing {self._block_shape} sized blocks...") for _ in trange(len(queued_slices), **tqdm_kwargs): cur_block, (rows, cols) = self.get_data() logger.debug(f"got data for {rows, cols}: {cur_block.shape}") diff --git a/src/dolphin/unwrap/_unwrap.py b/src/dolphin/unwrap/_unwrap.py index 0767ae63..041a2045 100644 --- a/src/dolphin/unwrap/_unwrap.py +++ b/src/dolphin/unwrap/_unwrap.py @@ -99,8 +99,10 @@ def run( output_path = Path(output_path) + ifg_suffixes = [full_suffix(f) for f in ifg_filenames] all_out_files = [ - (output_path / Path(f).name).with_suffix(UNW_SUFFIX) for f in ifg_filenames + (output_path / Path(f).name.replace(suf, UNW_SUFFIX)) + for f, suf in zip(ifg_filenames, ifg_suffixes) ] in_files, out_files = [], [] for inf, outf in zip(ifg_filenames, all_out_files): @@ -137,7 +139,8 @@ def run( for ifg_file, out_file, cor_file in zip(in_files, out_files, cor_filenames) ] for fut in tqdm(as_completed(futures)): - fut.result() + # We're not passing all the unw files in, so we need to tally up below + _unw_path, _cc_path = fut.result() conncomp_files = [ Path(str(outf).replace(UNW_SUFFIX, CONNCOMP_SUFFIX)) for outf in all_out_files diff --git a/src/dolphin/workflows/config/_common.py b/src/dolphin/workflows/config/_common.py index a3347e2f..1f8d7fe3 100644 --- a/src/dolphin/workflows/config/_common.py +++ b/src/dolphin/workflows/config/_common.py @@ -287,6 +287,18 @@ class OutputOptions(BaseModel, extra="forbid"): list(DEFAULT_TIFF_OPTIONS), description="GDAL creation options for GeoTIFF files", ) + add_overviews: bool = Field( + True, + description=( + "Whether to add overviews to the output GeoTIFF files. This will " + "increase file size, but can be useful for visualizing the data with " + "web mapping tools. See https://gdal.org/programs/gdaladdo.html for more." + ), + ) + overview_levels: list[int] = Field( + [4, 8, 16, 32, 64], + description="List of overview levels to create (if `add_overviews=True`).", + ) # validators @field_validator("bounds", mode="after") diff --git a/src/dolphin/workflows/stitching_bursts.py b/src/dolphin/workflows/stitching_bursts.py index 71c3f597..87e2b1e4 100644 --- a/src/dolphin/workflows/stitching_bursts.py +++ b/src/dolphin/workflows/stitching_bursts.py @@ -6,6 +6,7 @@ from dolphin import stitching from dolphin._log import get_log, log_runtime +from dolphin._overviews import ImageType, create_image_overviews, create_overviews from dolphin._types import Bbox from dolphin.interferogram import estimate_interferometric_correlations @@ -100,6 +101,13 @@ def run( out_bounds_epsg=output_options.bounds_epsg, ) + if output_options.add_overviews: + logger.info("Creating overviews for stitched images") + create_overviews(stitched_ifg_paths, image_type=ImageType.INTERFEROGRAM) + create_overviews(interferometric_corr_paths, image_type=ImageType.CORRELATION) + create_image_overviews(stitched_ps_file, image_type=ImageType.PS) + create_image_overviews(stitched_temp_coh_file, image_type=ImageType.CORRELATION) + return ( stitched_ifg_paths, interferometric_corr_paths, diff --git a/src/dolphin/workflows/unwrapping.py b/src/dolphin/workflows/unwrapping.py index 9a2982e4..56815cd0 100644 --- a/src/dolphin/workflows/unwrapping.py +++ b/src/dolphin/workflows/unwrapping.py @@ -6,6 +6,7 @@ from dolphin import io, stitching, unwrap from dolphin._log import get_log, log_runtime +from dolphin._overviews import ImageType, create_overviews from dolphin._types import PathOrStr from .config import UnwrapOptions @@ -20,6 +21,7 @@ def run( nlooks: float, unwrap_options: UnwrapOptions, mask_file: PathOrStr | None = None, + add_overviews: bool = True, ) -> tuple[list[Path], list[Path]]: """Run the displacement workflow on a stack of SLCs. @@ -37,6 +39,9 @@ def run( mask_file : PathOrStr, optional Path to boolean mask indicating nodata areas. 1 indicates valid data, 0 indicates missing data. + add_overviews : bool, default = True + If True, creates overviews of the unwrapped phase and connected component + labels. Returns ------- @@ -80,6 +85,12 @@ def run( scratchdir=unwrap_scratchdir, ) + if add_overviews: + logger.info("Creating overviews for unwrapped images") + create_overviews(unwrapped_paths, image_type=ImageType.UNWRAPPED) + create_overviews(conncomp_paths, image_type=ImageType.CONNCOMP) + create_overviews(unwrapped_paths, image_type=ImageType.CORRELATION) + return (unwrapped_paths, conncomp_paths) diff --git a/tests/conftest.py b/tests/conftest.py index 99ad2f8d..a0b9ce50 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -427,3 +427,39 @@ def dem_file(tmp_path, slc_stack): ds.GetRasterBand(1).WriteArray(dem) ds = None return fname + + +# For unwrapping/overviews +@pytest.fixture() +def list_of_gtiff_ifgs(tmp_path, raster_100_by_200): + ifg_list = [] + for i in range(3): + # Create a copy of the raster in the same directory + f = tmp_path / f"ifg_{i}.int.tif" + write_arr( + arr=np.ones((100, 200), dtype=np.complex64), + output_name=f, + like_filename=raster_100_by_200, + driver="GTiff", + ) + ifg_list.append(f) + + return ifg_list + + +@pytest.fixture() +def list_of_envi_ifgs(tmp_path, raster_100_by_200): + ifg_list = [] + for i in range(3): + # Create a copy of the raster in the same directory + f = tmp_path / f"ifg_{i}.int" + ifg_list.append(f) + write_arr( + arr=np.ones((100, 200), dtype=np.complex64), + output_name=f, + like_filename=raster_100_by_200, + driver="ENVI", + ) + ifg_list.append(f) + + return ifg_list diff --git a/tests/test_overviews.py b/tests/test_overviews.py new file mode 100644 index 00000000..7332a611 --- /dev/null +++ b/tests/test_overviews.py @@ -0,0 +1,55 @@ +import pytest +import rasterio as rio + +from dolphin._overviews import ( + ImageType, + Resampling, + create_image_overviews, + create_overviews, +) + +# Dataset has no geotransform, gcps, or rpcs. The identity matrix will be returned. +pytestmark = pytest.mark.filterwarnings( + "ignore::rasterio.errors.NotGeoreferencedWarning", + "ignore:.*io.FileIO.*:pytest.PytestUnraisableExceptionWarning", +) + + +def get_overviews(filename, band=1): + with rio.open(filename) as src: + return src.overviews(band) + + +def test_create_image_overviews(list_of_gtiff_ifgs): + f = list_of_gtiff_ifgs[0] + assert len(get_overviews(f)) == 0 + create_image_overviews(f, image_type=ImageType.INTERFEROGRAM) + assert len(get_overviews(f)) > 0 + + +def test_create_image_overviews_envi(list_of_envi_ifgs): + f = list_of_envi_ifgs[0] + assert len(get_overviews(f)) == 0 + create_image_overviews(f, image_type=ImageType.INTERFEROGRAM) + assert len(get_overviews(f)) > 0 + + +@pytest.mark.parametrize("resampling", list(Resampling)) +def test_resamplings(list_of_gtiff_ifgs, resampling): + f = list_of_gtiff_ifgs[0] + assert len(get_overviews(f)) == 0 + create_image_overviews(f, resampling=resampling) + assert len(get_overviews(f)) > 0 + + +@pytest.mark.parametrize("levels", [[2, 4], [4, 8, 18]]) +def test_levels(list_of_gtiff_ifgs, levels): + f = list_of_gtiff_ifgs[0] + assert len(get_overviews(f)) == 0 + create_image_overviews(f, resampling="nearest", levels=levels) + assert len(get_overviews(f)) == len(levels) + + +def test_create_overviews(list_of_gtiff_ifgs): + create_overviews(list_of_gtiff_ifgs, image_type=ImageType.INTERFEROGRAM) + assert all(len(get_overviews(f)) > 0 for f in list_of_gtiff_ifgs) diff --git a/tests/test_unwrap.py b/tests/test_unwrap.py index be8f0f11..3b628145 100644 --- a/tests/test_unwrap.py +++ b/tests/test_unwrap.py @@ -27,112 +27,104 @@ def corr_raster(raster_100_by_200): # Make a correlation raster of all 1s in the same directory as the raster d = Path(raster_100_by_200).parent - corr_raster = d / "corr_raster.bin" + corr_raster = d / "corr_raster.cor.tif" io.write_arr( arr=np.ones((100, 200), dtype=np.float32), output_name=corr_raster, like_filename=raster_100_by_200, - driver="ENVI", + driver="GTiff", ) return corr_raster -def test_unwrap_snaphu(tmp_path, raster_100_by_200, corr_raster): - unw_filename = tmp_path / "unwrapped.unw.tif" - unw_path, conncomp_path = dolphin.unwrap.unwrap( - ifg_filename=raster_100_by_200, - corr_filename=corr_raster, - unw_filename=unw_filename, - nlooks=1, - init_method="mst", - ) - assert unw_path == unw_filename - assert str(conncomp_path) == str(unw_filename).replace( - ".unw.tif", ".unw.conncomp.tif" - ) - assert io.get_raster_xysize(unw_filename) == io.get_raster_xysize(raster_100_by_200) - - # test other init_method - unw_path, conncomp_path = dolphin.unwrap.unwrap( - ifg_filename=raster_100_by_200, - corr_filename=corr_raster, - unw_filename=unw_filename, - nlooks=1, - init_method="mcf", - ) - - -@pytest.mark.parametrize("method", [UnwrapMethod.ICU, UnwrapMethod.PHASS]) -def test_unwrap_methods(tmp_path, raster_100_by_200, corr_raster, method): - unw_filename = tmp_path / f"{method.value}_unwrapped.unw.tif" - dolphin.unwrap.unwrap( - ifg_filename=raster_100_by_200, - corr_filename=corr_raster, - unw_filename=unw_filename, - nlooks=1, - unwrap_method=method, - ) - - -def test_unwrap_logfile(tmp_path, raster_100_by_200, corr_raster): - unw_filename = tmp_path / "unwrapped.unw.tif" - dolphin.unwrap.unwrap( - ifg_filename=raster_100_by_200, - corr_filename=corr_raster, - unw_filename=unw_filename, - nlooks=1, - unwrap_method="icu", - log_to_file=True, - ) - logfile_name = str(unw_filename).replace(".unw.tif", ".unw.log") - assert Path(logfile_name).exists() - - -@pytest.fixture() -def list_of_ifgs(tmp_path, raster_100_by_200): - ifg_list = [] - for i in range(3): - # Create a copy of the raster in the same directory - f = tmp_path / f"ifg_{i}.int" - ifg_list.append(f) - io.write_arr( - arr=np.ones((100, 200), dtype=np.complex64), - output_name=f, - like_filename=raster_100_by_200, - driver="ENVI", +class TestUnwrapSingle: + def test_unwrap_snaphu_default(self, tmp_path, list_of_gtiff_ifgs, corr_raster): + unw_filename = tmp_path / "unwrapped.unw.tif" + unw_path, conncomp_path = dolphin.unwrap.unwrap( + ifg_filename=list_of_gtiff_ifgs[0], + corr_filename=corr_raster, + unw_filename=unw_filename, + nlooks=1, ) - ifg_list.append(f) - - return ifg_list - - -def test_run(list_of_ifgs, corr_raster): - ifg_path = list_of_ifgs[0].parent - dolphin.unwrap.run( - ifg_filenames=list_of_ifgs, - cor_filenames=[corr_raster] * len(list_of_ifgs), - output_path=ifg_path, - nlooks=1, - init_method="mst", - max_jobs=1, - ) - - -@pytest.fixture() -def list_of_gtiff_ifgs(tmp_path, raster_100_by_200): - ifg_list = [] - for i in range(3): - # Create a copy of the raster in the same directory - f = tmp_path / f"ifg_{i}.int.tif" - io.write_arr( - arr=np.ones((100, 200), dtype=np.complex64), - output_name=f, - like_filename=raster_100_by_200, - driver="GTiff", + assert unw_path == unw_filename + assert str(conncomp_path) == str(unw_filename).replace( + ".unw.tif", ".unw.conncomp.tif" + ) + assert io.get_raster_xysize(unw_filename) == io.get_raster_xysize( + list_of_gtiff_ifgs[0] ) - ifg_list.append(f) - return ifg_list + @pytest.mark.parametrize("init_method", ["mst", "mcf"]) + def test_unwrap_snaphu( + self, tmp_path, list_of_gtiff_ifgs, corr_raster, init_method + ): + # test other init_method + unw_filename = tmp_path / "unwrapped.unw.tif" + unw_path, conncomp_path = dolphin.unwrap.unwrap( + ifg_filename=list_of_gtiff_ifgs[0], + corr_filename=corr_raster, + unw_filename=unw_filename, + nlooks=1, + init_method=init_method, + ) + assert unw_path.exists() + assert conncomp_path.exists() + + @pytest.mark.parametrize("method", [UnwrapMethod.ICU, UnwrapMethod.PHASS]) + def test_unwrap_methods(self, tmp_path, raster_100_by_200, corr_raster, method): + unw_filename = tmp_path / f"{method.value}_unwrapped.unw.tif" + u_path, c_path = dolphin.unwrap.unwrap( + ifg_filename=raster_100_by_200, + corr_filename=corr_raster, + unw_filename=unw_filename, + nlooks=1, + unwrap_method=method, + ) + assert u_path.exists() + assert c_path.exists() + + def test_unwrap_logfile(self, tmp_path, raster_100_by_200, corr_raster): + unw_filename = tmp_path / "unwrapped.unw.tif" + u_path, c_path = dolphin.unwrap.unwrap( + ifg_filename=raster_100_by_200, + corr_filename=corr_raster, + unw_filename=unw_filename, + nlooks=1, + unwrap_method="icu", + log_to_file=True, + ) + logfile_name = str(unw_filename).replace(".unw.tif", ".unw.log") + assert Path(logfile_name).exists() + assert u_path.exists() + assert c_path.exists() + + +class TestUnwrapRun: + def test_run_gtiff(self, list_of_gtiff_ifgs, corr_raster): + ifg_path = list_of_gtiff_ifgs[0].parent + u_paths, c_paths = dolphin.unwrap.run( + ifg_filenames=list_of_gtiff_ifgs, + cor_filenames=[corr_raster] * len(list_of_gtiff_ifgs), + output_path=ifg_path, + nlooks=1, + init_method="mst", + max_jobs=1, + ) + assert all(p.exists() for p in u_paths) + assert all(p.exists() for p in c_paths) + + def test_run_envi(self, list_of_envi_ifgs, corr_raster): + ifg_path = list_of_envi_ifgs[0].parent + u_paths, c_paths = dolphin.unwrap.run( + ifg_filenames=list_of_envi_ifgs, + cor_filenames=[corr_raster] * len(list_of_envi_ifgs), + output_path=ifg_path, + nlooks=1, + init_method="mst", + max_jobs=1, + ) + assert all(p.exists() for p in u_paths) + assert all(p.exists() for p in c_paths) @pytest.mark.skipif(