visium_hd.py

from __future__ import annotations

import json
import re
import warnings
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Any

import h5py
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import scanpy as sc
from dask_image.imread import imread
from geopandas import GeoDataFrame
from imageio import imread as imread2
from numpy.random import default_rng
from scipy.sparse import csc_matrix
from shapely.geometry import Polygon
from skimage.transform import ProjectiveTransform, warp
from spatialdata import (
    SpatialData,
    get_extent,
    rasterize_bins,
    rasterize_bins_link_table_to_labels,
)
from spatialdata.models import Image2DModel, ShapesModel, TableModel
from spatialdata.transformations import Affine, Identity, Scale, set_transformation
from xarray import DataArray

from spatialdata_io._constants._constants import VisiumHDKeys
from spatialdata_io._docs import inject_docs

if TYPE_CHECKING:
    from collections.abc import Mapping

    from anndata import AnnData
    from multiscale_spatial_image import MultiscaleSpatialImage
    from spatial_image import SpatialImage
    from spatialdata._types import ArrayLike

RNG = default_rng(0)


@inject_docs(vx=VisiumHDKeys)
def visium_hd(
    path: str | Path,
    dataset_id: str | None = None,
    filtered_counts_file: bool = True,
    load_segmentations_only: bool | None = None,
    load_nucleus_segmentations: bool = False,
    bin_size: int | list[int] | None = None,
    bins_as_squares: bool = True,
    annotate_table_by_labels: bool = False,
    fullres_image_file: str | Path | None = None,
    load_all_images: bool = False,
    var_names_make_unique: bool = True,
    imread_kwargs: Mapping[str, Any] = MappingProxyType({}),
    image_models_kwargs: Mapping[str, Any] = MappingProxyType({}),
    anndata_kwargs: Mapping[str, Any] = MappingProxyType({}),
) -> SpatialData:
    """Read *10x Genomics* Visium HD formatted dataset.

    Parameters
    ----------
    path
        Path to directory containing the *10x Genomics* Visium HD output.
    dataset_id
        Unique identifier of the dataset, used to name the elements of the `SpatialData` object. If `None`, it tries to
         infer it from the file name of the feature slice file.
    filtered_counts_file
        It sets the value of `counts_file` to ``{vx.FILTERED_COUNTS_FILE!r}`` (when `True`) or to
        ``{vx.RAW_COUNTS_FILE!r}`` (when `False`).
    load_segmentations_only
        If `True`, only the segmented cell boundaries and their associated counts will be loaded. All binned data
        will be skipped. If `False`, only the binned data will be loaded (which is consistent with legacy behavior).
        If `None` (default), it will be equivalent to `False`, but a deprecation warning will be raised to inform users that
        in future releases the default value will be changed to `True`. To avoid the warning, explicitly set this parameter to
        `False` or `True`.
    load_nucleus_segmentations
        If `True` and nucleus segmentation files are present, load nucleus segmentation polygons and the corresponding
        nucleus-filtered count table. The counts are aggregated from the 2 µm binned matrix using the provided
        barcode mappings so that only bins under segmented nuclei contribute to each cell’s counts. Requires all of:
        nucleus segmentation GeoJSON, barcode_mappings.parquet, and the 2 µm filtered_feature_bc_matrix.h5.
    bin_size
        When specified, load the data of a specific bin size, or a list of bin sizes. By default, it loads all the
        available bin sizes.
    bins_as_squares
        If `True`, the bins are represented as squares. If `False`, the bins are represented as circles. For a correct
        visualization one should use squares.
    annotate_table_by_labels
        If `True`, the tables will annotate labels layers representing the bins, if `False`, the tables will annotate
        shapes layer.
    fullres_image_file
        Path to the full-resolution image. By default the image is searched in the ``{vx.MICROSCOPE_IMAGE!r}``
        directory.
    load_all_images
        If `False`, load only the full resolution, high resolution and low resolution images. If `True`, also the
        following images: ``{vx.IMAGE_CYTASSIST!r}``.
    var_names_make_unique
        If `True`, call `.var_names_make_unique()` on each `AnnData` table.
    imread_kwargs
        Keyword arguments for :func:`imageio.imread`.
    image_models_kwargs
        Keyword arguments for :class:`spatialdata.models.Image2DModel`.
    anndata_kwargs
        Keyword arguments for :func:`anndata.io.read_h5ad`.

    Returns
    -------
    SpatialData object for the Visium HD data.
    """
    path = Path(path)
    all_files = [file for file in path.rglob("*") if file.is_file()]
    tables = {}
    shapes = {}
    images: dict[str, Any] = {}
    labels: dict[str, Any] = {}

    # Deprecation warning for load_segmentations_only default value
    if not load_segmentations_only:
        warnings.warn(
            "`load_segmentations_only` default value will change to `True` in future releases. Please set it "
            "explicitly to `True` or `False` to avoid this warning.",
            FutureWarning,
            stacklevel=2,
        )

    # Check for segmentation files
    SEGMENTED_OUTPUTS_PATH = path / VisiumHDKeys.SEGMENTATION_OUTPUTS
    COUNT_MATRIX_PATH = SEGMENTED_OUTPUTS_PATH / VisiumHDKeys.FILTERED_CELL_COUNTS_FILE
    CELL_GEOJSON_PATH = SEGMENTED_OUTPUTS_PATH / VisiumHDKeys.CELL_SEGMENTATION_GEOJSON_PATH
    NUCLEUS_GEOJSON_PATH = SEGMENTED_OUTPUTS_PATH / VisiumHDKeys.NUCLEUS_SEGMENTATION_GEOJSON_PATH
    SCALE_FACTORS_PATH = SEGMENTED_OUTPUTS_PATH / VisiumHDKeys.SPATIAL / VisiumHDKeys.SCALEFACTORS_FILE
    if not SCALE_FACTORS_PATH.exists():
        # Visium HD 3.0.0 does not have the SEGMENTATION_OUTPUTS folder
        scale_factors_file = next(
            (file for file in path.rglob("*") if file.name.endswith(VisiumHDKeys.SCALEFACTORS_FILE)),
            None,
        )
        assert scale_factors_file is not None, "Scale factors file not found in any of the subdirectories."
        SCALE_FACTORS_PATH = scale_factors_file
    BARCODE_MAPPINGS_PATH = next(
        (file for file in path.rglob("*") if file.name.endswith(VisiumHDKeys.BARCODE_MAPPINGS_FILE)),
        None,
    )
    FILTERED_MATRIX_2U_PATH = (
        path / VisiumHDKeys.BINNED_OUTPUTS / f"{VisiumHDKeys.BIN_PREFIX}002um" / VisiumHDKeys.FILTERED_COUNTS_FILE
    )
    cell_segmentation_files_exist = (
        COUNT_MATRIX_PATH.exists() and CELL_GEOJSON_PATH.exists() and SCALE_FACTORS_PATH.exists()
    )
    nucleus_segmentation_files_exist = (
        NUCLEUS_GEOJSON_PATH.exists()
        and (BARCODE_MAPPINGS_PATH is not None and BARCODE_MAPPINGS_PATH.exists())
        and FILTERED_MATRIX_2U_PATH.exists()
    )

    if dataset_id is None:
        dataset_id = _infer_dataset_id(path)

    filename_prefix = _get_filename_prefix(path, dataset_id)

    def load_image(path: Path, suffix: str, scale_factors: list[int] | None = None) -> None:
        _load_image(
            path=path,
            images=images,
            suffix=suffix,
            dataset_id=dataset_id,
            imread_kwargs=imread_kwargs,
            image_models_kwargs=image_models_kwargs,
            scale_factors=scale_factors,
        )

    metadata, hd_layout = _parse_metadata(path, filename_prefix)
    file_format = hd_layout[VisiumHDKeys.FILE_FORMAT]
    if file_format != "1.0":
        warnings.warn(
            f"File format {file_format} is not supported. A more recent file format may be supported in a newer version"
            f"of the spatialdata-io package.",
            UserWarning,
            stacklevel=2,
        )

    with open(SCALE_FACTORS_PATH) as file:
        scalefactors = json.load(file)

    transform_lowres = Scale(
        np.array(
            [
                scalefactors[VisiumHDKeys.SCALEFACTORS_LOWRES],
                scalefactors[VisiumHDKeys.SCALEFACTORS_LOWRES],
            ]
        ),
        axes=("x", "y"),
    )
    transform_hires = Scale(
        np.array(
            [
                scalefactors[VisiumHDKeys.SCALEFACTORS_HIRES],
                scalefactors[VisiumHDKeys.SCALEFACTORS_HIRES],
            ]
        ),
        axes=("x", "y"),
    )
    # scaling
    transform_original = Identity()
    transformations = {
        dataset_id: transform_original,
        f"{dataset_id}_downscaled_hires": transform_hires,
        f"{dataset_id}_downscaled_lowres": transform_lowres,
    }

    # Load Binned Data
    if not load_segmentations_only:

        def _get_bins(path_bins: Path) -> list[str]:
            return sorted(
                [
                    bin_size.name
                    for bin_size in path_bins.iterdir()
                    if bin_size.is_dir() and bin_size.name.startswith(VisiumHDKeys.BIN_PREFIX)
                ]
            )

        all_path_bins = [path_bin for path_bin in all_files if VisiumHDKeys.BINNED_OUTPUTS in str(path_bin)]
        if len(all_path_bins) != 0:
            path_bins_parts = all_path_bins[
                -1
            ].parts  # just choosing last one here as users might have tar file which would be first
            path_bins = Path(*path_bins_parts[: path_bins_parts.index(VisiumHDKeys.BINNED_OUTPUTS) + 1])
        else:
            path_bins = path
        all_bin_sizes = _get_bins(path_bins)

        bin_sizes = []
        if bin_size is not None and (isinstance(bin_size, int) or len(bin_size) > 0):
            if not isinstance(bin_size, list):
                bin_size = [bin_size]
            bin_sizes = [f"square_{bs:03}um" for bs in bin_size if f"square_{bs:03}um" in all_bin_sizes]
            if len(bin_sizes) < len(bin_size):
                warnings.warn(
                    f"Requested bin size {bin_size} (available {all_bin_sizes}); ignoring the bin sizes that are not "
                    "found.",
                    UserWarning,
                    stacklevel=2,
                )
        if bin_size is None or bin_sizes == []:
            bin_sizes = all_bin_sizes

        # iterate over the given bins and load the data
        for bin_size_str in bin_sizes:
            path_bin = path_bins / bin_size_str
            counts_file = VisiumHDKeys.FILTERED_COUNTS_FILE if filtered_counts_file else VisiumHDKeys.RAW_COUNTS_FILE
            adata = sc.read_10x_h5(
                path_bin / counts_file,
                gex_only=False,
                **anndata_kwargs,
            )

            path_bin_spatial = path_bin / VisiumHDKeys.SPATIAL

            # The scale factors of binned data are consistent to the global ones
            # (which are already loaded in "scalefactors"), but the json file in the
            # binned spatial folder contains also the bin size information
            with open(path_bin_spatial / VisiumHDKeys.SCALEFACTORS_FILE) as file:
                scalefactors_bins = json.load(file)

            # consistency check
            found_bin_size = re.search(r"\d{3}", bin_size_str)
            assert found_bin_size is not None
            assert float(found_bin_size.group()) == scalefactors_bins[VisiumHDKeys.SCALEFACTORS_BIN_SIZE_UM]
            assert np.isclose(
                scalefactors_bins[VisiumHDKeys.SCALEFACTORS_BIN_SIZE_UM]
                / scalefactors_bins[VisiumHDKeys.SCALEFACTORS_SPOT_DIAMETER_FULLRES],
                scalefactors_bins[VisiumHDKeys.SCALEFACTORS_MICRONS_PER_PIXEL],
            )

            tissue_positions_file = path_bin_spatial / VisiumHDKeys.TISSUE_POSITIONS_FILE

            # read coordinates and set up adata.obs and adata.obsm
            coords = pd.read_parquet(tissue_positions_file)
            assert all(
                coords.columns.values
                == [
                    VisiumHDKeys.BARCODE,
                    VisiumHDKeys.IN_TISSUE,
                    VisiumHDKeys.ARRAY_ROW,
                    VisiumHDKeys.ARRAY_COL,
                    VisiumHDKeys.LOCATIONS_Y,
                    VisiumHDKeys.LOCATIONS_X,
                ]
            )
            coords.set_index(VisiumHDKeys.BARCODE, inplace=True, drop=True)
            coords_filtered = coords.loc[adata.obs.index]
            adata.obs = pd.merge(
                adata.obs,
                coords_filtered,
                how="left",
                left_index=True,
                right_index=True,
            )
            # compatibility to legacy squidpy
            adata.obsm["spatial"] = adata.obs[[VisiumHDKeys.LOCATIONS_X, VisiumHDKeys.LOCATIONS_Y]].values
            # dropping the spatial coordinates (will be stored in shapes)
            adata.obs.drop(
                columns=[
                    VisiumHDKeys.LOCATIONS_X,
                    VisiumHDKeys.LOCATIONS_Y,
                ],
                inplace=True,
            )
            adata.obs[VisiumHDKeys.INSTANCE_KEY] = np.arange(len(adata))

            # parse shapes
            shapes_name = dataset_id + "_" + bin_size_str
            radius = scalefactors_bins[VisiumHDKeys.SCALEFACTORS_SPOT_DIAMETER_FULLRES] / 2.0
            circles = ShapesModel.parse(
                adata.obsm["spatial"],
                geometry=0,
                radius=radius,
                index=adata.obs[VisiumHDKeys.INSTANCE_KEY].copy(),
                transformations=transformations,
            )
            if not bins_as_squares:
                shapes[shapes_name] = circles
            else:
                squares_series = circles.buffer(radius, cap_style=3)
                shapes[shapes_name] = ShapesModel.parse(
                    GeoDataFrame(geometry=squares_series),
                    transformations=transformations,
                )

            # parse table
            adata.obs[VisiumHDKeys.REGION_KEY] = shapes_name
            adata.obs[VisiumHDKeys.REGION_KEY] = adata.obs[VisiumHDKeys.REGION_KEY].astype("category")

            tables[bin_size_str] = TableModel.parse(
                adata,
                region=shapes_name,
                region_key=str(VisiumHDKeys.REGION_KEY),
                instance_key=str(VisiumHDKeys.INSTANCE_KEY),
            )
            if var_names_make_unique:
                tables[bin_size_str].var_names_make_unique()

    # Integrate the segmentation data (skipped if segmentation files are not found)
    if cell_segmentation_files_exist:
        print("Found segmentation data. Incorporating cell_segmentations.")
        cell_adata_hd = sc.read_10x_h5(COUNT_MATRIX_PATH)
        cell_adata_hd.var_names_make_unique()

        cell_shapes_gdf = _extract_geometries_from_geojson(
            cell_adata_hd,
            geojson_path=CELL_GEOJSON_PATH,
        )

        SHAPES_KEY_HD = f"{dataset_id}_{VisiumHDKeys.CELL_SEG_KEY_HD}"
        cell_adata_hd.obs["cell_id"] = cell_adata_hd.obs.index
        cell_adata_hd.obs["region"] = SHAPES_KEY_HD
        cell_adata_hd.obs["region"] = cell_adata_hd.obs["region"].astype("category")
        cell_adata_hd = cell_adata_hd[cell_shapes_gdf.index].copy()

        shapes[SHAPES_KEY_HD] = ShapesModel.parse(cell_shapes_gdf, transformations=transformations)
        tables[VisiumHDKeys.CELL_SEG_KEY_HD] = TableModel.parse(
            cell_adata_hd,
            region=SHAPES_KEY_HD,
            region_key="region",
            instance_key="cell_id",
        )

        # load nucleus segmentations if available
        if nucleus_segmentation_files_exist and load_nucleus_segmentations:
            print("Found nucleus segmentation data. Incorporating nucleus_segmentations.")

            # we already ensure this by having nucleus_segmentation_files_exist True, but
            # mypy is not able to infer that
            assert BARCODE_MAPPINGS_PATH is not None

            nucleus_adata_hd = _make_filtered_nucleus_adata(
                filtered_matrix_h5_path=FILTERED_MATRIX_2U_PATH,
                barcode_mappings_parquet_path=BARCODE_MAPPINGS_PATH,
            )
            nucleus_shapes_gdf = _extract_geometries_from_geojson(
                adata=nucleus_adata_hd, geojson_path=NUCLEUS_GEOJSON_PATH
            )

            SHAPES_KEY_HD = f"{dataset_id}_{VisiumHDKeys.NUCLEUS_SEG_KEY_HD}"
            nucleus_adata_hd.obs["cell_id"] = nucleus_adata_hd.obs.index
            nucleus_adata_hd.obs["region"] = SHAPES_KEY_HD
            nucleus_adata_hd.obs["region"] = nucleus_adata_hd.obs["region"].astype("category")
            nucleus_adata_hd = nucleus_adata_hd[nucleus_shapes_gdf.index].copy()

            shapes[SHAPES_KEY_HD] = ShapesModel.parse(nucleus_shapes_gdf, transformations=transformations)
            tables[VisiumHDKeys.NUCLEUS_SEG_KEY_HD] = TableModel.parse(
                nucleus_adata_hd,
                region=SHAPES_KEY_HD,
                region_key="region",
                instance_key="cell_id",
            )

    # Read all images and add transformations
    fullres_image_file_paths = []
    if fullres_image_file is not None:
        fullres_image_file_paths.append(Path(fullres_image_file))
    else:
        path_fullres = path / VisiumHDKeys.MICROSCOPE_IMAGE
        if path_fullres.exists():
            fullres_image_paths = [file for file in path_fullres.iterdir() if file.is_file()]
        elif list((path_fullres := (path / f"{filename_prefix}tissue_image")).parent.glob(f"{path_fullres.name}.*")):
            fullres_image_paths = list(path_fullres.parent.glob(f"{path_fullres.name}.*"))
        else:
            fullres_image_paths = []
        if len(fullres_image_paths) > 1:
            warnings.warn(
                f"Multiple files found in {path_fullres}, using the first one: {fullres_image_paths[0].stem}. Please"
                " specify the path to the full resolution image manually using the `fullres_image_file` argument.",
                UserWarning,
                stacklevel=2,
            )
        if len(fullres_image_paths) == 0:
            warnings.warn(
                "No full resolution image found. If incorrect, please specify the path in the "
                "`fullres_image_file` parameter when calling the `visium_hd` reader function.",
                UserWarning,
                stacklevel=2,
            )
        fullres_image_file = fullres_image_paths[0] if len(fullres_image_paths) > 0 else None

    if fullres_image_file is not None:
        load_image(
            path=fullres_image_file_paths[0],
            suffix="_full_image",
            scale_factors=[2, 2, 2, 2],
        )
    else:
        warnings.warn(
            "No full resolution image found. If incorrect, please specify the path in the "
            "`fullres_image_file` parameter when calling the `visium_hd` reader function.",
            UserWarning,
            stacklevel=2,
        )

    # hires image
    hires_image_path = [path for path in all_files if VisiumHDKeys.IMAGE_HIRES_FILE in str(path)]
    if len(hires_image_path) > 0:
        load_image(
            path=hires_image_path[0],
            suffix="_hires_image",
        )
        set_transformation(
            images[dataset_id + "_hires_image"],
            {
                f"{dataset_id}_downscaled_hires": Identity(),
                dataset_id: transform_hires.inverse(),
            },
            set_all=True,
        )
    else:
        warnings.warn(
            f"No image path found containing the hires image: {VisiumHDKeys.IMAGE_HIRES_FILE}",
            UserWarning,
            stacklevel=2,
        )

    # lowres image
    lowres_image_path = [path for path in all_files if VisiumHDKeys.IMAGE_LOWRES_FILE in str(path)]
    if len(lowres_image_path) > 0:
        load_image(
            path=lowres_image_path[0],
            suffix="_lowres_image",
        )
        set_transformation(
            images[dataset_id + "_lowres_image"],
            {
                f"{dataset_id}_downscaled_lowres": Identity(),
                dataset_id: transform_lowres.inverse(),
            },
            set_all=True,
        )
    else:
        warnings.warn(
            f"No image path found containing the lowres image: {VisiumHDKeys.IMAGE_LOWRES_FILE}",
            UserWarning,
            stacklevel=2,
        )

    # cytassist image
    cytassist_path = [path for path in all_files if VisiumHDKeys.IMAGE_CYTASSIST in str(path)]
    if load_all_images and len(cytassist_path) > 0:
        load_image(
            path=cytassist_path[0],
            suffix="_cytassist_image",
        )
        image = images[dataset_id + "_cytassist_image"]
        transform_matrices = _get_transform_matrices(metadata, hd_layout)
        projective0 = transform_matrices["cytassist_colrow_to_spot_colrow"]
        projective1 = transform_matrices["spot_colrow_to_microscope_colrow"]
        projective = projective1 @ projective0
        projective /= projective[2, 2]
        if _projective_matrix_is_affine(projective):
            affine = Affine(projective, input_axes=("x", "y"), output_axes=("x", "y"))
            set_transformation(image, affine, dataset_id)
        else:
            # the projective matrix is not affine, we will separate the affine part and the projective shift, and apply
            # the projective shift to the image
            affine_matrix, projective_shift = _decompose_projective_matrix(projective)
            affine = Affine(affine_matrix, input_axes=("x", "y"), output_axes=("x", "y"))

            # determine the size of the transformed image
            bounding_box = get_extent(image, coordinate_system=dataset_id)
            x0, x1 = bounding_box["x"]
            y0, y1 = bounding_box["y"]
            x1 -= 1
            y1 -= 1
            corners = [(x0, y0), (x1, y0), (x1, y1), (x0, y1)]

            transformed_corners = []
            for x, y in corners:
                px, py = _projective_matrix_transform_point(projective_shift, x, y)
                transformed_corners.append((px, py))
            transformed_corners_array = np.array(transformed_corners)
            transformed_bounds = (
                np.min(transformed_corners_array[:, 0]),
                np.min(transformed_corners_array[:, 1]),
                np.max(transformed_corners_array[:, 0]),
                np.max(transformed_corners_array[:, 1]),
            )
            # the first two components are <= 0, we just discard them since the cytassist image has a lot of padding
            # and therefore we can safely discard pixels with negative coordinates
            transformed_shape = (
                np.ceil(transformed_bounds[2]),
                np.ceil(transformed_bounds[3]),
            )

            # flip xy
            transformed_shape = (transformed_shape[1], transformed_shape[0])

            # the cytassist image is a small, single-scale image, so we can compute it in memory
            numpy_data = image.transpose("y", "x", "c").data.compute()
            warped = warp(
                numpy_data,
                ProjectiveTransform(projective_shift).inverse,
                output_shape=transformed_shape,
                order=1,
            )
            warped = np.round(warped * 255).astype(np.uint8)
            warped = Image2DModel.parse(
                warped,
                dims=("y", "x", "c"),
                transformations={dataset_id: affine},
                rgb=True,
            )
            # we replace the cytassist image with the warped image
            images[dataset_id + "_cytassist_image"] = warped
    elif load_all_images:
        warnings.warn(
            f"No image path found containing the cytassist image: {VisiumHDKeys.IMAGE_CYTASSIST}",
            UserWarning,
            stacklevel=2,
        )

    sdata = SpatialData(tables=tables, images=images, shapes=shapes, labels=labels)

    if annotate_table_by_labels:
        for bin_size_str in bin_sizes:
            shapes_name = dataset_id + "_" + bin_size_str
            # add labels layer (rasterized bins).
            labels_name = f"{dataset_id}_{bin_size_str}_labels"
            labels_element = rasterize_bins(
                sdata,
                bins=shapes_name,
                table_name=bin_size_str,
                row_key=VisiumHDKeys.ARRAY_ROW,
                col_key=VisiumHDKeys.ARRAY_COL,
                value_key=None,
                return_region_as_labels=True,
            )
            sdata[labels_name] = labels_element
            rasterize_bins_link_table_to_labels(
                sdata=sdata, table_name=bin_size_str, rasterized_labels_name=labels_name
            )

    return sdata


def _infer_dataset_id(path: Path) -> str:
    suffix = f"_{VisiumHDKeys.FEATURE_SLICE_FILE.value}"
    files = [file.name for file in path.iterdir() if file.is_file() and file.name.endswith(suffix)]
    if len(files) == 0 or len(files) > 1:
        raise ValueError(
            f"Cannot infer `dataset_id` from the feature slice file in {path}, please pass `dataset_id` as an "
            f"argument. The `dataset_id` value will be used to name the elements in the `SpatialData` object."
        )
    return files[0].replace(suffix, "")


def _load_image(
    path: Path,
    images: dict[str, SpatialImage | MultiscaleSpatialImage],
    suffix: str,
    dataset_id: str,
    imread_kwargs: Mapping[str, Any],
    image_models_kwargs: Mapping[str, Any],
    scale_factors: list[int] | None,
) -> None:
    if path.exists():
        if path.suffix != ".btf":
            data = imread(path)
            if len(data.shape) == 4:
                # this happens for the cytassist, hires and lowres images; the umi image doesn't need processing
                data = data.squeeze()
        else:
            if "MAX_IMAGE_PIXELS" in imread_kwargs:
                from PIL import Image as ImagePIL

                ImagePIL.MAX_IMAGE_PIXELS = dict(imread_kwargs).pop("MAX_IMAGE_PIXELS")
            # dask_image doesn't recognize .btf automatically and imageio v3 throws error due to pixel limit -> use imageio v2
            data = imread2(path, **imread_kwargs).squeeze()

        if data.shape[-1] == 3:  # HE image in RGB format
            data = data.transpose(2, 0, 1)
        else:
            assert data.shape[0] == min(data.shape), (
                "When the image is not in RGB, the first dimension should be the number of channels."
            )

        image = DataArray(data, dims=("c", "y", "x"))
        parsed = Image2DModel.parse(
            image,
            scale_factors=scale_factors,
            rgb=None,
            transformations={dataset_id: Identity()},
            **image_models_kwargs,
        )
        images[dataset_id + suffix] = parsed
    else:
        warnings.warn(f"File {path} does not exist, skipping it.", UserWarning, stacklevel=2)
    return None


def _projective_matrix_transform_point(projective_shift: ArrayLike, x: float, y: float) -> tuple[float, float]:
    v = np.array([x, y, 1])
    v = projective_shift @ v
    v /= v[2]
    return v[0], v[1]


def _projective_matrix_is_affine(projective_matrix: ArrayLike) -> bool:
    assert np.allclose(projective_matrix[2, 2], 1), "A projective matrix should have a 1 in the bottom right corner."
    return np.allclose(projective_matrix[2, :2], [0, 0])


def _decompose_projective_matrix(
    projective_matrix: ArrayLike,
) -> tuple[ArrayLike, ArrayLike]:
    """Decompose a projective transformation matrix into an affine transformation and a projective shift.

    Parameters
    ----------
    projective_matrix
        Projective transformation matrix.

    Returns
    -------
    A tuple where the first element is the affine matrix and the second element is the projective shift.

    Let P be the initial projective matrix and A the affine matrix. The projective shift S is defined as: S = A^-1 @ P.
    """
    assert np.allclose(projective_matrix[2, 2], 1), "A projective matrix should have a 1 in the bottom right corner."
    affine_matrix = projective_matrix.copy()
    affine_matrix[2] = [0, 0, 1]
    # equivalent to np.linalg.inv(affine_matrix) @ projective_matrix, but more numerically stable
    projective_shift = np.linalg.solve(affine_matrix, projective_matrix)
    projective_shift /= projective_shift[2, 2]
    return affine_matrix, projective_shift


def _get_filename_prefix(path: Path, dataset_id: str) -> str:
    if (path / f"{dataset_id}_{VisiumHDKeys.FEATURE_SLICE_FILE.value}").exists():
        return f"{dataset_id}_"
    assert (path / VisiumHDKeys.FEATURE_SLICE_FILE.value).exists(), (
        f"Cannot locate the feature slice file, please ensure the file is present in the {path} directory and/or adjust"
        "the `dataset_id` parameter"
    )
    return ""


def _parse_metadata(path: Path, filename_prefix: str) -> tuple[dict[str, Any], dict[str, Any]]:
    with h5py.File(path / f"{filename_prefix}{VisiumHDKeys.FEATURE_SLICE_FILE.value}", "r") as f5:
        metadata = json.loads(dict(f5.attrs)[VisiumHDKeys.METADATA_JSON])
        hd_layout = json.loads(metadata[VisiumHDKeys.HD_LAYOUT_JSON])
    return metadata, hd_layout


def _get_transform_matrices(metadata: dict[str, Any], hd_layout: dict[str, Any]) -> dict[str, ArrayLike]:
    """Gets 4 projective transformation matrices, describing how to align the CytAssist, spots and microscope coordinates.

    Parameters
    ----------
    metadata
        Metadata of the Visium HD dataset parsed using `_parse_metadata()` from the feature slice file.
    hd_layout
        Layout of the Visium HD dataset parsed using `_parse_metadata()` from the feature slice file.

    Returns
    -------
    A dictionary containing four projective transformation matrices:
    - CytAssist col/row to Spot col/row
    - Spot col/row to CytAssist col/row
    - Microscope col/row to Spot col/row
    - Spot col/row to Microscope col/row
    """
    transform_matrices = {}

    # this transformation is parsed but not used in the current implementation
    transform_matrices["hd_layout_transform"] = np.array(hd_layout[VisiumHDKeys.TRANSFORM]).reshape(3, 3)

    for key in [
        VisiumHDKeys.CYTASSIST_COLROW_TO_SPOT_COLROW,
        VisiumHDKeys.SPOT_COLROW_TO_CYTASSIST_COLROW,
        VisiumHDKeys.MICROSCOPE_COLROW_TO_SPOT_COLROW,
        VisiumHDKeys.SPOT_COLROW_TO_MICROSCOPE_COLROW,
    ]:
        coefficients = metadata[VisiumHDKeys.TRANSFORM_MATRICES][key]
        transform_matrices[key] = np.array(coefficients).reshape(3, 3)

    return transform_matrices


def _make_filtered_nucleus_adata(
    filtered_matrix_h5_path: Path,
    barcode_mappings_parquet_path: Path,
    bin_col_name: str = "square_002um",
    aggregate_col_name: str = "cell_id",
) -> AnnData:
    """Generate a filtered AnnData object by aggregating 2um binned data based on nucleus segmentation.

    Uses filtered_feature_bc_matrix.h5 file and a barcode_mappings.parquet file containing
    barcode mappings, filters the data to include only valid nucleus mappings,
    and aggregates the data based on specified bin into cell IDs which only contain
    the 2um square data under segmented nuclei.

    Parameters:
    -----------
    filtered_matrix_h5_path
        Path to the 10x Genomics HDF5 matrix file.
    barcode_mappings_parquet_path
        Path to the Parquet file containing barcode mappings.
    bin_col_name
        Column name in the barcode mappings that specifies the spatial bin (default is 'square_002um').
    aggregate_col_name
        Column name in the barcode mappings that specifies the aggregate cell ID (default is 'cell_id').

    Returns:
    --------
    AnnData
        An AnnData object where the observations correspond to filtered cell IDs
        and the variables correspond to the original features from the input data.
    """
    # Read in the necessary files
    adata_2um = sc.read_10x_h5(filtered_matrix_h5_path)
    barcode_mappings = pq.read_table(barcode_mappings_parquet_path)

    # Filter to only include valid cell IDs that are in both nucleus and cell
    barcode_mappings = barcode_mappings.filter(
        (barcode_mappings["cell_id"].is_valid()) and barcode_mappings["in_nucleus"]
    )

    # Filter the 2um adata to only include squares present in the barcode mappings
    valid_squares = barcode_mappings[bin_col_name].unique()
    squares_to_keep = np.intersect1d(adata_2um.obs_names, valid_squares)
    adata_filtered = adata_2um[squares_to_keep, :].copy()

    # Map each square to its corresponding cell ID
    square_to_cell_map = dict(
        zip(
            barcode_mappings[bin_col_name].to_pylist(),
            barcode_mappings[aggregate_col_name].to_pylist(),
            strict=True,
        )
    )
    ordered_cell_ids = [square_to_cell_map[square] for square in adata_filtered.obs_names]
    unique_cells = list(dict.fromkeys(ordered_cell_ids).keys())
    cell_to_idx = {cell: i for i, cell in enumerate(unique_cells)}

    # Make the aggregation matrix
    col_indices = [cell_to_idx[cell] for cell in ordered_cell_ids]
    row_indices = np.arange(len(ordered_cell_ids))
    data = np.ones_like(row_indices)

    aggregation_matrix = csc_matrix(
        (data, (row_indices, col_indices)),
        shape=(adata_filtered.n_obs, len(unique_cells)),
    )

    # Make the final AnnData object where cell IDs are filtered
    # to the data under the segmented nuclei
    nucleus_matrix_sparse = adata_filtered.X.T.dot(aggregation_matrix)
    adata_nucleus = sc.AnnData(nucleus_matrix_sparse.T)
    adata_nucleus.obs_names = unique_cells
    adata_nucleus.var = adata_filtered.var

    return adata_nucleus


def _extract_geometries_from_geojson(
    adata: AnnData,
    geojson_path: Path,
) -> GeoDataFrame:
    """Extract geometries and create a GeoDataFrame from a GeoJSON features map.

    Parameters
    ----------
    cell_adata
        AnnData object containing cell data.
    geojson_path
        Path to the GeoJSON file containing cell segmentation geometries.

    Returns
    -------
    GeoDataFrame
        A GeoDataFrame containing cell IDs and their corresponding geometries.
    """
    with open(geojson_path) as f:
        geojson_data = json.load(f)
    geojson_features_map: dict[str, Any] = {
        f"cellid_{feature['properties']['cell_id']:09d}-1": feature for feature in geojson_data["features"]
    }

    geometries = []
    cell_ids_ordered = []

    for obs_index_str in adata.obs.index:
        feature = geojson_features_map.get(obs_index_str)
        if feature:
            polygon_coords = np.array(feature["geometry"]["coordinates"][0])
            geometries.append(Polygon(polygon_coords))
            cell_ids_ordered.append(obs_index_str)

    return GeoDataFrame({"cell_id": cell_ids_ordered, "geometry": geometries}, index=cell_ids_ordered)