Update app.py

app.py

@@ -65,6 +65,7 @@ else:
     st.write(cnn.apply({"params": params}, input).argmax(axis=1)[0])
 
 with st.expander("See first convolutional layer"):
+    first = np.array(params["Conv_0"]["kernel"].squeeze()[:, :, 0]+1)*127.
+    st.image(Image.fromarray(np.repeat(np.expand_dims(first, axis=0), repeats=3, axis=0).astype(np.uint8).transpose(1,2,0), mode="RGB"))
     grid = np.array(params["Conv_0"]["kernel"].reshape((3, 3, 4, 8)).transpose(2,0,3,1).reshape(12, 24)+1)*127.
-    rgb = np.repeat(np.expand_dims(grid, axis=0), repeats=3, axis=0).astype(np.uint8).transpose(1,2,0)
-    st.image(Image.fromarray(rgb, mode="RGB"))
+    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"))