post_processing_utils

dreem.inference.post_processing_utils

Functions:

Name	Description
get_principal_axis_with_fallback	Compute principal axis with automatic fallback.
register_principal_axis_method	Register a new principal axis computation method.

get_principal_axis_with_fallback(instance, front_nodes, back_nodes, crop, logger, frame_id)

Compute principal axis with automatic fallback.

Tries registered methods in priority order until one succeeds.

Parameters:

Name	Type	Description	Default
instance		Instance dict with pose keypoints	required
front_nodes	list[str] | None	List of front nodes	required
back_nodes	list[str] | None	List of back nodes	required
crop	Tensor | None	Optional crop tensor	required
logger		Logger instance	required
frame_id		Frame ID for logging	required

Returns:

Type	Description
tuple[Tensor, bool]	Tuple of (principal_axis_vector, success) where success indicates if the principal axis was successfully computed.

Source code in dreem/inference/post_processing_utils.py

def get_principal_axis_with_fallback(
    instance,
    front_nodes: list[str] | None,
    back_nodes: list[str] | None,
    crop: torch.Tensor | None,
    logger,
    frame_id,
) -> tuple[torch.Tensor, bool]:
    """Compute principal axis with automatic fallback.

    Tries registered methods in priority order until one succeeds.

    Args:
        instance: Instance dict with pose keypoints
        front_nodes: List of front nodes
        back_nodes: List of back nodes
        crop: Optional crop tensor
        logger: Logger instance
        frame_id: Frame ID for logging

    Returns:
        Tuple of (principal_axis_vector, success) where success indicates
        if the principal axis was successfully computed.
    """
    kwargs = {"front_nodes": front_nodes, "back_nodes": back_nodes, "crop": crop}

    for method_name, can_compute, compute in _PRINCIPAL_AXIS_METHODS:
        if can_compute(instance, **kwargs):
            result = compute(instance, **kwargs)
            if result is not None:
                return result, True
    return torch.full((2,), torch.nan, dtype=torch.float32), False

register_principal_axis_method(name, can_compute, compute, priority=None)

Register a new principal axis computation method.

Parameters:

Name	Type	Description	Default
can_compute	Callable	Function that checks if method can be used for an instance. Signature: can_compute(instance, **kwargs) -> bool	required
compute	Callable	Function that computes principal axis. Signature: compute(instance, **kwargs) -> torch.Tensor	required
name	str	Name of the method	required
priority	int | None	Optional priority index. If None, appends to end. Lower index = higher priority.	None

Source code in dreem/inference/post_processing_utils.py

def register_principal_axis_method(
    name: str, can_compute: Callable, compute: Callable, priority: int | None = None
) -> None:
    """Register a new principal axis computation method.

    Args:
        can_compute: Function that checks if method can be used for an instance.
                    Signature: can_compute(instance, **kwargs) -> bool
        compute: Function that computes principal axis.
                Signature: compute(instance, **kwargs) -> torch.Tensor
        name: Name of the method
        priority: Optional priority index. If None, appends to end. Lower index = higher priority.
    """
    entry = (name, can_compute, compute)
    if priority is None:
        _PRINCIPAL_AXIS_METHODS.append(entry)
    else:
        _PRINCIPAL_AXIS_METHODS.insert(priority, entry)