TTA Seemingly Doesn't Affect Predictions

[malje16]

tldr: Adding a resizing augmentation to your mapper and thereby your dataloader doesn't affect the predictions made on the images from the dataloader. I would like to know why, and if that is intended.

The Long Version:
I'm working on a pcba dataset, where the objective is to automatically recognize different categories of components. I don't have a massive number of images, so in an attempt to improve performance I decided to augment my dataset. By studying the tutorial (and the sourcecode) I came up with a custom dataloader that adds a mapper with augmentations. It looks like this:

def My_train_aug(cfg):
augs = [T.ResizeShortestEdge((1800, 2000), 2600) ]
#augs.append(T.RandomExtent(scale_range=[0.4, 0.7], shift_range=[0.8, 0.8]))
augs.append(T.RandomFlip(prob=0.5, horizontal=True, vertical=False))
augs.append(T.RandomFlip(prob=0.5, horizontal=False, vertical=True))
augs.append(T.RandomBrightness(0.8, 1.2))
augs.append(T.RandomContrast(0.8,1.2))
augs.append(T.RandomSaturation(0.8, 1.2))
return augs

def My_test_aug(cfg):
augs = [T.ResizeShortestEdge((1800, 2000), 2600)]
return augs

class Trainer(DefaultTrainer):

@classmethod
def build_train_loader(cls, cfg):
train_mapper = DatasetMapper(cfg, is_train=True, augmentations=My_train_aug(cfg))
return build_detection_train_loader(cfg, mapper=train_mapper)

@classmethod
def build_test_loader(cls, cfg, dataset_name):
test_mapper = DatasetMapper(cfg, is_train=False, augmentations=My_test_aug(cfg))
return build_detection_test_loader(cfg, dataset_name, mapper=test_mapper)

Now, you may notice that the training augmentations include a RandomExtent augmentation that isn't used. Instead that line is a comment. Thats because my AP immediately fell to 0% when I added that augmentation. This is not the problem however (although if you have any ideas why that happens, do let me know). The problem occurs when i attempt to visualize the output of the model. I wrote a little method to visualize output it looks like this:

class Trainer(DefaultTrainer):

@classmethod
def output_visualizer(cls, cfg, model):
model.eval()
    data_loader = cls.build_test_loader(cfg, cfg.DATASETS.TEST[0])
with torch.no_grad():
    for inputs in data_loader:
        outputs = model(inputs)

        input_img = inputs[0]['image']
        prediction_instance = outputs[0]['instances']

        MyVisual = Visualizer(input_img.permute(1,2,0), instance_mode=1)
        VisImg_with_anns = MyVisual.draw_instance_predictions(prediction_instance)
        img_with_anns = cv2.cvtColor(VisImg_with_anns.get_image(), cv2.COLOR_BGR2RGB)
        plt.imshow(img_with_anns)
        plt.show()

And then I run a mainfile that creates a model from a "model_final.pth" file and then runs my method.

main():
register_my_datasets() # A function that registers and loads the custom datasets I'm using
cfg = get_cfg()
cfg.merge_from_file("<config_file_path>")

model = Trainer.build_model(cfg)
MyCheckpointer = DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR)
MyCheckpointer.load("<checkpoint_file_path/model_final.pth>")

Trainer.output_visualizer(cfg, model)

When I use the checkpoint from my RandomExtent augmentation session the predictions are useless, which is not really surprising given the 0% AP.
But then I went troubleshooting and created a dataset with a single image sampled from my original PCBA dataset. Training this tiny dataset naturally yields near perfect precision as I am massively overfitting, but when i run my output visualization I get this:

The bounding boxes should fit perfectly on the visible components as they almost do in the top-left-most boxes but the more we move away from the top-left corner the worse they get. Essentially it looks like the boxes would fit the image if it was resized. So I tried changing the test augmentation, so it resizes to smaller images ie:

def My_test_aug(cfg):
augs = [T.ResizeShortestEdge(1400, 2000)]
return augs

And that makes the visualization even worse:

The way I would expect this to work is that the dataloader reads images from the dataset and applies the augmentation from the mapper. The model then takes that image and makes predictions based on it. What I see here instead, is the predictions being unaffected by the augmentations given to the dataloader. Why?

[malje16]

Figured it out by looking at the sourcecode.
When you call the model on your input (Its used as a functor in the default trainer) It performs inference and then resizes the output, the predictions, back to the original image dimensions. Its done in a function called post-processing. you can specify when you call the model that you do not want it to postprocess.
So this behaviour is by design, although it seems strange to me to hide this transformation inside the model call. It makes it hard to find.