Repository: kornia/kornia

Issue:
Implement kmeans
## 🚀 Feature

Add a `kornia.contrib.kmeans` functionality that clusters batches of n-dimensional sets of features.

Reference: https://github.com/subhadarship/kmeans_pytorch/

#### Consider also to contribute to Kornia universe projects :)

<sub>

- [**Tutorials**](https://github.com/kornia/tutorials): our repository containing the tutorials.

</sub>



Discussion:
Hi @edgarriba I've raised a draft PR for this issue containing the implementation. Will push tests soon. 

Resolving patch:
diff --git a/kornia/contrib/__init__.py b/kornia/contrib/__init__.py
--- a/kornia/contrib/__init__.py
+++ b/kornia/contrib/__init__.py
@@ -13,6 +13,7 @@
 from .face_detection import *
 from .histogram_matching import histogram_matching, interp
 from .image_stitching import ImageStitcher
+from .kmeans import KMeans
 from .lambda_module import Lambda
 from .models.tiny_vit import TinyViT
 from .object_detection import ObjectDetector
@@ -39,4 +40,5 @@
     "DistanceTransform",
     "diamond_square",
     "ObjectDetector",
+    "KMeans",
 ]
diff --git a/kornia/contrib/kmeans.py b/kornia/contrib/kmeans.py
new file mode 100644
--- /dev/null
+++ b/kornia/contrib/kmeans.py
@@ -0,0 +1,179 @@
+# based on https://github.com/subhadarship/kmeans_pytorch
+
+from __future__ import annotations
+
+import torch
+
+from kornia.core import Tensor
+from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_SHAPE
+
+
+class KMeans:
+    """Implements the kmeans clustering algorithm with euclidean distance as similarity measure.
+
+    Args:
+        num_clusters: number of clusters the data has to be assigned to
+        cluster_centers: tensor of starting cluster centres can be passed instead of num_clusters
+        tolerance: float value. the algorithm terminates if the shift in centers is less than tolerance
+        max_iterations: number of iterations to run the algorithm for
+        seed: number to set torch manual seed for reproducibility
+
+    Example:
+        >>> kmeans = kornia.contrib.KMeans(3, None, 10e-4, 100, 0)
+        >>> kmeans.fit(torch.rand((1000, 5)))
+        >>> predictions = kmeans.predict(torch.rand((10, 5)))
+    """
+
+    def __init__(
+        self,
+        num_clusters: int,
+        cluster_centers: Tensor | None,
+        tolerance: float = 10e-4,
+        max_iterations: int = 0,
+        seed: int | None = None,
+    ) -> None:
+        KORNIA_CHECK(num_clusters != 0, "num_clusters can't be 0")
+
+        # cluster_centers should have only 2 dimensions
+        if cluster_centers is not None:
+            KORNIA_CHECK_SHAPE(cluster_centers, ["C", "D"])
+
+        self.num_clusters = num_clusters
+        self._cluster_centers = cluster_centers
+        self.tolerance = tolerance
+        self.max_iterations = max_iterations
+
+        self._final_cluster_assignments: None | Tensor = None
+        self._final_cluster_centers: None | Tensor = None
+
+        if seed is not None:
+            torch.manual_seed(seed)
+
+    @property
+    def cluster_centers(self) -> Tensor:
+        if isinstance(self._final_cluster_centers, Tensor):
+            return self._final_cluster_centers
+        if isinstance(self._cluster_centers, Tensor):
+            return self._cluster_centers
+        else:
+            raise TypeError("Model has not been fit to a dataset")
+
+    @property
+    def cluster_assignments(self) -> Tensor:
+        if isinstance(self._final_cluster_assignments, Tensor):
+            return self._final_cluster_assignments
+        else:
+            raise TypeError("Model has not been fit to a dataset")
+
+    def _initialise_cluster_centers(self, X: Tensor, num_clusters: int) -> Tensor:
+        """Chooses num_cluster points from X as the initial cluster centers.
+
+        Args:
+            X: 2D input tensor to be clustered
+            num_clusters: number of desired cluster centers
+
+        Returns:
+            2D Tensor with num_cluster rows
+        """
+        num_samples: int = len(X)
+        perm = torch.randperm(num_samples)
+        idx = perm[:num_clusters]
+        initial_state = X[idx]
+        return initial_state
+
+    def _pairwise_euclidean_distance(self, data1: Tensor, data2: Tensor) -> Tensor:
+        """Computes pairwise squared distance between 2 sets of vectors.
+
+        Args:
+            data1: 2D tensor of shape N, D
+            data2: 2D tensor of shape C, D
+
+        Returns:
+            2D tensor of shape N, C
+        """
+        # N*1*D
+        A = data1[:, None, ...]
+        # 1*C*D
+        B = data2[None, ...]
+        distance = (A - B) ** 2.0
+        # return N*C matrix for pairwise distance
+        distance = distance.sum(dim=-1)
+        return distance
+
+    def fit(self, X: Tensor) -> None:
+        """Iterative KMeans clustering till a threshold for shift in cluster centers or a maximum no of iterations
+        have reached.
+
+        Args:
+            X: 2D input tensor to be clustered
+        """
+        # X should have only 2 dimensions
+        KORNIA_CHECK_SHAPE(X, ["N", "D"])
+
+        if self._cluster_centers is None:
+            self._cluster_centers = self._initialise_cluster_centers(X, self.num_clusters)
+        else:
+            # X and cluster_centers should have same number of columns
+            KORNIA_CHECK(
+                X.shape[1] == self._cluster_centers.shape[1],
+                f"Dimensions at position 1 of X and cluster_centers do not match. \
+                {X.shape[1]} != {self._cluster_centers.shape[1]}",
+            )
+
+        # X = X.to(self.device)
+        current_centers = self._cluster_centers
+
+        previous_centers: Tensor | None = None
+        iteration: int = 0
+
+        while True:
+            # find distance between X and current_centers
+            distance: Tensor = self._pairwise_euclidean_distance(X, current_centers)
+
+            cluster_assignment = distance.argmin(-1)
+
+            previous_centers = current_centers.clone()
+
+            for index in range(self.num_clusters):
+                selected = torch.nonzero(cluster_assignment == index).squeeze()
+                selected = torch.index_select(X, 0, selected)
+                # edge case when a certain cluster centre has no points assigned to it
+                # just choose a random point as it's update
+                if selected.shape[0] == 0:
+                    selected = X[torch.randint(len(X), (1,))]
+                current_centers[index] = selected.mean(dim=0)
+
+            # sum of distance of how much the newly computed clusters have moved from their previous positions
+            center_shift = torch.sum(torch.sqrt(torch.sum((current_centers - previous_centers) ** 2, dim=1)))
+
+            iteration = iteration + 1
+
+            if self.tolerance is not None and center_shift**2 < self.tolerance:
+                break
+
+            if self.max_iterations != 0 and iteration >= self.max_iterations:
+                break
+
+        self._final_cluster_assignments = cluster_assignment
+        self._final_cluster_centers = current_centers
+
+    def predict(self, x: Tensor) -> Tensor:
+        """Find the cluster center closest to each point in x.
+
+        Args:
+            x: 2D tensor
+
+        Returns:
+            1D tensor containing cluster id assigned to each data point in x
+        """
+
+        # x and cluster_centers should have same number of columns
+        KORNIA_CHECK(
+            x.shape[1] == self.cluster_centers.shape[1],
+            f"Dimensions at position 1 of x and cluster_centers do not match. \
+                {x.shape[1]} != {self.cluster_centers.shape[1]}",
+        )
+
+        distance = self._pairwise_euclidean_distance(x, self.cluster_centers)
+        cluster_assignment = distance.argmin(-1)
+        return cluster_assignment


Test patch:
diff --git a/tests/contrib/test_kmeans.py b/tests/contrib/test_kmeans.py
new file mode 100644
--- /dev/null
+++ b/tests/contrib/test_kmeans.py
@@ -0,0 +1,133 @@
+import pytest
+import torch
+
+import kornia
+
+from testing.base import BaseTester
+
+
+class TestKMeans(BaseTester):
+    @pytest.mark.parametrize("num_clusters", [3, 10, 1])
+    @pytest.mark.parametrize("tolerance", [10e-4, 10e-5, 10e-1])
+    @pytest.mark.parametrize("max_iterations", [10, 1000])
+    def test_smoke(self, device, dtype, num_clusters, tolerance, max_iterations):
+        N = 1000
+        D = 2
+
+        kmeans = kornia.contrib.KMeans(num_clusters, None, tolerance, max_iterations, 0)
+        kmeans.fit(torch.rand((N, D), dtype=dtype, device=device))
+
+        out1 = kmeans.cluster_assignments
+        out2 = kmeans.cluster_centers
+
+        # output is of type tensor
+        assert isinstance(out1, torch.Tensor)
+        assert isinstance(out2, torch.Tensor)
+
+        # output is of same dtype
+        assert out1.dtype == torch.int64
+        assert out2.dtype == dtype
+
+    @pytest.mark.parametrize("num_clusters", [3])
+    @pytest.mark.parametrize("tolerance", [10e-4])
+    @pytest.mark.parametrize("max_iterations", [100])
+    def test_cardinality(self, device, dtype, num_clusters, tolerance, max_iterations):
+        N = 1000
+        D = 2
+
+        kmeans = kornia.contrib.KMeans(num_clusters, None, tolerance, max_iterations, 0)
+        kmeans.fit(torch.rand((N, D), dtype=dtype))
+
+        out1 = kmeans.cluster_assignments
+        out2 = kmeans.cluster_centers
+
+        # output is of correct shape
+        assert out1.shape == (N,)
+        assert out2.shape == (num_clusters, D)
+
+    def test_exception(self, device, dtype):
+        # case: cluster_center = 0:
+        with pytest.raises(Exception) as errinfo:
+            kornia.contrib.KMeans(0, None, 10e-4, 10, 0)
+        assert "num_clusters can't be 0" in str(errinfo)
+
+        # case: cluster centers is not a 2D tensor
+        with pytest.raises(TypeError) as errinfo:
+            starting_centers = torch.rand((2, 3, 5), device=device, dtype=dtype)
+            kmeans = kornia.contrib.KMeans(None, starting_centers, 10e-4, 100, 0)
+        assert "shape must be [['C', 'D']]. Got torch.Size([2, 3, 5])" in str(errinfo)
+
+        # case: input data is not a 2D tensor
+        with pytest.raises(TypeError) as errinfo:
+            kmeans = kornia.contrib.KMeans(3, None, 10e-4, 100, 0)
+            kmeans.fit(torch.rand((1000, 5, 60), dtype=dtype, device=device))
+        assert "shape must be [['N', 'D']]. Got torch.Size([1000, 5, 60])" in str(errinfo)
+
+        # case: column dimensions of cluster centers and data to be predicted do not match
+        with pytest.raises(Exception) as errinfo:
+            kmeans = kornia.contrib.KMeans(3, None, 10e-4, 100, 0)
+            kmeans.fit(torch.rand((1000, 5), dtype=dtype))
+            kmeans.predict(torch.rand((10, 7), dtype=dtype))
+        assert "7 != 5" in str(errinfo)
+
+    @staticmethod
+    def _create_data(device, dtype):
+        # create example dataset
+        torch.manual_seed(2023)
+        x = 5 * torch.randn((500, 2), dtype=dtype, device=device) + torch.tensor((-13, 17), dtype=dtype, device=device)
+        x = torch.vstack(
+            [x, torch.randn((500, 2), dtype=dtype, device=device) + torch.tensor((15, -12), dtype=dtype, device=device)]
+        )
+        x = torch.vstack(
+            [
+                x,
+                13 * torch.randn((500, 2), dtype=dtype, device=device)
+                + torch.tensor((35, 15), dtype=dtype, device=device),
+            ]
+        )
+        return x
+
+    def test_module(self, device, dtype):
+        x = TestKMeans._create_data(device, dtype)
+
+        kmeans = kornia.contrib.KMeans(3, None, 10e-4, 10000, 2023)
+        kmeans.fit(x)
+
+        centers = kmeans.cluster_centers
+        prediciton = kmeans.predict(torch.tensor([[-14, 16], [45, 12]], dtype=dtype, device=device)).tolist()
+
+        expected_centers = torch.tensor([[-13, 17], [15, -12], [35, 15]], dtype=dtype, device=device)
+        expected_prediciton = [0, 2]
+
+        # sorting centers using dimension 0 as key so that they can be checked for equalness
+        order = torch.argsort(centers[:, 0]).tolist()
+        new_classes = {old_class: new_class for new_class, old_class in enumerate(order)}
+
+        ordered_centers = centers[order]
+        oredered_prediciton = [new_classes[predicted_class] for predicted_class in prediciton]
+
+        self.assert_close(ordered_centers, expected_centers, atol=2, rtol=0.1)
+        assert oredered_prediciton == expected_prediciton
+
+    def test_dynamo(self, device, dtype, torch_optimizer):
+        x = TestKMeans._create_data(device, dtype)
+        kmeans_params = (3, None, 10e-4, 10000, 2023)
+        predict_param = torch.tensor([[-14, 16], [45, 12]])
+
+        kmeans = kornia.contrib.KMeans(*kmeans_params)
+        kmeans.fit(x)
+
+        centers = kmeans.cluster_centers
+        prediciton = kmeans.predict(predict_param)
+
+        kmeans_op = kornia.contrib.KMeans(*kmeans_params)
+        kmeans_op.fit = torch_optimizer(kmeans_op.fit)
+        kmeans_op.predict = torch_optimizer(kmeans_op.predict)
+
+        kmeans_op.fit(x)
+
+        centers_op = kmeans_op.cluster_centers
+        prediciton_op = kmeans_op.predict(predict_param)
+
+        self.assert_close(centers, centers_op)
+        self.assert_close(prediciton, prediciton_op)