Receptive Field for dilated convolutions

hdc.py

"""
Motivation: Understanding Convolution for Semantic Segmentation (https://arxiv.org/pdf/1702.08502.pdf) 
          : https://stats.stackexchange.com/questions/265462/whats-the-receptive-field-of-a-stack-of-dilated-convolutions
"""

import pdb
import traceback
import numpy as np
import matplotlib.pyplot as plt

cmap     = 'viridis' # ['gray', 'viridis']

def get_kernel(ksize, dilrate):
    """
    For square kernels only
    """

    # Step 1 - Get empty kernel
    kernel = np.zeros((1 + (ksize-1)*dilrate, 1 + (ksize-1)*dilrate), dtype=np.uint8)
    
    # Step 2 - Fill the kernel depending on dilrate
    allidxs = np.argwhere(kernel == 0)
    idxs_kernel = allidxs[(allidxs[:,0] % dilrate==0) & (allidxs[:,1] % dilrate == 0)]
    kernel[idxs_kernel[:,0], idxs_kernel[:,1]] = 1

    return kernel

def get_image(img_size):

    # Step 1 - Get empty image
    image = np.zeros((img_size,img_size), dtype=np.uint8)

    # Step 2 - Set middle pixel as 1 since we will calculate its receptive field 
    image[img_size//2,img_size//2] = 1

    return image

def do_conv(image, ksize, dilrates):

    f,axarr = plt.subplots(2,len(dilrates))
    image_copy = np.array(image, copy=True)
    for dilid, dilrate in enumerate(dilrates):
        
        print (' - ksize: {}, dilrate={}'.format(ksize, dilrate))
        try:
            # Step 1 - Get a kernel
            kernel = get_kernel(ksize, dilrate)
            image_copy2 = np.array(image_copy, copy=True)
            idx_sets = np.argwhere(image_copy > 0)
            for idx_set in idx_sets:
                image_copy2[idx_set[0]-kernel.shape[0]//2:idx_set[0]+kernel.shape[0]//2+1,idx_set[1]-kernel.shape[0]//2:idx_set[1]+kernel.shape[0]//2+1] += kernel
            axarr[0, dilid].imshow(image_copy2, cmap=cmap)
            axarr[0, dilid].grid(True); _ = axarr[0, dilid].set_xticks(np.arange(image_copy2.shape[0]) + 0.5); _ = axarr[0, dilid].set_xticklabels([]); _ = axarr[0, dilid].set_yticks(np.arange(image_copy2.shape[1]) + 0.5); _ = axarr[0, dilid].set_yticklabels([]);        
            axarr[0, dilid].set_title('k={}, dil={}'.format(ksize, dilrate))

            # Step 2 - Do conv for middle pixel
            idx_sets = np.argwhere(image > 0)
            # print (' - len(idx_sets): ', len(idx_sets))
            for idx_set in idx_sets:
                image[idx_set[0]-kernel.shape[0]//2:idx_set[0]+kernel.shape[0]//2+1,idx_set[1]-kernel.shape[0]//2:idx_set[1]+kernel.shape[0]//2+1] += kernel
            
            # Step 3 - Plot every step 
            if 0:
                rgba = plt.get_cmap(cmap)(image)
                # rgba[image == 0, :] = [1,1,1,1]
                axarr[1, dilid].imshow(rgba)
            else:
                axarr[1, dilid].imshow(image, cmap=cmap)
            axarr[1, dilid].grid(True); _ = axarr[1, dilid].set_xticks(np.arange(image.shape[0]) + 0.5); _ = axarr[1, dilid].set_xticklabels([]); _ = axarr[1, dilid].set_yticks(np.arange(image.shape[1]) + 0.5); _ = axarr[1, dilid].set_yticklabels([]);
            axarr[1, dilid].set_title('Effective Receptive Field'.format(ksize, dilrate))
            print ('')
        except:
            traceback.print_exc()
            pdb.set_trace()

    plt.show()

if __name__ == "__main__":

    if 0:
        img_size = 25
        ksize    = 3
        dilrates = [1,2,3]
    elif 0:
        img_size = 25
        ksize    = 3
        dilrates = [1,2,5]
    elif 0:
        img_size = 25
        ksize    = 3
        dilrates = [2,2,2]
    
    elif 1:
        img_size = 25
        ksize    = 3
        dilrates = [2,3,5]
    
    elif 0:
        img_size = 40
        ksize    = 3
        dilrates = [3,3,6,6]
    
    elif 0:
        img_size = 40
        ksize    = 3
        dilrates = [1,2,5,9]
    
    elif 0:
        img_size = 40
        ksize    = 3
        dilrates = [1,2,3,5]
    
    elif 0:
        img_size = 25
        ksize    = 3
        dilrates = [1,1]
    
    # Step 1 - Get image
    image = get_image(img_size=img_size)
    do_conv(image, ksize, dilrates)