Skip to content

rslearn.vis.utils

utils

Utility functions and constants for visualization.

read_vector_layer

read_vector_layer(window: Window, layer_name: str, layer_config: LayerConfig, group_idx: int = 0) -> list[Any]

Read a vector layer for visualization.

Parameters:

Name Type Description Default
window Window

The window to read from

required
layer_name str

The layer name

required
layer_config LayerConfig

The layer configuration

required
group_idx int

The item group index (default 0)

0

Returns:

Type Description
list[Any]

List of Feature objects

Source code in rslearn/vis/utils.py
def read_vector_layer(
    window: Window,
    layer_name: str,
    layer_config: LayerConfig,
    group_idx: int = 0,
) -> list[Any]:
    """Read a vector layer for visualization.

    Args:
        window: The window to read from
        layer_name: The layer name
        layer_config: The layer configuration
        group_idx: The item group index (default 0)

    Returns:
        List of Feature objects
    """
    if layer_config.type != LayerType.VECTOR:
        raise ValueError(f"Layer {layer_name} is not a vector layer")

    vector_format: VectorFormat = layer_config.instantiate_vector_format()
    logger.info(
        f"Reading vector layer {layer_name} group {group_idx} for window "
        f"{window.name}, bounds: {window.bounds}, projection: {window.projection}"
    )

    features = window.data.read_vector(
        layer_name,
        vector_format,
        group_idx=group_idx,
    )
    logger.info(f"Decoded {len(features)} features from vector layer {layer_name}")
    return features

generate_label_colors_for_layer

generate_label_colors_for_layer(layer_config: LayerConfig) -> dict[str, tuple[int, int, int]] | None

Generate label colors from a layer config's class_names.

Returns None if class_names is not set, signaling that the caller should fall back to DEFAULT_COLORS by index.

Parameters:

Name Type Description Default
layer_config LayerConfig

LayerConfig to read class_names from

required

Returns:

Type Description
dict[str, tuple[int, int, int]] | None

Dictionary mapping label class names to RGB color tuples, or None

Source code in rslearn/vis/utils.py
def generate_label_colors_for_layer(
    layer_config: LayerConfig,
) -> dict[str, tuple[int, int, int]] | None:
    """Generate label colors from a layer config's class_names.

    Returns None if class_names is not set, signaling that the caller should fall back
    to DEFAULT_COLORS by index.

    Args:
        layer_config: LayerConfig to read class_names from

    Returns:
        Dictionary mapping label class names to RGB color tuples, or None
    """
    if not layer_config.class_names:
        return None

    label_colors = {}
    for color_idx, label in enumerate(layer_config.class_names):
        label_colors[label] = DEFAULT_COLORS[color_idx % len(DEFAULT_COLORS)]
    return label_colors

format_window_info

format_window_info(window: Window) -> tuple[tuple[datetime, datetime] | None, float | None, float | None]

Extract window metadata for display.

Parameters:

Name Type Description Default
window Window

Window object

required

Returns:

Type Description
tuple[tuple[datetime, datetime] | None, float | None, float | None]

Tuple of (time_range, lat, lon) where time_range is a tuple of (start, end) datetime objects

Source code in rslearn/vis/utils.py
def format_window_info(
    window: Window,
) -> tuple[tuple[datetime, datetime] | None, float | None, float | None]:
    """Extract window metadata for display.

    Args:
        window: Window object

    Returns:
        Tuple of (time_range, lat, lon) where time_range is a tuple of (start, end) datetime objects
    """
    lat = None
    lon = None

    geom_wgs84 = window.get_geometry().to_projection(WGS84_PROJECTION)
    centroid = geom_wgs84.shp.centroid
    lon = float(centroid.x)
    lat = float(centroid.y)

    return window.time_range, lat, lon

array_to_bytes

array_to_bytes(array: ndarray, resampling: Resampling = NEAREST) -> bytes

Convert a numpy array to PNG bytes.

Parameters:

Name Type Description Default
array ndarray

Array with shape (height, width, channels) or (height, width) as uint8

required
resampling Resampling

PIL Image resampling method (default NEAREST)

NEAREST

Returns:

Type Description
bytes

PNG image bytes

Source code in rslearn/vis/utils.py
def array_to_bytes(
    array: np.ndarray, resampling: Image.Resampling = Image.Resampling.NEAREST
) -> bytes:
    """Convert a numpy array to PNG bytes.

    Args:
        array: Array with shape (height, width, channels) or (height, width) as uint8
        resampling: PIL Image resampling method (default NEAREST)

    Returns:
        PNG image bytes
    """
    if array.ndim == 2:
        img = Image.fromarray(array, mode="L")
    elif array.ndim == 3:
        if array.shape[-1] == 1:
            img = Image.fromarray(array[:, :, 0], mode="L")
        elif array.shape[-1] == 3:
            img = Image.fromarray(array, mode="RGB")
        else:
            img = Image.fromarray(array[:, :, :3], mode="RGB")
    else:
        raise ValueError(f"Unsupported array shape: {array.shape}")

    img = img.resize(VISUALIZATION_IMAGE_SIZE, resampling)

    buf = BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()