dataset.preprocessing.cog_conversion

cog_conversion

HiRISE preprocessing utilities.

Provides:

• :func:jp2_to_cog — Convert a single JP2 to a Cloud-Optimized GeoTIFF (COG) sidecar. COGs use internal 512×512 tiling so rasterio can decompress only the tiles that overlap a small query window, giving 10–100× faster random-access reads compared to JPEG2000 during ML training.

• :func:convert_all — Batch-convert all JP2 files under a root directory using a process pool.

CLI usage::

# Convert all JP2s to COG GeoTIFFs (run once before training)
uv run python -m src.preprocessing --root /scratch/mars_hirise --workers 4

# Overwrite existing .tif sidecars
uv run python -m src.preprocessing --root /scratch/mars_hirise --workers 4 --overwrite

jp2_to_cog

jp2_to_cog(jp2_path: Path, overwrite: bool = False) -> pathlib.Path | None

Convert a single HiRISE JP2 to a Cloud-Optimized GeoTIFF sidecar.

The COG is written alongside the source JP2 with the same stem and a .tif extension. :meth:~temp.MarsHiRISE.prefer_cog will automatically use it when it exists, bypassing the slower JP2 path.

The conversion proceeds in two passes:

1. Write a temporary intermediate GeoTIFF so that overviews can be built on a writeable dataset (rasterio requires this).
2. Copy the intermediate file to the final COG path with copy_src_overviews=True to embed the overviews efficiently.

Parameters:

Name	Type	Description	Default
jp2_path	Path	Path to the source JPEG2000 file.	required
overwrite	bool	If False (default), skip files that already have a .tif sidecar.	False

Returns:

Type	Description
Path | None	Path to the output COG on success, or None if conversion failed.

Source code in src/dataset/preprocessing/cog_conversion.py

def jp2_to_cog(jp2_path: pathlib.Path, overwrite: bool = False) -> pathlib.Path | None:
    """Convert a single HiRISE JP2 to a Cloud-Optimized GeoTIFF sidecar.

    The COG is written alongside the source JP2 with the same stem and a
    ``.tif`` extension.  :meth:`~temp.MarsHiRISE.prefer_cog` will
    automatically use it when it exists, bypassing the slower JP2 path.

    The conversion proceeds in two passes:

    1. Write a temporary intermediate GeoTIFF so that overviews can be built
       on a writeable dataset (rasterio requires this).
    2. Copy the intermediate file to the final COG path with
       ``copy_src_overviews=True`` to embed the overviews efficiently.

    Args:
        jp2_path: Path to the source JPEG2000 file.
        overwrite: If ``False`` (default), skip files that already have a
            ``.tif`` sidecar.

    Returns:
        Path to the output COG on success, or ``None`` if conversion failed.
    """
    cog = _cog_path(jp2_path)

    if cog.exists() and not overwrite:
        logger.debug("COG sidecar already exists, skipping: %s", cog.name)
        return cog

    tmp = cog.with_suffix(".tmp.tif")
    try:
        # Some HiRISE JP2s carry no embedded geotransform; rasterio emits
        # NotGeoreferencedWarning on open (identity matrix assumed) and again
        # on the intermediate write.  Both are expected — we copy whatever
        # spatial metadata is present — so suppress them for the whole pass.
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                category=UserWarning,
                message=".*geotransform.*|.*identity matrix.*",
            )
            with rasterio.open(jp2_path) as src:
                profile = src.profile.copy()
                profile.update(
                    driver="GTiff",
                    tiled=True,
                    blockxsize=512,
                    blockysize=512,
                    bigtiff="IF_SAFER",  # switch to 64-bit offsets when >4 GiB
                )
                # The intermediate is a scratch file deleted in `finally` —
                # compressing it with deflate is the main conversion bottleneck
                # (CPU-intensive on gigabytes of data that are immediately
                # discarded).  Compression is applied only in the final
                # rasterio.shutil.copy call below.
                for key in ("lossless", "quality", "compress", "predictor", "zstd_level"):
                    profile.pop(key, None)

                logger.debug("Writing intermediate GeoTIFF to RAM for %s …", jp2_path.name)

                with rasterio.MemoryFile() as memfile:
                    with memfile.open(**profile) as dst:
                        for band_idx in src.indexes:
                            band_data = src.read(band_idx)
                            dst.write(band_data, band_idx)
                            del band_data

                        dst.build_overviews(_OVERVIEW_LEVELS, _OVERVIEW_RESAMPLING)
                        dst.update_tags(ns="rio_overview", resampling=_OVERVIEW_RESAMPLING.name)

                        # Second pass: copy to final COG with overviews embedded.
                        # Use _COG_CREATION_OPTIONS directly — it already carries compress,
                        # predictor, tiling, and copy_src_overviews.  Deriving opts from
                        # `profile` would omit compression because we stripped it above.

                        # Second pass: copy directly from RAM to the final NVMe file
                        rasterio.shutil.copy(dst, cog, **_COG_CREATION_OPTIONS)

        # Source JP2 is now closed; release any lingering references before the
        # second pass so the decompressed pixel data can be reclaimed.
        gc.collect()

        logger.info("COG written: %s", cog.name)
        return cog

    except rasterio.errors.RasterioIOError as exc:
        if _is_corrupt_jp2_error(exc):
            logger.warning(
                "Corrupted JP2 detected — deleting so it can be re-downloaded: %s",
                jp2_path.name,
            )
            jp2_path.unlink(missing_ok=True)
        else:
            logger.error("COG conversion failed for %s: %s", jp2_path.name, exc)
        cog.unlink(missing_ok=True)
        return None
    except Exception as exc:
        logger.error("COG conversion failed for %s: %s", jp2_path.name, exc)
        cog.unlink(missing_ok=True)
        return None

    finally:
        tmp.unlink(missing_ok=True)

img_to_cog

img_to_cog(img_path: Path, overwrite: bool = False) -> pathlib.Path | None

Convert a single HiRISE DTM .IMG (PDS3 float32) to a Cloud-Optimized GeoTIFF.

DTM .IMG files are flat binary rasters with attached PDS3 labels. Unlike JP2 orthoimages, they are not compressed, so the conversion is mainly about adding internal 512×512 tiling and overviews for fast random-access reads during ML training.

Float32 elevation data uses predictor=3 (floating-point differencing) for better deflate compression. The nodata sentinel -3.4028226550889045e+38 is preserved in the output GeoTIFF metadata.

Parameters:

Name	Type	Description	Default
img_path	Path	Path to the PDS3 .IMG file.	required
overwrite	bool	If False (default), skip files that already have a .tif sidecar.	False

Returns:

Type	Description
Path | None	Path to the output COG on success, or None if conversion failed.

Source code in src/dataset/preprocessing/cog_conversion.py

def img_to_cog(img_path: pathlib.Path, overwrite: bool = False) -> pathlib.Path | None:
    """Convert a single HiRISE DTM .IMG (PDS3 float32) to a Cloud-Optimized GeoTIFF.

    DTM ``.IMG`` files are flat binary rasters with attached PDS3 labels.
    Unlike JP2 orthoimages, they are not compressed, so the conversion is
    mainly about adding internal 512×512 tiling and overviews for fast
    random-access reads during ML training.

    Float32 elevation data uses ``predictor=3`` (floating-point differencing)
    for better deflate compression.  The nodata sentinel
    ``-3.4028226550889045e+38`` is preserved in the output GeoTIFF metadata.

    Args:
        img_path: Path to the PDS3 ``.IMG`` file.
        overwrite: If ``False`` (default), skip files that already have a
            ``.tif`` sidecar.

    Returns:
        Path to the output COG on success, or ``None`` if conversion failed.
    """
    cog = _img_cog_path(img_path)

    if cog.exists() and not overwrite:
        logger.debug("COG sidecar already exists, skipping: %s", cog.name)
        return cog

    tmp = cog.with_suffix(".tmp.tif")
    try:
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                category=UserWarning,
                message=".*geotransform.*|.*identity matrix.*",
            )
            with rasterio.open(img_path) as src:
                profile = src.profile.copy()
                profile.update(
                    driver="GTiff",
                    tiled=True,
                    blockxsize=512,
                    blockysize=512,
                    bigtiff="IF_SAFER",
                )
                # Strip JP2-specific keys that don't apply
                for key in ("lossless", "quality", "compress", "predictor", "zstd_level"):
                    profile.pop(key, None)

                # Ensure float32 dtype for elevation data
                if profile.get("dtype") is None:  # pragma: no cover - rasterio always reports a dtype
                    profile["dtype"] = "float32"

                # Preserve nodata
                src_nodata = src.nodata
                if src_nodata is None:
                    # HiRISE DTMs use IEEE float32 min as nodata
                    src_nodata = _DTM_NODATA
                profile["nodata"] = src_nodata

                logger.debug(
                    "Writing intermediate GeoTIFF for DTM %s "
                    "(dtype=%s, %dx%d, %d band(s)) …",
                    img_path.name,
                    profile.get("dtype"),
                    src.width,
                    src.height,
                    src.count,
                )
                logger.debug("Writing intermediate DTM GeoTIFF to RAM for %s …", img_path.name)

                # THE MAGIC: Use RAM instead of NVMe
                with rasterio.MemoryFile() as memfile:
                    with memfile.open(**profile) as dst:
                        for band_idx in src.indexes:
                            band_data = src.read(band_idx)
                            dst.write(band_data, band_idx)
                            del band_data

                        dst.build_overviews(_OVERVIEW_LEVELS, _OVERVIEW_RESAMPLING)
                        dst.update_tags(ns="rio_overview", resampling=_OVERVIEW_RESAMPLING.name)

                        # Copy directly from RAM to NVMe
                        rasterio.shutil.copy(dst, cog, **_COG_CREATION_OPTIONS_FLOAT)

            gc.collect()

            logger.info("DTM COG written: %s", cog.name)
        return cog

    except rasterio.errors.RasterioIOError as exc:
        logger.error("DTM COG conversion failed for %s: %s", img_path.name, exc)
        cog.unlink(missing_ok=True)
        return None
    except Exception as exc:
        logger.error("DTM COG conversion failed for %s: %s", img_path.name, exc)
        cog.unlink(missing_ok=True)
        return None
    finally:
        tmp.unlink(missing_ok=True)

convert_all

convert_all(root: Path, workers: int = 4, overwrite: bool = False, skip_jp2: bool = False, skip_dtm: bool = False) -> dict[str, int]

Convert all JP2 and DTM .IMG files under root to COG GeoTIFF sidecars.

Uses a :class:~concurrent.futures.ProcessPoolExecutor to parallelise the CPU-bound conversion. Each worker calls :func:jp2_to_cog or :func:img_to_cog.

.. note:: Large JP2 files (up to 2.5 GB) are fully decompressed in memory during the intermediate write step. The actual worker count is automatically capped by :func:_safe_worker_count based on available RAM and the estimated decompressed size of the JP2 files found under root; the workers argument is therefore treated as an upper bound.

Parameters:

Name	Type	Description	Default
root	Path	Dataset root directory containing JP2 and/or IMG files.	required
workers	int	Number of parallel conversion processes.	4
overwrite	bool	Re-convert files that already have a .tif sidecar.	False
skip_jp2	bool	Skip JP2 orthoimage conversion (only process DTM .IMG).	False
skip_dtm	bool	Skip DTM .IMG conversion (only process JP2 orthoimages).	False

Returns:

Type	Description
dict[str, int]	Dict with keys "converted", "skipped", and "failed".

Source code in src/dataset/preprocessing/cog_conversion.py

def convert_all(
        root: pathlib.Path,
        workers: int = 4,
        overwrite: bool = False,
        skip_jp2: bool = False,
        skip_dtm: bool = False,
) -> dict[str, int]:
    """Convert all JP2 and DTM .IMG files under *root* to COG GeoTIFF sidecars.

    Uses a :class:`~concurrent.futures.ProcessPoolExecutor` to parallelise
    the CPU-bound conversion.  Each worker calls :func:`jp2_to_cog` or
    :func:`img_to_cog`.

    .. note::
        Large JP2 files (up to 2.5 GB) are fully decompressed in memory during
        the intermediate write step.  The actual worker count is automatically
        capped by :func:`_safe_worker_count` based on available RAM and the
        estimated decompressed size of the JP2 files found under *root*; the
        ``workers`` argument is therefore treated as an upper bound.

    Args:
        root: Dataset root directory containing JP2 and/or IMG files.
        workers: Number of parallel conversion processes.
        overwrite: Re-convert files that already have a ``.tif`` sidecar.
        skip_jp2: Skip JP2 orthoimage conversion (only process DTM .IMG).
        skip_dtm: Skip DTM .IMG conversion (only process JP2 orthoimages).

    Returns:
        Dict with keys ``"converted"``, ``"skipped"``, and ``"failed"``.
    """
    # Collect files to convert
    tasks: list[tuple[pathlib.Path, Callable]] = []

    if not skip_jp2:
        jp2_files = list(_iter_jp2_files(root))
        logger.info("Found %d JP2 file(s) under %s.", len(jp2_files), root)
        tasks.extend((p, jp2_to_cog) for p in jp2_files)
    else:
        jp2_files = []

    if not skip_dtm:
        img_files = list(_iter_img_files(root))
        logger.info("Found %d DTM .IMG file(s) under %s.", len(img_files), root)
        tasks.extend((p, img_to_cog) for p in img_files)
    else:
        img_files = []

    if not tasks:
        logger.warning("No files found to convert under %s.", root)
        return {"converted": 0, "skipped": 0, "failed": 0}

    counts: dict[str, int] = {"converted": 0, "skipped": 0, "failed": 0}

    # Memory-safe worker count is computed from JP2s (the larger files);
    # IMG files are typically smaller in memory since they're already raw.
    size_reference = jp2_files if jp2_files else img_files
    actual_workers = _safe_worker_count(size_reference, workers)

    pool = concurrent.futures.ProcessPoolExecutor(
        max_workers=actual_workers,
        mp_context=multiprocessing.get_context("spawn"),
        initializer=_worker_init,
        max_tasks_per_child=1,
    )
    future_to_path = {
        pool.submit(convert_fn, p, overwrite): p
        for p, convert_fn in tasks
    }
    try:
        for future in concurrent.futures.as_completed(future_to_path):
            src_path = future_to_path[future]
            try:
                result = future.result()
                if result is None:
                    counts["failed"] += 1
                else:
                    # Check if the COG sidecar was freshly written
                    cog = src_path.with_suffix(".tif")
                    if cog.exists() and cog.stat().st_mtime >= src_path.stat().st_mtime:
                        counts["converted"] += 1
                    else:
                        counts["skipped"] += 1
            except Exception as exc:
                logger.error("Worker error for %s: %s", src_path.name, exc)
                counts["failed"] += 1
    except KeyboardInterrupt:
        pool.shutdown(wait=False, cancel_futures=True)
        logger.warning("Interrupted — %d converted so far.", counts["converted"])
        raise SystemExit(130)
    else:
        pool.shutdown(wait=True)

    logger.info(
        "COG conversion complete — converted: %d, skipped: %d, failed: %d.",
        counts["converted"],
        counts["skipped"],
        counts["failed"],
    )
    return counts