Update app.py

app.py

@@ -60,12 +60,29 @@ if uploaded_img is not None:
     fm = func_model.predict(input_img)
     fm = fm[0]
 
-    fig1 = plt.figure(figsize=(12, 4))
-    for i in range(6):
-        plt.subplot(2, 3, i + 1)
-        plt.imshow(fm[:, :, i], cmap='gray')
-        plt.title(f"Feature Map {i+1}")
-        plt.axis('off')
+    if layer_ind == 0:
+        fig1 = plt.figure(figsize=(12, 4))
+        for i in range(6):
+            plt.subplot(2, 3, i + 1)
+            plt.imshow(fm[:, :, i], cmap='gray')
+            plt.title(f"Feature Map {i+1}")
+            plt.axis('off')
+    
+    elif layer_ind == 1:
+        fig1 = plt.figure(figsize=(12, 4))
+        for i in range(16):
+            plt.subplot(2, 8, i + 1)
+            plt.imshow(fm[:, :, i], cmap='gray')
+            plt.title(f"Feature Map {i+1}")
+            plt.axis('off')
+    
+    elif layer_ind == 2:
+        fig1 = plt.figure(figsize=(12, 4))
+        for i in range(120):
+            plt.subplot(30, 4, i + 1)
+            plt.imshow(fm[:, :, i], cmap='gray')
+            plt.title(f"Feature Map {i+1}")
+            plt.axis('off')
 
     plt.tight_layout()
     st.pyplot(fig1)