Fixed a bug in which gray images cause crash

[CallShaul]

Now functions support 2 dimensions gray images, not only 3 dimensions color images.

[nekhtiari]

Nice, LGTM

[chrieke]

Hey @CallShaul, thank you for the PR and sorry about the longer inactivity.

Wouldn't this also need to be applied on the second image?

[CallShaul]

@chrieke Hi, I'm not sure what you mean by "second image".

[nikooc2s]

@chrieke Hi, I'm not sure what you mean by "second image".

I think he meant the predicted one. Your changes are only for original image.

[CallShaul]

@nikooc2s @chrieke No, there is no need to do anything else, using the expand_dims this way, adds a third dimension with size 1, so the argument in the for loop org_img.shape[2] will not fail, yet the number of iterations will be 1 in the case of a gray image. (image size will change from [row, col] to [row, col, 1])