bbbb

pages/3_Oxford_Vases_Disentanglement.py

@@ -144,9 +144,10 @@ else:
 
 print(ann_df.head())
 img = generate_original_image(original_image_vec, model, latent_space=st.session_state.space_id)
-print(ann_df.loc[st.session_state.image_id])
-print(ann_df.loc[st.session_state.image_id, list(ann_df.columns).remove('ID')])
-top_pred = ann_df.loc[st.session_state.image_id, list(ann_df.columns).remove('ID')].idxmax()
+cols = list(ann_df.columns)
+cols.remove('Unnamed: 0')
+cols.remove('ID')
+top_pred = ann_df.loc[st.session_state.image_id, cols].idxmax()
 # input_image = original_image_dict['image']
 # input_label = original_image_dict['label']
 # input_id = original_image_dict['id']