How to do sliding window for detection?

[jhairgallardo]

Hello! I am new using Monai. I managed to follow the retinanet tutorial example and created my own detection code based on DETR. For a input image, my network output consist of a dictionary that contains class probabilities and bounding boxes.

For example, for batch size 1 (single input), my input shape is (1,1,96,96,96) , which means (batch, channel, w,h,d). The output of my network would be (assuming it outputs 2 bounding boxes and there are only 2 classes):
{ "pred_logits": this would be a tensor with shape (1,2,2), "boxes": this would be a tensor with shape (1,2,6) }

Where pre_logits is (batch, num_boxes, num_classes) and boxes is (batch, num_boxes, box_tensor_dimension).

I want to get predictions for a larger input image doing sliding window approach. Let's say for a (1,1,300,300,100). I tried using the sliding window method (SlidingWindowInferer) done in the detection tutorial for retinanet, but I get an error regarding the output dimensional size of my network not matching the input one. (input is 5 dimensional, values in the output dictionary are 3 dimensional). In the function definition (https://docs.monai.io/en/stable/inferers.html#sliding-window-inference-function) it says the following:

"e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes could be ([128,64,256], [64,32,128])."

Is there a way to get around this? I only need the sliding window to give me the outputs for each patch and the locations of each patch so I can modify the bounding boxes accordingly. Not sure if there is another function in Monai that does this.

Thank you!

[KumoLiu]

Hi @jhairgallardo, thanks for your interest here.
You could do something like this:

import torch
import monai

class NaiveNet(torch.nn.Module):
    def __init__(self, ):
        super().__init__()

    def forward(self, images: torch.Tensor):
        return (torch.randn(2, 1, 1, 2, 2), torch.randn(2, 1, 1, 2, 6)) # ensure they have the same spatial dim with the input.

# create a predictor
network = NaiveNet()
inferer = monai.inferers.SlidingWindowInferer(
    roi_size = (256, 256, 256),
    overlap = 0,
    cache_roi_weight_map = True,
)
images = torch.randn((2, 3, 512, 512, 512))  # a large input
out = inferer(images, network)

[jhairgallardo]

I managed to use SlidingWindowSplitter to create a generator that gives me each window subvolume and location. I did a for loop over that generator and take care of the box shift stuff. My code looks something like this (I am assuming batch size =1):

import torch
import monai

class NaiveNet(torch.nn.Module):
    def __init__(self, ):
        super().__init__()

    def forward(self, image: torch.Tensor):
        # assume boxes are in cccwhd format and that it predicts 5 boxes per image
        return {'pred_logits': torch.randn(1, 5, 2), 
                'pred_boxes': torch.randn(1,5,6)} 
                
model = NaiveNet()

splitter_funct = monai.inferers.SlidingWindowSplitter(patch_size=96, overlap=0.25)

input = torch.randn((1, 1, 512, 512, 512)) # (batch, channel, W,H,D)
split_input_generator = splitter_funct(input)

out= {'pred_logits': [], 'pred_boxes': []}
for subvolume, location in split_input_generator:

    output_sub = model(subvolume)

    # correct box according to subvolume location.
    pred_boxes = output_sub['pred_boxes']
    pred_boxes = pred_boxes + torch.tensor([location[0],location[1],location[2],0,0,0], dtype=torch.float32)

    # accumulate
    out['pred_boxes'].append(pred_boxes)
    out['pred_logits'].append(output_sub['pred_logits'])

# concat
out['pred_logits'] = torch.cat(out['pred_logits'], dim=0)
out['pred_boxes'] = torch.cat(out['pred_boxes'], dim=0)

I will also try the way you mentioned. I am not sure how SlidingWindowInferer would know to shift the boxes predictions accordinly to subvolumes locations.