3D convolution in python

问题

I need to wite a code to perform a 3D convolution in python using numpy, with 3x3 kernels. I've done it right for 2D arrays like B&W images but when i try to extend it to 3D arrays like RGB is a mess. I need help to improve my method. Here is the 2D code:

def convolucion_3x3(arreglo, kernel):
  (dim_x, dim_y) = arreglo.shape
  (ker_x, ker_y) = kernel.shape

  matriz_convolucionada = np.zeros((dim_x, dim_y))

  for i in range(dim_x):
    for j in range(dim_y):
      resultado = 0
      for x in range(-1, 2):
        try:
          if i + x not in range(dim_x):
              raise ValueError()
          for y in range(-1, 2):
            try:
              if j + y not in range(dim_y):
                  raise ValueError()

              resultado += arreglo[i + x, j + y] * kernel[x + 1][y + 1]
              '''
              Para el kernel sumo un 1 a cada índice para que lo corra desde 0 hasta 2 y no de -1 a 1
              '''
            except ValueError:
                pass
        except ValueError:
            pass
      matriz_convolucionada[i][j] = resultado
  return matriz_convolucionada

The next one is my attempt to the RGB images:

def convolucion(arreglo, kernel): (dim_x, dim_y, dim_z) = arreglo.shape (ker_x, ker_y) = kernel.shape

matriz_convolucionada = np.zeros((dim_x, dim_y, dim_z))

for k in range(dim_z):
    for i in range(dim_x):
        for j in range(dim_y):
            resultado = 0
            for x in range(-1, 2):
                try:
                    if i + x not in range(dim_x):
                        raise ValueError()

                    for y in range(-1, 2):
                        try:
                            if j + y not in range(dim_y):
                                raise ValueError()

                            resultado += arreglo[i + x, j + y, k] * kernel[x + 1][y + 1]

                            '''
                            Para el kernel sumo un 1 a cada índice para que lo corra desde 0 hasta 2 y no de -1 a 1
                            '''

                        except ValueError:
                            pass

                except ValueError:
                    pass

            matriz_convolucionada[i][j][k] = resultado

return matriz_convolucionada

回答1:

While looping through would work, it can also be difficult to follow the nested loops. You might consider invoking the convolution theorem to perform the convolution easier. See here.

Using numpy's fft module, you can compute an n-dimensional discrete Fourier transform of the original stack of images and multiply it by the n-dimensional Fourier transform (documentation found here)of a kernel of the same size. Since your 2D kernel is a 3x3 array, it's a 3x3xz square 'pillar.' You can just pad this array with zeros to increase the dimensions accordingly.

Try this:

import numpy as np
import math

radius = 2
r2 = np.arange(-radius, radius+1)**2
sphere = r2[:, None, None] + r2[:, None] + r2
sphere -= np.max(sphere)
sphere = -sphere*2
array_len = 10*radius

array = np.zeros((array_len, array_len, array_len))
center = slice(array_len//2-radius,
               array_len//2+radius+1), slice(array_len//2-radius,
                                             array_len//2+radius+1),slice(array_len//2-radius,
                                                                          array_len//2+radius+1)
array[center] = sphere


k_len = 3
kernel_2D = np.ones((k_len,k_len))
kernel = np.zeros_like(array)

center_k = slice(array_len//2-math.ceil(k_len/2),
           array_len//2+k_len//2), slice(array_len//2-math.ceil(k_len/2),
                                         array_len//2+k_len//2)
for i in range(kernel.shape[2]):
    kernel[center_k+(i,)] = kernel_2D

def fft(array):
  fft = np.fft.ifftshift(np.fft.fftn(np.fft.fftshift(array)))
  return fft

def ifft(array):
  ifft = np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(array)))
  return ifft

def conv_3D(array, kernel):
  conv = np.abs(ifft(fft(array)*fft(kernel)))
  return conv

conv = conv_3D(array, kernel)

This convolves a sphere of radius 2 with a pillar of side length 3.

来源：https://stackoverflow.com/questions/47441952/3d-convolution-in-python