Repository: mlcommons/GaNDLF

Issue:
Add ability to handle high compression file formats (e.g., jpg) for segmentation output
**Is your feature request related to a problem? Please describe.**
Currently, writing masks are written as 16-bit images at inference time: 
https://github.com/CBICA/GaNDLF/blob/f6d26038a9d6203941748e43a593663594a3b0c0/GANDLF/compute/forward_pass.py#L330

However, this produces blank images because only JPG-200 supports 16-bit pixels properly [[ref](https://users.wfu.edu/yipcw/atg/jpeg2000/)].

**Describe the solutions you'd like**
There are a few ways around this problem:
- do not support jpg output at all, since it seems like it could be an issue if a system does not have the appropriate jpeg2000 libraries installed, and if one of `[".jpg", ".jpeg"]` (or more) are detected in the following function, return `".png"`:
https://github.com/CBICA/GaNDLF/blob/f6d26038a9d6203941748e43a593663594a3b0c0/GANDLF/utils/generic.py#L71
- if jpg needs to be written, use `tifffile` (which has internal support for JPEG2000 standard since version `2022.2.2`) to write image instead of the default (i.e., `SimpleITK`)
- if jpg needs to be written, cast image to 8-bit, but ONLY if jpg is detected:
https://github.com/CBICA/GaNDLF/blob/f6d26038a9d6203941748e43a593663594a3b0c0/GANDLF/compute/forward_pass.py#L330

**Describe alternatives you've considered**
Easiest would be the first choice above. Second would be the more robust solution but take more time to implement, and the last would be somewhere in the middle but prone to breakage for cases with large number of labels.

**Additional context**
N.A.



Discussion:

Resolving patch:
diff --git a/GANDLF/compute/forward_pass.py b/GANDLF/compute/forward_pass.py
--- a/GANDLF/compute/forward_pass.py
+++ b/GANDLF/compute/forward_pass.py
@@ -306,6 +306,9 @@ def validate_network(
                         affine=subject["1"]["affine"].squeeze(0),
                     ).as_sitk()
                     ext = get_filename_extension_sanitized(subject["1"]["path"][0])
+                    jpg_detected = False
+                    if ext in [".jpg", ".jpeg"]:
+                        jpg_detected = True
                     pred_mask = output_prediction.numpy()
                     # '0' because validation/testing dataloader always has batch size of '1'
                     pred_mask = reverse_one_hot(
@@ -317,7 +320,11 @@ def validate_network(
                     for postprocessor in params["data_postprocessing"]:
                         pred_mask = global_postprocessing_dict[postprocessor](
                             pred_mask, params
-                        )
+                        ).numpy()
+                    if jpg_detected:
+                        pred_mask = pred_mask.astype(np.uint8)
+                    else:
+                        pred_mask = pred_mask.astype(np.uint16)
 
                     ## special case for 2D
                     if image.shape[-1] > 1:
@@ -326,8 +333,6 @@ def validate_network(
                         result_image = sitk.GetImageFromArray(pred_mask.squeeze(0))
                     result_image.CopyInformation(img_for_metadata)
 
-                    # cast as the same data type
-                    result_image = sitk.Cast(result_image, sitk.sitkUInt16)
                     # this handles cases that need resampling/resizing
                     if "resample" in params["data_preprocessing"]:
                         result_image = resample_image(
diff --git a/GANDLF/compute/inference_loop.py b/GANDLF/compute/inference_loop.py
--- a/GANDLF/compute/inference_loop.py
+++ b/GANDLF/compute/inference_loop.py
@@ -296,7 +296,7 @@ def inference_loop(
                         out_probs_map[n, ...] * 255 / out_probs_map[n, ...].max(),
                         dtype=np.uint8,
                     )
-                    heatmap = cv2.applyColorMap(image, cv2.COLORMAP_HOT)
+                    heatmap = cv2.applyColorMap(image, cv2.COLORMAP_JET)
                     cv2.imwrite(file_to_write, heatmap)
 
 
diff --git a/GANDLF/data/post_process/tensor.py b/GANDLF/data/post_process/tensor.py
--- a/GANDLF/data/post_process/tensor.py
+++ b/GANDLF/data/post_process/tensor.py
@@ -28,4 +28,4 @@ def get_mapped_label(input_tensor, params):
     for key, value in mapping.items():
         output[input_arr == key] = value
 
-    return output
+    return torch.from_numpy(output)


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -1537,7 +1537,7 @@ def test_one_hot_logic():
     parameters = {"data_postprocessing": {"mapping": {0: 0, 1: 1, 2: 5}}}
     mapped_output = get_mapped_label(
         torch.from_numpy(img_tensor_oh_rev_array), parameters
-    )
+    ).numpy()
 
     for key, value in parameters["data_postprocessing"]["mapping"].items():
         comparison = (img_tensor_oh_rev_array == key) == (mapped_output == value)