(Custom App) DeepEdit clicks are not transformed correctly

[matt3o]

Hey guys! I am currently trying to convert my sliding interactive model based on DeepEdit to work with MONAILabel. I got as far that the input is working and the label output are matching in terms of Nifti.
However while debugging I realized the clicks on the input image are not set / transformed correctly. More specifically when I click on a tumor on the right side, the network added a prediction right opposite to that on the left side.
So obviously the clicks do not fit the transformed input image. If I include the clicks in the transforms (setting the signal to 1 in the according input channel, DeepEdit style, image then has 1+len(labels) channels), I do get even weirder results where the click gets distorted among many pixels and in the end the network does no longer recognize it as a click.
I don't understand however how the previous DeepEdit Code is dealing with this problem.
I'll append my current preprocessing code. The setting below is not working where the guidance signal is added before the transforms. If I move AddGuidanceSignal to the end of the transforms, it works, but the clicks added are based on the original image whereas AddGuidanceSignal works on the transformed one..

Any help here would be appreciated, I'll try to find a working solution with the current transforms by only converting the signal or something.

def get_pre_transforms_val_as_list_monailabel(labels: Dict, device, args, input_keys=("image")):
    spacing = AUTOPET_SPACING if args.dataset == "AutoPET" else MSD_SPLEEN_SPACING
    cpu_device = torch.device("cpu")
    
    # Input keys have to be ["image", "label"] for train, and least ["image"] for val
    if args.dataset == "AutoPET":
        t_val = [ 
            # Initial transforms on the inputs done on the CPU which does not hurt since they are executed asynchronously and only once
            InitLoggerd(
                args
            ),  # necessary if the dataloader runs in an extra thread / process
            LoadImaged(keys=input_keys, reader="ITKReader", image_only=False),
            #ToTensord(keys=input_keys, device=cpu_device, track_meta=True),
            EnsureChannelFirstd(keys=input_keys),
            NormalizeLabelsInDatasetd(
                keys="label", labels=labels, device=cpu_device
            ),
            ScaleIntensityRanged(
                keys="image", a_min=0, a_max=43, b_min=0.0, b_max=1.0, clip=True
            ),  # 0.05 and 99.95 percentiles of the spleen HUs
            EnsureTyped(keys=input_keys, device=device, data_type='tensor'),
            AddEmptySignalChannels(keys=input_keys, device=device),
            AddGuidanceSignal(
                keys=input_keys,
                sigma=1,
                disks=True,
                device=device,
            ),
            Orientationd(keys=input_keys, axcodes="RAS"),
            Spacingd(
                keys=input_keys, pixdim=spacing
            ),  # 2-factor because of the spatial size
            CenterSpatialCropd(keys=input_keys, roi_size=args.val_crop_size)
            if args.val_crop_size is not None
            else NoOpd(),def get_pre_transforms_val_as_list_monailabel(labels: Dict, device, args, input_keys=("image")):

            DivisiblePadd(keys=input_keys, k=64, value=0)
            if args.inferer == "SimpleInferer"
            else NoOpd(),
            # EnsureTyped(keys=("image", "label"), device=cpu_device, track_meta=False),
            # PrintGPUUsaged(device=device, name="pre"),
        ]
    return t_val

[diazandr3s]

Hi @matt3o,

This is a great initiative! DeepEdit has a disadvantage in that it needs the whole volume to train. Using a sliding window or patch-based style for training is very much needed. However, this needs some major changes in the transforms and training workflow.

Here is the DeepEdit architecture: https://github.com/Project-MONAI/MONAILabel/wiki/DeepEdit#training-schema

We use two loops for training. The transforms for the click simulation after a first forward pass are the following: https://github.com/Project-MONAI/MONAILabel/blob/main/sample-apps/radiology/lib/trainers/deepedit.py#L83-L98

With regards to your comment:

I realized the clicks on the input image are not set / transformed correctly. More specifically when I click on a tumor on the right side, the network added a prediction right opposite to that on the left side.

Are you sure you have considered the correct image axis or image orientation for processing the clicks? Otherwise, I'm not understanding why this is happening in your case.

I'm happy to further discuss this over a video call if you think it is faster. We can later comment on what is discussed here :)

[matt3o]

Hey @diazandr3s , thanks for the quick response! I might have figured it out by changing the order of the transforms :)

I'll try to explain the problem again before I explain the solution: the clicks in the UI are being done on the initial non-transformed image. However when I add these clicks on the transformed image, of course they don't make any sense since quite a few transforms have been run. So what happened in my case when I clicked on the left side of the person, the network registered a click on the right side of the person (image got tranformed, clicks however not). You can see that well on the appended image. Clicks was on the right, however the small pixel on the left popped up after the inference -> For me this could only mean the Network got the click at the wrong position (which I verified with some debug Nifti output)

I actually already pasted the working transform solution above, yesterday evening I did not know it would work. The import step was moving ScaleIntensityRanged() above adding the Guidance signal.
Now the clicks get transformed together with the image and clicking works as expected.

I'll publish the code soon as well, it's part of my master thesis. I spent most of the time speeding up the previous code to do the transforms on the GPU, so that now the network runs the inference in less than 5 seconds for the full volume. Dice is at 87% on validation set with 10 clicks per label as additional input on AutoPET, which I am pretty happy about :)

[matt3o]

Btw I also saw the MONAILabel BasicInferTask always clears the CUDA cache. I am not sure if you have to do that. I don't know if you are following MONAI as well but I opened this bug report there: Project-MONAI/MONAI#6626
In essence: OOMs with MONAI in my case do not occur due to torch not clearing it's cache. Rather the gc does not clean up the GPU tensor pointers in time so that they keep on using memory without being used.
I would expect that the reason for the torch.cuda.empty_cache() is here for a similar reason - to prevent OOMs. I believe that a similar solution would be better here, to run gc.collect() twice, as the GarbageCollector does.
Maybe I am wrong, just wanted to tell you once, so that you have heard about it :)
Also if it is the torch memory itself, which is the issue, you might be able to fix by setting a lower gc threshold for torch with PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.8
which activates the torch GPU garbage collector automatically if the GPU memory usage is above 80%.

[diazandr3s]

Many thanks for reporting this, @matt3o.
Looking forward to seeing your code public available!

Thanks again!

[matt3o]

@diazandr3s Is there any functionality to let monailabel track how long it takes for a user to finish an annotation (time between next sample and writing)?
I know that there is the log file from which I could probably extract the information, but that would be super tedious. Optimal would be a csv file or similar with the time per image and a total time at the bottom. If possible I'd like to go one step further by also measuring how much interaction vs waiting / machine time there was. So to separate the time the person has to annotate vs the time to wait until MONAILabel server returns the new segmentation.
Also if there isn't, do you think it is rather easy to implement something like that?

Bonus: I was wondering why 3D Slicer blocks while waiting for the new segmentation / inference result. Are there any plans to make that non-blocking or does Slicer not support this behaviour?

[diazandr3s]

Hi @matt3o,

This is a great question. The current implementation only tracks the time it takes for the pre-transform, inference, inverse transform and post transforms: https://github.com/Project-MONAI/MONAILabel/blob/main/monailabel/tasks/infer/basic_infer.py#L345-L353

The time taken between the next sample and writing is not implemented. But I think it may be a good idea if the annotator is always in front of the screen annotating. You should also consider the time it takes to transfer the predicted mask to the viewer, which may be a bit difficult to track. You can assume that time is small though.

[matt3o]

So I dug a bit around in the code and it is kind of implemented for the MONAILabel Slicer Plugin. At least it does record how long Deepgrow took in the logs which matches the waiting time that I have had. Even though the network (server side) evaluates in 3 seconds there is still the transfer time of the segmentation which is at 7 seconds for the 50 Mb image that I am seeing. Maybe there is some potential for improvement here by compressing the file, I don't know.

Since we start our user study this week, I will just evaluate the logs by hand. For the future a "Record annotation time" button (plus a select csv output file) would be super cool. If I still have the time at the end of my thesis I might implement that.

For the server side, I did look into what the functions outside of basic_infer do, since that is the only place where the timing is collected. However all the request overhead and the writing in the end are completely neglectable, so those don't matter. The transit time for the image matters a lot but as far as I can see there is no easy way to see how long it takes from the server side.
Think it makes sense to make this a client-side functionality.

[diazandr3s]

Thanks for sharing all your findings, @matt3o.
I agree with you. It makes sense to have a function or something in the Slicer plugin for this. Those times can be combined with the ones tracked for inference in the server, which can also be obtained in the client side.