Skip to content

kmeans_img

Performs KMeans clustering on the input image.

Parameters:

Name Type Description Default
img ndarray

Image to cluster. Shape (H, W, 3).

required
n_clust int

Number of clusters.

3
seed int

Random seed.

42

Returns:

Type Description
ndarray

np.ndarray: Label image. Shape (H, W).

Source code in src/histolytics/stroma_feats/utils.py
def kmeans_img(img: np.ndarray, n_clust: int = 3, seed: int = 42) -> np.ndarray:
    """Performs KMeans clustering on the input image.

    Parameters:
        img (np.ndarray):
            Image to cluster. Shape (H, W, 3).
        n_clust (int):
            Number of clusters.
        seed (int):
            Random seed.

    Returns:
        np.ndarray:
            Label image. Shape (H, W).
    """
    pixels = img.reshape(-1, 3)
    labs = np.zeros(pixels.shape[0])
    nonzero_inds = np.where(np.all(pixels != 0, axis=1))[0]
    nonzero_pixels = pixels[nonzero_inds]

    kmeans = KMeans(n_clusters=n_clust, random_state=seed).fit(nonzero_pixels)
    labs[nonzero_inds] = kmeans.labels_ + 1

    # Reshape the labels to the original image shape
    return labs.reshape(img.shape[:2])