Computer Vision News - February 2021

20 Computer Vision Tool # Random 3D kernel - HWDIO layout kernel = np.array([ [[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, -1, 0], [-1, 0, -1], [0, -1, 0]], [[0, 0, 0], [0, 1, 0], [0, 0, 0]]], dtype=jnp.float32)[:, :, :, np.newaxis, np.newaxis] # 3D data - NHWDC layout data = np.zeros((1, 30, 30, 30, 1), dtype=jnp.float32) x, y, z = np.mgrid[0:1:30j, 0:1:30j, 0:1:30j] data += (np.sin(2*x *jnp.pi) *np.cos(2*y *jnp.pi )*np.cos(2*z *jnp.pi ))[None,:,:,:,None] print("in shapes:", data.shape, kernel.shape) dn = lax.conv_dimension_numbers(data.shape, kernel.shape, ('NHWDC', 'HWDIO', 'NHWDC')) print(dn) out = lax.conv_general_dilated(data, # lhs = image tensor kernel, # rhs = conv kernel tensor (1,1,1), # window strides 'SAME', # padding mode (1,1,1), # lhs/image dilation (1,1,1), # rhs/kernel dilation dn) # dimension_numbers print("out shape: ", out.shape) # Make some simple 3d density plots: from mpl_toolkits.mplot3d import Axes3D def make_alpha(cmap): my_cmap = cmap(jnp.arange(cmap.N)) my_cmap[:,-1] = jnp.linspace(0, 1, cmap.N)**3 return mpl.colors.ListedColormap(my_cmap) my_cmap = make_alpha(plt.cm.viridis) fig = plt.figure() ax = fig.gca(projection='3d') ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=data.ravel(), cmap=my_cmap) ax.axis('off') ax.set_title('input') fig = plt.figure() ax = fig.gca(projection='3d') ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=out.ravel(), cmap=my_cmap) ax.axis('off') ax.set_title('3D conv output'); in shapes: (1, 30, 30, 30, 1) (3, 3, 3, 1, 1) ConvDimensionNumbers(lhs_spec=(0, 4, 1, 2, 3), rhs_spec=(4, 3, 0, 1, 2), out_spec=(0, 4, 1, 2, 3)) out shape: (1, 30, 30, 30, 1) [45]: Text(0.5, 0.92, '3D conv output')