OpenCV cv2 Affine augment segmentation / heatmap broken when #classes > 3

[Erotemic]

I'm getting an issue when I use iaa.Affine to warp segmentation masks or heat maps.

error: OpenCV(4.1.0) /io/opencv/modules/imgproc/src/imgwarp.cpp:2597: error: (-215:Assertion failed) _src.channels() <= 4 || (interpolation != INTER_LANCZOS4 && interpolation != INTER_CUBIC) in function 'warpAffine'

This only happens when I force the cv2 backend and I have more than 3 classes in my heatmap / segmentation mask. For whatever reason opencv doesn't like to do cubic or lanczos resampling for many channel images (I don't see why it doesn't perhaps opencv needs a PR on this front).

Code to reproduce is as follows

def mwe():
    import imgaug
    import imgaug.augmenters as iaa
    import numpy as np
    augmenter = iaa.Affine(
        scale=(1, 1.2),
        order=1,
        cval=0,
        backend='cv2',
    )

    classes = ['class_{}'.format(i) for i in range(10)]
    sample_dims = (224, 224)

    # SOFT SEGMENTATION MASK
    # (C, H, W): independent per-class probabilities
    target_energy = np.random.rand(len(classes), *sample_dims).astype(np.float32)

    # HARD SEGMENTATION MASK
    # (H, W): build a mutually exclusive segmentation mask
    target_cidx = target_energy.argmax(axis=0).astype(np.int32)

    augdet = augmenter.to_deterministic()

    # Augment the hard mask
    cidx_segmap_oi = imgaug.SegmentationMapOnImage(target_cidx, target_cidx.shape, nb_classes=len(classes))
    cidx_segmap_oi = augdet.augment_segmentation_maps([cidx_segmap_oi])[0]
    target_cidx = cidx_segmap_oi.arr.argmax(axis=2)

    # Augment the soft mask
    class_target_hwc = target_energy.transpose(1, 2, 0)
    class_target_oi = imgaug.HeatmapsOnImage(class_target_hwc, class_target_hwc.shape)
    class_target_oi = augdet.augment_heatmaps([class_target_oi])[0]
    class_target_hwc = class_target_oi.get_arr()
    target_energy = class_target_hwc.transpose(2, 0, 1)

It took me awhile to understand why this was happening (because I had set order explicitly to 1, which is linear interpolation). I found out that this happens because in _augment_heatmaps on line 717: the interpolation order is hard coded to 3 (which is cubic).

I'm not sure why this isn't using the same order as the regular image augmentation. Perhaps there is a reason? I did a quick test and changed the hard coded 3 to a 1 and everything worked, but I haven't tried simply removing this line (which I think may be the correct fix).

Thoughts?

[aleju]

The interpolation order was hard-coded to 3, because the heatmap outputs are intended for ground truth and are therefore deteriorated as little as possible. So in the case of affine with e.g. nearest neighbour interpolation, the model would get an (augmented) input image with some information lost due to the interpolation method, but would still be trained to guess the correct heatmap as if no information had been lost. I recall thinking about the correct interpolation order there quite a bit and not coming to a clear decision whether it should always be 3 or always be the corresponding image array's interpolation order. You can come up with arguments for both cases, but ultimately it was fixed to 3.

It is however surprising that this fails for heatmaps with more than 4 channels. Affine was particularly adapted to automatically fall back to scikit-image in these cases -- which isn't a very efficient way of solving this as you could just split up the array into chunks of 4 channels and then feed them into OpenCV, but at least it shouldn't crash here. I hope I have the time to look into the issue this weekend.

[aleju]

The channel-issue should now be fixed in master, thanks for reporting the problem. Making order configurable for heatmaps is left to a future patch.

[aleju]

It should now be possible to change the interpolation order of heatmaps in Affine to "use the same as for images":

aug = Affine(...)
aug._order_heatmaps = None  # or an integer to always use a fixed value

Not tested though.

[Erotemic]

This is a reasonable workaround. Thanks for adding that.