labels updated to load all the feature vector elements in the mat files

Nolte-Group · Oct 1, 2025 · a94fffa · a94fffa
1 parent c751d57
commit a94fffa
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pipe/extract.py b/pipe/extract.py
@@ -27,7 +27,8 @@ def extract(spark: SparkSession) -> DataFrame:
     """
 
     path = Path("/app/workdir")
-    labels = []
+    labels = [DataKind.BB, DataKind.FPS, DataKind.NSD, DataKind.NCNT,
+              DataKind.SPEC, DataKind.TARGET, DataKind.TREATMENT]
     with open(path / "train.csv", "r") as file:
         for line in file:
             labels.append(line.strip().split(",")[0])