Merge 45083ac into cff2153

nicodv · Apr 14, 2019 · 4df4a16 · 4df4a16
2 parents cff2153 + 45083ac
commit 4df4a16
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/kmodes/kmodes.py b/kmodes/kmodes.py
@@ -423,6 +423,11 @@ def predict(self, X, **kwargs):
         labels : array, shape [n_samples,]
             Index of the cluster each sample belongs to.
         """
+
+        # Convert pandas objects to numpy arrays.
+        if 'pandas' in str(X.__class__):
+            X = X.values
+
         assert hasattr(self, '_enc_cluster_centroids'), "Model not yet fitted."
 
         if self.verbose and self.cat_dissim == ng_dissim:

diff --git a/kmodes/kprototypes.py b/kmodes/kprototypes.py
@@ -478,6 +478,11 @@ def predict(self, X, categorical=None):
             Index of the cluster each sample belongs to.
         """
         assert hasattr(self, '_enc_cluster_centroids'), "Model not yet fitted."
+
+        # Convert pandas objects to numpy arrays.
+        if 'pandas' in str(X.__class__):
+            X = X.values
+
         if categorical is not None:
             assert isinstance(categorical, (int, list, tuple)), "The 'categorical' \
                 argument needs to be an integer with the index of the categorical \