2D convolution using only numpy

np_convolve2d.py

def np_convolve2d(a, v, mode='valid', keep_dims=False):  # mode='valid'
    """Convolve an image with filter_n using valid mode
    In valid mode the output consists only of those elements that do not rely on the zero-padding.
    # TODO: full and same modes
    
    Parameters
    ----------
    a : (N,) array_like
        First one-dimensional input array.
    v : (M,) array_like
        Second one-dimensional input array.
    mode : {'full', 'valid', 'same'}, optional
        'full':
          By default, mode is 'full'.  This returns the convolution
          at each point of overlap, with an output shape of (N+M-1,). At
          the end-points of the convolution, the signals do not overlap
          completely, and boundary effects may be seen.
        'same':
          Mode 'same' returns output of length ``max(M, N)``.  Boundary
          effects are still visible.
        'valid':
          Mode 'valid' returns output of length
          ``max(M, N) - min(M, N) + 1``.  The convolution product is only given
          for points where the signals overlap completely.  Values outside
          the signal boundary have no effect.
    keep_dims : str
        If True, keep the size of the original image and replace only the affected pixels by the filtered results.
        
    Returns
    -------
    out : ndarray
        Discrete, linear convolution of `a` and `v`.
    """
    sub_shape = conv_filter.shape + tuple(np.subtract(image.shape, conv_filter.shape) + 1)
    strider = np.lib.stride_tricks.as_strided
    sub_matrices = strider(image, shape=sub_shape, strides=image.strides*2)
    
    c = np.einsum('ij,ijkl->kl', conv_filter, sub_matrices)
    
    if keep_dims:
        h = np.floor(np.array(v.shape)/2).astype(np.int)
        c_keep_dims = a.copy()
        c_keep_dims[h[0]:-h[0],h[1]:-h[1]] = c
        return c_keep_dims
    return c