Advice on "fill_missing data" function

[skyshi0]

Hi,
I'm using the fill_missing function in the analysis example. However, I found sometimes the plot is 'jumpy'. I checked the prediction and I found it was because some frames in the new video are predicted very far from the actual location or previous frame location. Since the wrong prediction is only for few frames, and usually the next frame is predicted correctly, I wonder if there's code can help to address this problem, or if SLEAP can avoid this problem(like set a range for the position change in pixels between each frame).

from scipy.interpolate import interp1d

def fill_missing(Y, kind="linear"):
    """Fills missing values independently along each dimension after the first."""

    # Store initial shape.
    initial_shape = Y.shape

    # Flatten after first dim.
    Y = Y.reshape((initial_shape[0], -1))

    # Interpolate along each slice.
    for i in range(Y.shape[-1]):
        y = Y[:, i]

        # Build interpolant.
        x = np.flatnonzero(~np.isnan(y))
        f = interp1d(x, y[x], kind=kind, fill_value=np.nan, bounds_error=False)

        # Fill missing
        xq = np.flatnonzero(np.isnan(y))
        y[xq] = f(xq)

        # Fill leading or trailing NaNs with the nearest non-NaN values
        mask = np.isnan(y)
        y[mask] = np.interp(np.flatnonzero(mask), np.flatnonzero(~mask), y[~mask])

        # Save slice
        Y[:, i] = y

    # Restore to initial shape.
    Y = Y.reshape(initial_shape)

    return Y

[talmo]

Hi @skyshi0,

Apologies that we missed your post! November was an exceptionally busy month over here.

That function probably won't handle all the cases but this kind of artifact can definitely be handled in postprocessing.

What I'd do here is to apply a median filter (window = 3 or 5) which should take care of these point discontinuities. Here's one that ChatGPT generated in the same style as fill_missing:

from scipy.ndimage import median_filter
import numpy as np

def apply_median_filter(Y, window_size=3):
    """Applies a median filter independently along each dimension after the first.

    Parameters:
    Y : np.ndarray
        Input array, assumed to have missing values as NaNs.
    window_size : int
        Size of the median filter window. Must be a positive odd integer.

    Returns:
    np.ndarray
        Median-filtered array with the same shape as input.
    """
    if window_size < 1 or window_size % 2 == 0:
        raise ValueError("Window size must be a positive odd integer.")

    # Store initial shape.
    initial_shape = Y.shape

    # Flatten after the first dimension.
    Y = Y.reshape((initial_shape[0], -1))

    # Process each slice independently.
    for i in range(Y.shape[-1]):
        y = Y[:, i]

        # Apply median filter.
        filtered = median_filter(y, size=window_size)

        # Save slice back.
        Y[:, i] = filtered

    # Restore the original shape.
    Y = Y.reshape(initial_shape)

    return Y

This will break if you have NaNs though. You could customize the behavior to take that into account, but at that point I'd recommend just using a library like movement which can load SLEAP data and has a great toolkit for exactly these kinds of tasks.

One of their example notebooks even does what you're looking for: https://movement.neuroinformatics.dev/examples/smooth.html

Cheers,

Talmo