Update app.py

app.py

@@ -23,12 +23,25 @@ if st.button("Train"):
     model.compile(optimizer="sgd",loss="sparse_categorical_crossentropy",metrics=["accuracy"])
     model.fit(x_train,y_train,epochs=10,batch_size=128,validation_split=0.2)
     fig, axs = plt.subplots(6, 1, figsize=(8, 6))
-    for i in range(6):
-        axs[i].imshow(
-            m.fit_transform(model.layers[0].weights[0][:, :, :, i][:, :, 0]),
-            cmap="gray"
-        )
-        axs[i].axis("off")  # remove axes for cleaner look
-    
-    st.pyplot(fig)
-
+    col1,col2,col3 = st.columns(3)
+    with col1:
+        st.write("Filters")
+        for i in range(6):
+            axs[i].imshow(
+                m.fit_transform(model.layers[0].weights[0][:, :, :, i][:, :, 0]),
+                cmap="gray"
+            )
+            axs[i].axis("off")
+    with col2:
+        st.write("Conv2D Layer")
+        sub = Model(inputs=model.inputs[0],outputs=model.layers[0].output)
+        fig,ax = plt.subplots(6,1,figsize=(8,6))
+        for i in range(6):
+            ax[i].imshow(sub.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
+            ax[i].axis("off")
+    with col3:
+        st.write("Average Pooling Layer")
+        max1 = Model(inputs=model.inputs[0],outputs=model.layers[1].output)
+        fig,ax1 = plt.subplots(6,1,figsize=(8,6))
+        for i in range(6):
+            ax1[i].imshow(max1.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")