CODEWITHSUNDEEP
pythonnumpy
Roberto
Roberto

Reputation: 287

Elegantly iterate the planes of an image - if more than one

I come from the Matlab world and I'm relatively new to Python, so I think I might be approaching this from a completely wrong perspective.

Anyway, I often find myself to write code that needs to operate separately on the R,G,B planes of an image, but needs to be general enough that if the image is in greyscale it will still work. Now the non-so-clever way I started off is:

if im_in.ndim == 2:
  im_out = signal.convolve2d(im_in, filt, 'same')
else:
  im_out = np.empty_like(im_in)
  for kk in range(im_in.shape[2]):
    im_out[:,:,kk] = signal.convolve2d(im_in, filt, 'same')

Never mind the actual operation - I'm using signal.convolve2d just as an example here. And let's assume that ndim can only be 2 or 3 here, for simplicity.

Now where Matlab was quite clever is that I could just loop on the third dimension of the 3D array representing the image regardless of the number of planes.

The obvious alternative to the above is something like:

if im_in.ndim == 2:
  im_in.shape = (im_in.shape[0], im_in.shape[1], 1)

Then I can just loop in the third dimension (just as in the else case above), but still it seems to me this is a bit of a hack, and I'd have to reshape im_out before the end. Is there a proper, elegant way to deal with this case?

Upvotes: 2

Views: 76

Answers (3)

Sven Marnach
Sven Marnach

Reputation: 602315

You can use numpy.atleast_3d() to unconditionally get a view of your input array with the right number of dimensions, and then iterate over the planes. This has the advantage of saving the if statement and not modifying the input image.

Upvotes: 3

syntonym
syntonym

Reputation: 7384

You could automate your first solution (or what ever solution you coose):

from functools import wraps

def rgb_seperate(f):
    @wraps(f)
    def seperate(array, *args, **kwargs):
        im_out = np.empty_like(im_in)
        if im_in.ndim == 2:
             im_out = f(array, *args, **kwargs))
        else:
             for kk in range(im_in.shape[2]):
                 im_out[:,:,kk] = f(array[:,:,kk], *args, **kwargs)
        return(im_out)

Usage:

>>> rgb_seperate(signal.convolve2d)(im_in, filt, 'same')
[output]
>>>@rgb_seperate
>>>def operator(array, [someargs]):
>>>     ...
>>> operator(im_in, [someargs])
[output]

Of course this doesn't help with making the solution nicer internally, but might help with the usage of it. You wouldn't have to remember to type your solution everytime (however you will solve it in the end), but only to wrap it in rgb_seperate.

Upvotes: 0

Julien Spronck
Julien Spronck

Reputation: 15433

It is not much better than your own solution (although I think your solution is fine) but you can reshape your image without the if statement:

oldshape = im_in.shape
im_in = reshape(im_in, (oldshape[0], oldshape[1], im_in.size/(oldshape[0]*oldshape[1])))

That way, the shape of your image is a three-element tuple and you can iterate over im_in.shape[2]

Upvotes: 0

Related Questions