fix warning message persistence and filter miss

app.py

@@ -105,28 +105,42 @@ def main():
 
     # Filter index selector
     if st.session_state.layer_name:
+        warn_ph = st.empty()
+        layer_ph = st.empty()
+
         filter_select = st.selectbox("Visualize -", VIS_OPTION.keys())
 
         if VIS_OPTION[filter_select] == 0:
             loss, img = visualize_filter(st.session_state.feat_extract, 0)
             st.image(img)
         else:
-            st.warning(":exclamation: Calculating the gradients can take a while..")
+            layer = st.session_state.model.get_layer(name=st.session_state.layer_name)
+            num_filters = layer.get_output_at(0).get_shape().as_list()[-1]
+
+            warn_ph.warning(
+                ":exclamation: Calculating the gradients can take a while.."
+            )
+            if num_filters < 64:
+                layer_ph.info(
+                    f"{st.session_state.layer_name} has only {num_filters} filters, visualizing only those filters.."
+                )
+
             prog_bar = st.progress(0)
             fig, axis = plt.subplots(nrows=8, ncols=8, figsize=(14, 14))
-            for filter_index, ax in enumerate(axis.ravel()):
-                prog_bar.progress((filter_index + 1) / 64)
-                if filter_index < 65:
-                    loss, img = visualize_filter(
-                        st.session_state.feat_extract, filter_index
-                    )
-                    ax.imshow(img)
-                    ax.set_title(filter_index + 1)
-                    ax.set_axis_off()
-                else:
+            for filter_index, ax in enumerate(axis.ravel()[: min(num_filters, 64)]):
+                prog_bar.progress((filter_index + 1) / min(num_filters, 64))
+                loss, img = visualize_filter(
+                    st.session_state.feat_extract, filter_index
+                )
+                ax.imshow(img)
+                ax.set_title(filter_index + 1)
+                ax.set_axis_off()
+            else:
+                for ax in axis.ravel()[num_filters:]:
                     ax.set_axis_off()
 
             st.write(fig)
+            warn_ph.empty()
 
 
 if __name__ == "__main__":