Upload app.py

app.py

@@ -30,9 +30,9 @@ MRISEQ_CAT = ["T1", "T2-FLAIR"]  # 0,1
 SEX_CAT = ["female", "male"]  # 0,1
 HEIGHT, WIDTH = 270, 270
 # chest
-SEX_CAT_CHEST = ["male", "female"]  # 0,1
-RACE_CAT = ["white", "asian", "black"]  # 0,1,2
-FIND_CAT = ["no disease", "pleural effusion"]
+SEX_CAT_CHEST = ["female", "male"]  # 0,1
+RACE_CAT = ["white", "black", "asian"]  # 0,1,2
+FIND_CAT = ["no disease", "consolidation", "opacity"] # 0,1,2
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 
@@ -42,6 +42,24 @@ class Hparams:
             setattr(self, k, v)
 
 
+# def get_paths(dataset_id):
+#     if "MNIST" in dataset_id:
+#         data_path = "./data/morphomnist"
+#         pgm_path = "./checkpoints/t_i_d/sup_pgm/checkpoint.pt"
+#         vae_path = "./checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt"
+#     elif "Brain" in dataset_id:
+#         data_path = "./data/ukbb_subset"
+#         pgm_path = "./checkpoints/m_b_v_s/sup_pgm/checkpoint.pt"
+#         vae_path = "./checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt"
+#     elif "Chest" in dataset_id:
+#         data_path = "./data/mimic_subset"
+#         pgm_path = "./checkpoints/a_r_s_f/pgm/checkpoint.pt"
+#         vae_path = [
+#             "./checkpoints/a_r_s_f/base_vae/checkpoint.pt",  # base vae
+#             "./checkpoints/a_r_s_f/cf_vae/6500_checkpoint.pt",  # cf trained DSCM
+#         ]
+#     return data_path, vae_path, pgm_path
+
 def get_paths(dataset_id):
     if "MNIST" in dataset_id:
         data_path = "./data/morphomnist"
@@ -53,14 +71,15 @@ def get_paths(dataset_id):
         vae_path = "./checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt"
     elif "Chest" in dataset_id:
         data_path = "./data/mimic_subset"
-        pgm_path = "./checkpoints/a_r_s_f/pgm/checkpoint.pt"
+        pgm_path = "./checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt"
         vae_path = [
-            "./checkpoints/a_r_s_f/base_vae/checkpoint.pt",  # base vae
-            "./checkpoints/a_r_s_f/cf_vae/6500_checkpoint.pt",  # cf trained DSCM
+            "./checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt",  # base vae
+            "./checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt",  # cf trained DSCM
         ]
     return data_path, vae_path, pgm_path
 
 
+
 def load_pgm(dataset_id, pgm_path):
     checkpoint = torch.load(pgm_path, map_location=DEVICE)
     args = Hparams()
@@ -202,7 +221,7 @@ def get_chest_obs(idx=None):
     idx, obs = get_obs_item(dataset_id, idx)
     x = get_fig_arr(postprocess(obs["x"].clone()))
     s = SEX_CAT_CHEST[int(obs["sex"].clone().squeeze().numpy())]
-    f = FIND_CAT[int(obs["finding"].clone().squeeze().numpy())]
+    f = FIND_CAT[obs["finding"].clone().squeeze().numpy().argmax(-1)]
     r = RACE_CAT[obs["race"].clone().squeeze().numpy().argmax(-1)]
     a = (obs["age"].clone().squeeze().numpy() + 1) * 50
     return (idx, x, r, s, f, float(np.round(a, 1)))
@@ -360,7 +379,10 @@ def infer_chest_cf(*args):
         if do_s:
             do_pa["sex"] = torch.tensor(SEX_CAT_CHEST.index(s)).view(1, 1)
         if do_f:
-            do_pa["finding"] = torch.tensor(FIND_CAT.index(f)).view(1, 1)
+            do_pa["finding"] = F.one_hot(
+                torch.tensor(FIND_CAT.index(f)), num_classes=3
+            ).view(1, 3)
+            # torch.tensor(FIND_CAT.index(f)).view(1, 1)
         if do_r:
             do_pa["race"] = F.one_hot(
                 torch.tensor(RACE_CAT.index(r)), num_classes=3
@@ -387,7 +409,7 @@ def infer_chest_cf(*args):
     rec_x = postprocess(rec_x)
     cf_r = RACE_CAT[cf_r.argmax(-1)]
     cf_s = SEX_CAT_CHEST[int(cf_s.item())]
-    cf_f = FIND_CAT[cf_f.argmax(-1)]
+    cf_f = FIND_CAT[cf_f.argmax(-1)] #FIND_CAT[int(cf_f.item())]
     cf_a = (cf_a.item() + 1) * 50
     # plots
     # plt.close('all')
@@ -426,7 +448,7 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                         label="Direct Causal Effect", interactive=False
                     , height=HEIGHT) #).style(height=HEIGHT)
             with gr.Row(): #.style(equal_height=True):
-                with gr.Column(scale=1.75):
+                with gr.Column(scale=1):#.75):
                     gr.Markdown(
                         "**Intervention**"
                         + 20 * " "
@@ -487,7 +509,7 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                         label="Direct Causal Effect", interactive=False
                     , height=HEIGHT) #).style(height=HEIGHT)
             with gr.Row():
-                with gr.Column(scale=2.55):
+                with gr.Column(scale=2):#.55):
                     gr.Markdown(
                         "**Intervention**"
                         + 20 * " "
@@ -571,7 +593,7 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                     , height=HEIGHT) #).style(height=HEIGHT)
 
             with gr.Row():
-                with gr.Column(scale=2.55):
+                with gr.Column(scale=2):#.55):
                     gr.Markdown(
                         "**Intervention**"
                         + 20 * " "