Update app.py

app.py

@@ -64,9 +64,12 @@ else:
     input = brightness.reshape(1, 28, 28, 1) / 255.
     st.write(cnn.apply({"params": params}, input).argmax(axis=1)[0])
 
-with st.expander("See first convolutional layer"):
-    gridify(params["Conv_0"]["kernel"], grid=(4,8), kernel_size=(3,3))
-
 def gridify(kernel, grid, kernel_size, scaling=5, padding=1):
     grid = np.pad(np.array(np.pad(np.repeat(np.repeat(kernel, repeats=scaling, axis=0), repeats=5, axis=1), ((padding,padding),(padding,padding),(0,0),(0,0)), 'constant', constant_values=(-1,)).reshape((kernel_size[0]*scaling+2*padding, kernel_size[1]*scaling+2*padding, grid[0], grid[1])).transpose(2,0,3,1).reshape(grid[0]*(kernel_size[0]*scaling+2*padding), grid[1]*(kernel_size[1]*scaling+2*padding))+1)*127., (1,), 'constant', constant_values=(0,))
     st.image(Image.fromarray(np.repeat(np.expand_dims(grid, axis=0), repeats=3, axis=0).astype(np.uint8).transpose(1,2,0), mode="RGB"))
+
+with st.expander("See first convolutional layer"):
+    gridify(params["Conv_0"]["kernel"], grid=(4,8), kernel_size=(3,3))
+    
+with st.expander("See second convolutional layer"):
+    gridify(params["Conv_1"]["kernel"], grid=(8,8), kernel_size=(3,3))