argmax maybe will work=

pages/3_Oxford_Vases_Disentanglement.py

@@ -90,7 +90,7 @@ smoothgrad_col_1, smoothgrad_col_2, smoothgrad_col_3, smoothgrad_col_4, smoothgr
 
 # ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
 with output_col_1:
-    separation_vector, number_important_features, imp_nodes, performance = get_separation_space(concept_ids, annotations, ann_df, latent_space=st.session_state.space_id, samples=200)
+    separation_vector, number_important_features, imp_nodes, performance = get_separation_space(concept_ids, annotations, ann_df, latent_space=st.session_state.space_id, samples=100)
     # st.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.3f}% confidence')
     st.write('Concept vector', separation_vector)
     header_col_1.write(f'Concept {st.session_state.concept_ids} - Space {st.session_state.space_id} - Number of relevant nodes: {number_important_features} - Val classification performance: {performance}')# - Nodes {",".join(list(imp_nodes))}')
@@ -143,7 +143,8 @@ else:
     original_image_vec = annotations['w_vectors'][st.session_state.image_id]
 
 img = generate_original_image(original_image_vec, model, latent_space=st.session_state.space_id)
-top_pred = ann_df.iloc[st.session_state.image_id].idxmax()
+print(ann_df.iloc[st.session_state.image_id, list(ann_df.column) - 'ID'])
+top_pred = ann_df.iloc[st.session_state.image_id, list(ann_df.columns) - 'ID'].idxmax()
 # input_image = original_image_dict['image']
 # input_label = original_image_dict['label']
 # input_id = original_image_dict['id']