Repository: mlcommons/GaNDLF

Issue:
Upgrade reverse one hot logic to get unique labels as output for combined class list 
**Is your feature request related to a problem? Please describe.**
For `class_list: '[0,1||2||3,1||4,4]'`, the current reverse one hot logic does output each class as a unique index, which causes an issue when doing downstream analysis, because the current logic involves adding the mask [[ref](https://github.com/CBICA/GaNDLF/blob/02b3778517225fe7b6c78fe6f6b374c031f38b1a/GANDLF/utils/tensor.py#L137)]: 
```python
final_mask += np.asarray(array_to_consider[i, ...], dtype=int)
```

**Describe the solution you'd like**
This should be updated to include the index of the `class_list`, which would give unique values as output:

For `class_list: '[0,1||2||3,1||4,4]'`, the final predicted mask should be formed with the following logic:

```
for each class in class_list
   unless the class is 0, the output will be (index * predicted_mask_array at that class)
```

**Describe alternatives you've considered**
N.A.

**Additional context**
The predicted masks can then be post-processed (this step can be specific for the application at hand, since LiTS can have a different un-combine logic than BraTS).



Discussion:

Resolving patch:
diff --git a/GANDLF/parseConfig.py b/GANDLF/parseConfig.py
--- a/GANDLF/parseConfig.py
+++ b/GANDLF/parseConfig.py
@@ -495,12 +495,9 @@ def parseConfig(config_file_path, version_check_flag=True):
                 "WARNING: This is a special case for multi-class computation, where different labels are processed together, `reverse_one_hot` will need mapping information to work correctly"
             )
             temp_classList = params["model"]["class_list"]
-            temp_classList = temp_classList.replace(
-                "[", ""
-            )  # we don't need the brackets
-            temp_classList = temp_classList.replace(
-                "]", ""
-            )  # we don't need the brackets
+            # we don't need the brackets
+            temp_classList = temp_classList.replace("[", "")
+            temp_classList = temp_classList.replace("]", "")
             params["model"]["class_list"] = temp_classList.split(",")
         else:
             try:
diff --git a/GANDLF/utils/tensor.py b/GANDLF/utils/tensor.py
--- a/GANDLF/utils/tensor.py
+++ b/GANDLF/utils/tensor.py
@@ -103,38 +103,34 @@ def reverse_one_hot(predmask_array, class_list):
             if isinstance(_class, str):
                 if case in _class:  # check if any of the special cases are present
                     special_case_detected = True
-                    # # if present, then split the sub-class
-                    # class_split = _class.split(case)
-                    # for i in class_split:  # find the max for computation later on
-                    #     if int(i) > max_current:
-                    #         max_current = int(i)
 
     final_mask = None
     if special_case_detected:
-        for i, _ in enumerate(class_list):
+        # for this case, the output mask will be formed with the following logic:
+        # for each class in class_list
+        #   unless the class is 0, the output will be index * predmask_array at that class
+        array_to_consider_bool = array_to_consider.astype(bool)
+        for idx, i in enumerate(class_list):
             initialize_mask = False
-            if isinstance(class_list[i], str):
+            if isinstance(class_list[idx], str):
                 for case in special_cases_to_check:
                     initialize_mask = False
-                    if case in class_list[i]:
+                    if case in class_list[idx]:
                         initialize_mask = True
                     else:
-                        if class_list[i] != "0":
+                        if class_list[idx] != "0":
                             initialize_mask = True
             else:
-                if class_list[i] != 0:
+                if class_list[idx] != 0:
                     initialize_mask = True
 
             if initialize_mask:
                 if final_mask is None:
-                    final_mask = np.asarray(array_to_consider[i, ...], dtype=int)
+                    final_mask = idx * np.asarray(
+                        array_to_consider_bool[idx, ...], dtype=int
+                    )
                 else:
-                    # TODO: this should be replaced with information coming from the config that defines what the specific union of labels should be defined as
-                    # for example:
-                    # class_list = "[0,1||2]"
-                    # output_classes = [0,1]
-                    # this would make "1||2" be mapped to "1", and this mechanism can be extended for N possible combinations
-                    final_mask += np.asarray(array_to_consider[i, ...], dtype=int)
+                    final_mask[array_to_consider_bool[idx, ...]] = idx
     else:
         for i, _ in enumerate(class_list):
             if final_mask is None:


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -1223,21 +1223,24 @@ def test_one_hot_logic():
     class_list = [0, "1||2||3", 4]
     img_tensor_oh = one_hot(img_tensor, class_list)
     img_tensor_oh_rev_array = reverse_one_hot(img_tensor_oh[0], class_list)
+
+    # check for background
     comparison = (random_array == 0) == (img_tensor_oh_rev_array == 0)
     assert comparison.all(), "Arrays at '0' are not equal"
 
-    random_array_sp = (random_array == 1) + (random_array == 2) + (random_array == 3)
-    img_tensor_oh_rev_array_sp = img_tensor_oh_rev_array == 1
-    img_tensor_oh_rev_array_sp[random_array == 4] = False
-    comparison = random_array_sp == img_tensor_oh_rev_array_sp
-    assert comparison.all(), "Special arrays are not equal"
-
-    # check for '4'
-    img_tensor_oh_rev_array_sp = img_tensor_oh_rev_array == 1
-    for i in [1, 2, 3]:
-        img_tensor_oh_rev_array_sp[random_array == i] = False
-    comparison = (random_array == 4) == img_tensor_oh_rev_array_sp
-    assert comparison.all(), "Arrays at '4' are not equal"
+    # check last foreground
+    comparison = (random_array == np.max(random_array)) == (
+        img_tensor_oh_rev_array == len(class_list) - 1
+    )
+    assert comparison.all(), "Arrays at final foreground are not equal"
+
+    # check combined foreground
+    combined_array = np.logical_or(
+        np.logical_or((random_array == 1), (random_array == 2)), (random_array == 3)
+    )
+    comparison = combined_array == (img_tensor_oh_rev_array == 1)
+    assert comparison.all(), "Arrays at the combined foreground are not equal"
+
     print("passed")