added cv2 and show sample image

backend/disentangle_concepts.py

@@ -45,3 +45,13 @@ def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3
         img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
         images.append(PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB'))
     return images
+
+def generate_original_image(z, model):
+    device = torch.device('cpu')
+    G = model.to(device) # type: ignore
+    # Labels.
+    label = torch.zeros([1, G.c_dim], device=device)
+    z = torch.from_numpy(z.copy()).to(device)
+    img = G(z, label, truncation_psi=0.7, noise_mode='random')
+    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    return PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')

pages/1_Disentanglement.py

@@ -128,11 +128,18 @@ with input_col_3:
         epsilon_button = st.form_submit_button('Choose the defined epsilon')
 
 # ---------------------------- DISPLAY COL 2 ROW 1 ------------------------------
-# original_image_dict = load_image(st.session_state.image_id)
+
+model = torch.load('./data/model_files/pytorch_model.bin')
+original_image_vec = annotations['z_vectors'][st.session_state.image_id]
+img = generate_original_image(original_image_vec, model)
 # input_image = original_image_dict['image']
 # input_label = original_image_dict['label']
 # input_id = original_image_dict['id']
 
+with smoothgrad_col_3:
+    st.image(img)
+    header_col_1.write(f'Base image')
+
 
 
 # if pred_class_id == (input_id-1) and (epsilon_button or find_epsilon or setting_button):

requirements.txt

@@ -12,4 +12,5 @@ torchvision==0.11.2
 tqdm==4.64.1
 transformers==4.25.1
 scikit-learn
-altair==4.0
+altair==4.0
+opencv-python