Fix for loading of single sample

Author: Svdvoort

PrognosAIs/IO/LabelParser.py

@@ -54,8 +54,18 @@ def get_labels(self) -> list:
         Returns:
             labels: List of labels
         """
+        labels = np.squeeze(self.label_data.values)
+
+        if isinstance(labels, np.ndarray) and labels.size > 1:
+            labels = labels.tolist()
+        elif isinstance(labels, np.ndarray):
+            # Otherwise if it is 1 element it will remove the list,
+            # and return only a string
+            labels = [labels.tolist()]
+        else:
+            labels = [labels]
 
-        return np.squeeze(self.label_data.values).tolist()
+        return labels
 
     def get_samples(self) -> list:
         """Get all labels of all samples
@@ -65,7 +75,14 @@ def get_samples(self) -> list:
         Returns:
             samples: List of samples
         """
-        return np.squeeze(self.label_data.index).tolist()
+
+        samples = np.squeeze(self.label_data.index)
+
+        if isinstance(samples, np.ndarray):
+            samples = samples.tolist()
+        else:
+            samples = [samples]
+        return samples
 
     def get_data(self) -> dict:
         """Get all data from the label file

tests/test_data/NPZ_Data/labels_file_single_class_single_sample.txt

@@ -0,0 +1,2 @@
+Image	Label_1
+./PrognosAIs/tests/test_data/NPZ_Data/Samples/Test_sample_1.npz	./PrognosAIs/tests/test_data/NPZ_Data/Samples/Test_sample_1_label.npz

tests/test_labelparser_new.py

@@ -15,3 +15,22 @@ def test_one_hot_encoding_missing_labels():
     result = label_parser.get_labels_from_category("Label_1")
 
     assert result == pytest.approx(np.asarray([[0, 1], [-1, -1], [1, 0], [0, 1]]))
+
+def test_sample_loading_single_sample():
+    label_file = os.path.join(FIXTURE_DIR, "labels_file_single_class_single_sample.txt")
+    label_parser = LabelParser.LabelLoader(label_file)
+
+    result = label_parser.get_samples()
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+
+def test_label_loading_single_sample():
+    label_file = os.path.join(FIXTURE_DIR, "labels_file_single_class_single_sample.txt")
+    label_parser = LabelParser.LabelLoader(label_file)
+
+    result = label_parser.get_labels()
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+