fix injector to use preprocessed data

mitre · Jul 28, 2023 · 5fa9d2e · 5fa9d2e
1 parent ea015f1
commit 5fa9d2e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/menelaus/injection/label_manipulation.py b/menelaus/injection/label_manipulation.py
@@ -158,7 +158,9 @@ class probability for 1 or more desired classes
 
         # locate each class in window
         for cls in all_classes:
-            cls_idx = np.where(data[:, target_col] == cls)[0]
+
+            cls_idx = None
+            cls_idx = np.where(ret[:, target_col] == cls)[0]
             cls_idx = cls_idx[(cls_idx < to_index) & (cls_idx >= from_index)]
 
             # each member has p_class / class_size chance, represented as bool to avoid div/0
@@ -178,7 +180,7 @@ class probability for 1 or more desired classes
         sample_idxs = np.random.choice(
             sample_idxs_grouped, to_index - from_index, True, self._p_distribution
         )
-        ret[from_index:to_index] = data[sample_idxs]
+        ret[from_index:to_index] = ret[sample_idxs]
 
         # handle data type and return
         ret = self._postprocess(ret)