Preprocessing fix for binary classification

ITMO-NSS-team · Nov 13, 2023 · c0055f9 · c0055f9
1 parent a5a9dd3
commit c0055f9
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nas/data/loader.py b/nas/data/loader.py
@@ -6,7 +6,7 @@
 import cv2
 import numpy as np
 from fedot.core.data.data import InputData
-from sklearn.preprocessing import OneHotEncoder
+from sklearn.preprocessing import OneHotEncoder, LabelEncoder
 
 
 class BaseDataLoader(ABC):
@@ -25,7 +25,9 @@ def target(self):
 
         if self.dataset.num_classes > 2:
             encoder = OneHotEncoder(handle_unknown='error', dtype=int, sparse_output=False)
-            new_targets = encoder.fit_transform(new_targets)
+        else:
+            encoder = LabelEncoder()
+        new_targets = encoder.fit_transform(new_targets)
         return new_targets
 
     @abstractmethod