[Python] sample_near_center_of_mass(mask, num_samples, distance)

sample_near_center_of_mass.py

from scipy.ndimage import center_of_mass
from typing import List

Vec2 = Tuple[float, float]


def sample_near_center_of_mass(
    mask: NDArray[np.uint8], num_samples: int = 10, distance: float = 5
) -> List[Vec2]:
    """
    Sample points near the center of mass of a binary mask.

    Parameters:
    - mask: 2D numpy array representing the binary mask.
    - num_samples: Number of points to sample.
    - distance: Maximum distance from the center of mass for the points to be sampled.

    Returns:
    - samples: List of (x, y) tuples representing the sampled points.
    """
    # Get the center of mass of the mask
    y, x = cast(Vec2, center_of_mass(mask))

    # Generate random points within a certain distance from the center of mass
    samples: List[Vec2] = []
    for _ in range(num_samples):
        dx, dy = np.random.uniform(-distance, distance, size=2)
        sample_x, sample_y = int(x + dx), int(y + dy)

        # Check if the sampled point is within the mask
        if (
            0 <= sample_x < mask.shape[1]
            and 0 <= sample_y < mask.shape[0]
            and mask[sample_y, sample_x]
        ):
            samples.append((sample_x, sample_y))

    return samples