ENHANCE-PET / PUMA

Tool to multiplex different tracer PET data from the same patient.
10 stars 2 forks source link

Feature: User-Defined Color Assignment for PET Tracer Images in Multiplexing Process #13

Closed LalithShiyam closed 11 months ago

LalithShiyam commented 11 months ago

Summary

I propose the addition of a feature in the PUMAZ tool that enables users to manually assign specific color channels (Red, Green, or Blue) to PET tracer images during the multiplexing process. This feature, which is optional and activated via a command line argument, should bypass the alignment step when used.

Detailed Explanation

In the current version of PUMAZ, color channels for tracer images are assigned automatically during the multiplexing process. However, in certain cases, the automatic assignments may not satisfy specific analytical requirements. To improve this, I recommend a feature that allows users to select a color channel (R, G, or B) for each tracer, thus offering enhanced control and flexibility in image visualization.

Key Features:

Use Case

This feature is crucial for instances where default color assignments do not provide clear or meaningful visual distinctions, allowing for tailored visualizations essential for enhanced interpretation and analysis of PET images.

Implementation Suggestion

Impact

This enhancement will significantly increase the usability of PUMAZ in medical image analysis, particularly in customizing the visualization of PET images for improved clarity and interpretability.

Proposed Code Modifications

Addition in pumaz.py (Main Script)

# Add a new command-line argument for custom color assignments
parser.add_argument("--custom-colors", action="store_true", 
                    help="Manually assign colors to tracer images")

args = parser.parse_args()

if args.custom_colors:
    # Bypass alignment and directly proceed to custom color multiplexing
    aligned_pt_dir = os.path.join(puma_dir, constants.ALIGNED_PET_FOLDER) ## get this from the user via argparse.
    color_channel_assignments = get_color_channel_assignments(aligned_pt_dir)
    image_processing.multiplex(aligned_pt_dir, '*nii*', 'PET', 
                               os.path.join(aligned_pt_dir, constants.MULTIPLEXED_COMPOSITE_IMAGE), 
                               color_channel_assignments=color_channel_assignments)
else:
    # Regular workflow
    # ... existing code for alignment and default multiplexing ...

Function for Obtaining Custom Color Channel Assignments in pumaz.py

def get_color_channel_assignments(aligned_pt_dir):
    tracer_images = [img for img in os.listdir(aligned_pt_dir) if img.endswith(('.nii', '.nii.gz'))]
    color_channel_assignments = {}
    channel_map = {'R': 0, 'G': 1, 'B': 2}
    for tracer in tracer_images:
        color_channel = input(f"Assign a color channel (R, G, B) for {tracer}: ").upper()
        while color_channel not in channel_map:
            print("Invalid input. Please enter 'R', 'G', or 'B'.")
            color_channel = input(f"Assign a color channel (R, G, B) for {tracer}: ").upper()
        color_channel_assignments[tracer] = channel_map[color_channel]
    return color_channel_assignments

Modifications in image_processing.py

def multiplex(directory, extension, modality, output_image_path, color_channel_assignments=None):
    nifti_files = get_files(directory, extension)
    if color_channel_assignments:
        blend_images_with_color_channels(nifti_files, color_channel_assignments, output_image_path)
    else:
        modalities = [modality] * len(nifti_files)
        blend_images(nifti_files, modalities, output_image_path)

_New Function for Custom Color Channel Blending in image_processing.py_


def blend_images_with_color_channels(image_paths, color_channel_assignments, output_path):
    if len(image_paths) != len(color_channel_assignments):
        raise ValueError("Number of images and color channel assignments do not match.")

    images = np.zeros((*nib.load(image_paths[0]).get_fdata().shape, 3))  # Initialize composite image

    for path in image_paths:
        img = nib.load(path).get_fdata()
        channel_index = color_channel_assignments[path]
        images[..., channel_index] += normalize_data(img)

    images = np.clip(images, 0, 1)  # Ensure values are within range after accumulation
    composite_image = nib.Nifti1Image(images, nib.load(image_paths[0]).affine)
    nib.save(composite_image, output_path)
LalithShiyam commented 11 months ago

@Keyn34 Kindly have a look at this. I think it's a reasonable way to provide users with some flexibility and don't take these codes as cut in stone.

Keyn34 commented 11 months ago

I like the idea of implementing this as a "post-processing" step instead of bending the regular workflow of pumaz to respect assignments. This design makes our lives and the lives of users significantly easier.

I would also propose to make the names of the generated aligned images easier to understand. For example, the name could become something like {tracer}_{channel}{colour}.nii.gz. This would also help users to understand the assignments better.

I will work on the implementation.

LalithShiyam commented 11 months ago

I would also propose to make the names of the generated aligned images easier to understand. For example, the name could become something like {tracer}{channel}{colour}.nii.gz. This would also help users to understand the assignments better._

The aligned names are generated from the folder name and I am not sure if we can get the tracer name without dicom and I am not a fan of making the folder structure naming rigid (after our lessons from moose). I would let the users decide on the color channel they prefer by asking them what they want from Rgb. We are just deriving a multiplexed image and not providing the individual rgb label anyway.

LalithShiyam commented 11 months ago

Regarding the postprocessing, sure - as long as the user can use puma as a CLI tool or a library. Either we can pass the directory with the aligned images and trigger an interactive terminal session or create an input argument with -- custom_colors tracer_01.nii.gz:red tracer_02.nii.gz:green tracer_03.nii.gz:blue But I would prefer to wrap this up by next week. Let's reconvene during the mid of next week to discuss progress, if there is a hiccup, let's sort it out so that we meet the deadlines.

Keyn34 commented 11 months ago

Fair point. We can't get tracer association when we don't have the dicom. And even if the dicom does exist, it might not be clear by tags.

As far as I understood it, we now have two options:

  1. We let the user process the images by puma as usual. After alignment is done, the user can run puma with the --custom_colors, followed by the <aligned_by_puma.nii.gz>:<colour> dictionary again (like your original post).
  2. The user runs puma with --custom_colors and is prompted during multiplexing which image receives which colour channel (like your second comment).

    Either way we do it, the implementation is similar. I will prepare the functions, and we take it from there on Wednesday.

Keyn34 commented 11 months ago

I will close this for now, as this was implemented by commit 39eea74.