Spaces:
Sleeping
Sleeping
Commit
·
8c4fe8b
0
Parent(s):
Duplicate from fabio-deep/counterfactuals
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- .gitattributes +34 -0
- .gitignore +2 -0
- README.md +14 -0
- app.py +556 -0
- app_utils.py +373 -0
- checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt +3 -0
- checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt +3 -0
- checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt +3 -0
- checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint_current.pt +3 -0
- checkpoints/m_b_v_s/sup_pgm/checkpoint.pt +3 -0
- checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt +3 -0
- checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt +3 -0
- checkpoints/t_i_d/sup_pgm/checkpoint.pt +3 -0
- data/mimic_subset/0.jpg +0 -0
- data/mimic_subset/1.jpg +0 -0
- data/mimic_subset/10.jpg +0 -0
- data/mimic_subset/11.jpg +0 -0
- data/mimic_subset/12.jpg +0 -0
- data/mimic_subset/13.jpg +0 -0
- data/mimic_subset/14.jpg +0 -0
- data/mimic_subset/15.jpg +0 -0
- data/mimic_subset/16.jpg +0 -0
- data/mimic_subset/17.jpg +0 -0
- data/mimic_subset/18.jpg +0 -0
- data/mimic_subset/19.jpg +0 -0
- data/mimic_subset/2.jpg +0 -0
- data/mimic_subset/20.jpg +0 -0
- data/mimic_subset/21.jpg +0 -0
- data/mimic_subset/22.jpg +0 -0
- data/mimic_subset/23.jpg +0 -0
- data/mimic_subset/24.jpg +0 -0
- data/mimic_subset/25.jpg +0 -0
- data/mimic_subset/26.jpg +0 -0
- data/mimic_subset/27.jpg +0 -0
- data/mimic_subset/28.jpg +0 -0
- data/mimic_subset/29.jpg +0 -0
- data/mimic_subset/3.jpg +0 -0
- data/mimic_subset/30.jpg +0 -0
- data/mimic_subset/31.jpg +0 -0
- data/mimic_subset/32.jpg +0 -0
- data/mimic_subset/33.jpg +0 -0
- data/mimic_subset/34.jpg +0 -0
- data/mimic_subset/35.jpg +0 -0
- data/mimic_subset/36.jpg +0 -0
- data/mimic_subset/37.jpg +0 -0
- data/mimic_subset/38.jpg +0 -0
- data/mimic_subset/39.jpg +0 -0
- data/mimic_subset/4.jpg +0 -0
- data/mimic_subset/40.jpg +0 -0
- data/mimic_subset/41.jpg +0 -0
.gitattributes
ADDED
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
28 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
29 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
30 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
31 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
32 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
33 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
34 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
__pycache__
|
2 |
+
*.pyc
|
README.md
ADDED
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
title: Counterfactuals
|
3 |
+
emoji: 🌖
|
4 |
+
colorFrom: purple
|
5 |
+
colorTo: green
|
6 |
+
sdk: gradio
|
7 |
+
sdk_version: 3.27.0
|
8 |
+
app_file: app.py
|
9 |
+
pinned: false
|
10 |
+
license: mit
|
11 |
+
duplicated_from: fabio-deep/counterfactuals
|
12 |
+
---
|
13 |
+
|
14 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
ADDED
@@ -0,0 +1,556 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import numpy as np
|
3 |
+
import gradio as gr
|
4 |
+
import matplotlib.pylab as plt
|
5 |
+
import torch.nn.functional as F
|
6 |
+
|
7 |
+
from vae import HVAE
|
8 |
+
from datasets import morphomnist, ukbb, mimic, get_attr_max_min
|
9 |
+
from pgm.flow_pgm import MorphoMNISTPGM, FlowPGM, ChestPGM
|
10 |
+
from app_utils import (
|
11 |
+
mnist_graph, brain_graph, chest_graph, vae_preprocess, normalize, \
|
12 |
+
preprocess_brain, get_fig_arr, postprocess, MidpointNormalize
|
13 |
+
)
|
14 |
+
|
15 |
+
DATA, MODELS = {}, {}
|
16 |
+
for k in ['Morpho-MNIST', 'Brain MRI', 'Chest X-ray']:
|
17 |
+
DATA[k], MODELS[k] = {}, {}
|
18 |
+
|
19 |
+
# mnist
|
20 |
+
DIGITS = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]
|
21 |
+
# brain
|
22 |
+
MRISEQ_CAT = ['T1', 'T2-FLAIR'] # 0,1
|
23 |
+
SEX_CAT = ['female', 'male'] # 0,1
|
24 |
+
HEIGHT, WIDTH = 270, 270
|
25 |
+
# chest
|
26 |
+
SEX_CAT_CHEST = ['male', 'female'] # 0,1
|
27 |
+
RACE_CAT = ['white', 'asian', 'black'] # 0,1,2
|
28 |
+
FIND_CAT = ['no disease', 'pleural effusion']
|
29 |
+
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
30 |
+
|
31 |
+
|
32 |
+
class Hparams:
|
33 |
+
def update(self, dict):
|
34 |
+
for k, v in dict.items():
|
35 |
+
setattr(self, k, v)
|
36 |
+
|
37 |
+
|
38 |
+
def get_paths(dataset_id):
|
39 |
+
if 'MNIST' in dataset_id:
|
40 |
+
data_path = './data/morphomnist'
|
41 |
+
pgm_path = './checkpoints/t_i_d/sup_pgm/checkpoint.pt'
|
42 |
+
vae_path = './checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt'
|
43 |
+
elif 'Brain' in dataset_id:
|
44 |
+
data_path = './data/ukbb_subset'
|
45 |
+
pgm_path = './checkpoints/m_b_v_s/sup_pgm/checkpoint.pt'
|
46 |
+
vae_path = './checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt'
|
47 |
+
elif 'Chest' in dataset_id:
|
48 |
+
data_path = './data/mimic_subset'
|
49 |
+
pgm_path = './checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt'
|
50 |
+
vae_path = [
|
51 |
+
'./checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt', # base vae
|
52 |
+
'./checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt' # cf trained DSCM
|
53 |
+
]
|
54 |
+
return data_path, vae_path, pgm_path
|
55 |
+
|
56 |
+
|
57 |
+
def load_pgm(dataset_id, pgm_path):
|
58 |
+
checkpoint = torch.load(pgm_path, map_location=DEVICE)
|
59 |
+
args = Hparams()
|
60 |
+
args.update(checkpoint['hparams'])
|
61 |
+
args.device = DEVICE
|
62 |
+
if 'MNIST' in dataset_id:
|
63 |
+
pgm = MorphoMNISTPGM(args).to(args.device)
|
64 |
+
elif 'Brain' in dataset_id:
|
65 |
+
pgm = FlowPGM(args).to(args.device)
|
66 |
+
elif 'Chest' in dataset_id:
|
67 |
+
pgm = ChestPGM(args).to(args.device)
|
68 |
+
pgm.load_state_dict(checkpoint['ema_model_state_dict'])
|
69 |
+
MODELS[dataset_id]['pgm'] = pgm
|
70 |
+
MODELS[dataset_id]['pgm_args'] = args
|
71 |
+
|
72 |
+
|
73 |
+
def load_vae(dataset_id, vae_path):
|
74 |
+
if 'Chest' in dataset_id:
|
75 |
+
vae_path, dscm_path = vae_path[0], vae_path[1]
|
76 |
+
checkpoint = torch.load(vae_path, map_location=DEVICE)
|
77 |
+
args = Hparams()
|
78 |
+
args.update(checkpoint['hparams'])
|
79 |
+
# backwards compatibility hack
|
80 |
+
if not hasattr(args, 'vae'):
|
81 |
+
args.vae = 'hierarchical'
|
82 |
+
if not hasattr(args, 'cond_prior'):
|
83 |
+
args.cond_prior = False
|
84 |
+
if hasattr(args, 'free_bits'):
|
85 |
+
args.kl_free_bits = args.free_bits
|
86 |
+
args.device = DEVICE
|
87 |
+
vae = HVAE(args).to(args.device)
|
88 |
+
|
89 |
+
if 'Chest' in dataset_id:
|
90 |
+
dscm_ckpt = torch.load(dscm_path, map_location=DEVICE)
|
91 |
+
vae.load_state_dict({k[4:]: v for k, v in dscm_ckpt['ema_model_state_dict'].items() if 'vae.' in k})
|
92 |
+
else:
|
93 |
+
vae.load_state_dict(checkpoint['ema_model_state_dict'])
|
94 |
+
MODELS[dataset_id]['vae'] = vae
|
95 |
+
MODELS[dataset_id]['vae_args'] = args
|
96 |
+
|
97 |
+
|
98 |
+
def get_dataloader(dataset_id, data_path):
|
99 |
+
MODELS[dataset_id]['pgm_args'].data_dir = data_path
|
100 |
+
args = MODELS[dataset_id]['pgm_args']
|
101 |
+
if 'MNIST' in dataset_id:
|
102 |
+
datasets = morphomnist(args)
|
103 |
+
elif 'Brain' in dataset_id:
|
104 |
+
datasets = ukbb(args)
|
105 |
+
elif 'Chest' in dataset_id:
|
106 |
+
datasets = mimic(args)
|
107 |
+
DATA[dataset_id]['test'] = torch.utils.data.DataLoader(
|
108 |
+
datasets['test'], shuffle=False, batch_size=args.bs, num_workers=4)
|
109 |
+
|
110 |
+
|
111 |
+
def load_dataset(dataset_id):
|
112 |
+
data_path, _, pgm_path = get_paths(dataset_id)
|
113 |
+
checkpoint = torch.load(pgm_path, map_location=DEVICE)
|
114 |
+
args = Hparams()
|
115 |
+
args.update(checkpoint['hparams'])
|
116 |
+
args.device = DEVICE
|
117 |
+
MODELS[dataset_id]['pgm_args'] = args
|
118 |
+
get_dataloader(dataset_id, data_path)
|
119 |
+
|
120 |
+
|
121 |
+
def load_model(dataset_id):
|
122 |
+
_, vae_path, pgm_path = get_paths(dataset_id)
|
123 |
+
load_pgm(dataset_id, pgm_path)
|
124 |
+
load_vae(dataset_id, vae_path)
|
125 |
+
|
126 |
+
|
127 |
+
@torch.no_grad()
|
128 |
+
def counterfactual_inference(dataset_id, obs, do_pa):
|
129 |
+
pa = {k: v.clone() for k, v in obs.items() if k != 'x'}
|
130 |
+
cf_pa = MODELS[dataset_id]['pgm'].counterfactual(obs=pa, intervention=do_pa, num_particles=1)
|
131 |
+
args, vae = MODELS[dataset_id]['vae_args'], MODELS[dataset_id]['vae']
|
132 |
+
_pa = vae_preprocess(args, {k: v.clone() for k, v in pa.items()})
|
133 |
+
_cf_pa = vae_preprocess(args , {k: v.clone() for k, v in cf_pa.items()})
|
134 |
+
z_t = 0.1 if 'mnist' in args.hps else 1.0
|
135 |
+
z = vae.abduct(x=obs['x'], parents=_pa, t=z_t)
|
136 |
+
if vae.cond_prior:
|
137 |
+
z = [z[j]['z'] for j in range(len(z))]
|
138 |
+
px_loc, px_scale = vae.forward_latents(latents=z, parents=_pa)
|
139 |
+
cf_loc, cf_scale = vae.forward_latents(latents=z, parents=_cf_pa)
|
140 |
+
u = (obs['x'] - px_loc) / px_scale.clamp(min=1e-12)
|
141 |
+
u_t = 0.1 if 'mnist' in args.hps else 1.0 # cf sampling temp
|
142 |
+
cf_scale = cf_scale * u_t
|
143 |
+
cf_x = torch.clamp(cf_loc + cf_scale * u, min=-1, max=1)
|
144 |
+
return {'cf_x': cf_x, 'rec_x': px_loc, 'cf_pa': cf_pa}
|
145 |
+
|
146 |
+
|
147 |
+
def get_obs_item(dataset_id, idx=None):
|
148 |
+
if idx is None:
|
149 |
+
n_test = len(DATA[dataset_id]['test'].dataset)
|
150 |
+
idx = torch.randperm(n_test)[0]
|
151 |
+
idx = int(idx)
|
152 |
+
return idx, DATA[dataset_id]['test'].dataset.__getitem__(idx)
|
153 |
+
|
154 |
+
|
155 |
+
def get_mnist_obs(idx=None):
|
156 |
+
dataset_id = 'Morpho-MNIST'
|
157 |
+
if not DATA[dataset_id]:
|
158 |
+
load_dataset(dataset_id)
|
159 |
+
idx, obs = get_obs_item(dataset_id, idx)
|
160 |
+
x = get_fig_arr(obs['x'].clone().squeeze().numpy())
|
161 |
+
t = (obs['thickness'].clone() + 1) / 2 * (6.255515 - 0.87598526) + 0.87598526
|
162 |
+
i = (obs['intensity'].clone() + 1) / 2 * (254.90317 - 66.601204) + 66.601204
|
163 |
+
y = DIGITS[obs['digit'].clone().argmax(-1)]
|
164 |
+
return (idx, x, float(np.round(t, 2)), float(np.round(i, 2)), y)
|
165 |
+
|
166 |
+
|
167 |
+
def get_brain_obs(idx=None):
|
168 |
+
dataset_id = 'Brain MRI'
|
169 |
+
if not DATA[dataset_id]:
|
170 |
+
load_dataset(dataset_id)
|
171 |
+
idx, obs = get_obs_item(dataset_id, idx)
|
172 |
+
x = get_fig_arr(obs['x'].clone().squeeze().numpy())
|
173 |
+
m = MRISEQ_CAT[int(obs['mri_seq'].clone().item())]
|
174 |
+
s = SEX_CAT[int(obs['sex'].clone().item())]
|
175 |
+
a = obs['age'].clone().item()
|
176 |
+
b = obs['brain_volume'].clone().item() / 1000 # in ml
|
177 |
+
v = obs['ventricle_volume'].clone().item() / 1000 # in ml
|
178 |
+
return (idx, x, m, s, a, float(np.round(b, 2)), float(np.round(v, 2)))
|
179 |
+
|
180 |
+
|
181 |
+
def get_chest_obs(idx=None):
|
182 |
+
dataset_id = 'Chest X-ray'
|
183 |
+
if not DATA[dataset_id]:
|
184 |
+
load_dataset(dataset_id)
|
185 |
+
idx, obs = get_obs_item(dataset_id, idx)
|
186 |
+
x = get_fig_arr(postprocess(obs['x'].clone()))
|
187 |
+
s = SEX_CAT_CHEST[int(obs['sex'].clone().squeeze().numpy())]
|
188 |
+
f = FIND_CAT[int(obs['finding'].clone().squeeze().numpy())]
|
189 |
+
r = RACE_CAT[obs['race'].clone().squeeze().numpy().argmax(-1)]
|
190 |
+
a = (obs['age'].clone().squeeze().numpy()+1)*50
|
191 |
+
return (idx, x, r, s, f, float(np.round(a, 1)))
|
192 |
+
|
193 |
+
|
194 |
+
def infer_mnist_cf(*args):
|
195 |
+
dataset_id = 'Morpho-MNIST'
|
196 |
+
idx, _, t, i, y, do_t, do_i, do_y = args
|
197 |
+
n_particles = 32
|
198 |
+
# preprocess
|
199 |
+
obs = DATA[dataset_id]['test'].dataset.__getitem__(int(idx))
|
200 |
+
obs['x'] = (obs['x'] - 127.5) / 127.5
|
201 |
+
for k, v in obs.items():
|
202 |
+
obs[k] = v.view(1, 1) if len(v.shape) < 1 else v.unsqueeze(0)
|
203 |
+
obs[k] = obs[k].to(MODELS[dataset_id]['vae_args'].device).float()
|
204 |
+
if n_particles > 1:
|
205 |
+
ndims = (1,)*3 if k == 'x' else (1,)
|
206 |
+
obs[k] = obs[k].repeat(n_particles, *ndims)
|
207 |
+
# intervention(s)
|
208 |
+
do_pa = {}
|
209 |
+
if do_t:
|
210 |
+
do_pa['thickness'] = torch.tensor(normalize(t, x_max=6.255515, x_min=0.87598526)).view(1, 1)
|
211 |
+
if do_i:
|
212 |
+
do_pa['intensity'] = torch.tensor(normalize(i, x_max=254.90317, x_min=66.601204)).view(1, 1)
|
213 |
+
if do_y:
|
214 |
+
do_pa['digit'] = F.one_hot(torch.tensor(DIGITS.index(y)), num_classes=10).view(1, 10)
|
215 |
+
|
216 |
+
for k, v in do_pa.items():
|
217 |
+
do_pa[k] = v.to(MODELS[dataset_id]['vae_args'].device).float().repeat(n_particles, 1)
|
218 |
+
# infer counterfactual
|
219 |
+
out = counterfactual_inference(dataset_id, obs, do_pa)
|
220 |
+
# avg cf particles
|
221 |
+
cf_x = out['cf_x'].mean(0)
|
222 |
+
cf_x_std = out['cf_x'].std(0)
|
223 |
+
rec_x = out['rec_x'].mean(0)
|
224 |
+
cf_t = out['cf_pa']['thickness'].mean(0)
|
225 |
+
cf_i = out['cf_pa']['intensity'].mean(0)
|
226 |
+
cf_y = out['cf_pa']['digit'].mean(0)
|
227 |
+
# post process
|
228 |
+
cf_x = postprocess(cf_x)
|
229 |
+
cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
|
230 |
+
rec_x = postprocess(rec_x)
|
231 |
+
cf_t = np.round((cf_t.item() + 1) / 2 * (6.255515 - 0.87598526) + 0.87598526, 2)
|
232 |
+
cf_i = np.round((cf_i.item() + 1) / 2 * (254.90317 - 66.601204) + 66.601204, 2)
|
233 |
+
cf_y = DIGITS[cf_y.argmax(-1)]
|
234 |
+
# plots
|
235 |
+
# plt.close('all')
|
236 |
+
effect = cf_x - rec_x
|
237 |
+
effect = get_fig_arr(effect, cmap='RdBu_r',
|
238 |
+
norm=MidpointNormalize(vmin=-255, midpoint=0, vmax=255))
|
239 |
+
cf_x = get_fig_arr(cf_x)
|
240 |
+
cf_x_std = get_fig_arr(cf_x_std, cmap='jet')
|
241 |
+
return (cf_x, cf_x_std, effect, cf_t, cf_i, cf_y)
|
242 |
+
|
243 |
+
|
244 |
+
def infer_brain_cf(*args):
|
245 |
+
dataset_id = 'Brain MRI'
|
246 |
+
idx, _, m, s, a, b, v = args[:7]
|
247 |
+
do_m, do_s, do_a, do_b, do_v = args[7:]
|
248 |
+
n_particles = 16
|
249 |
+
# preprocessing
|
250 |
+
obs = DATA[dataset_id]['test'].dataset.__getitem__(int(idx))
|
251 |
+
obs.pop('pa')
|
252 |
+
obs = preprocess_brain(MODELS[dataset_id]['vae_args'], obs)
|
253 |
+
for k, _v in obs.items():
|
254 |
+
if n_particles > 1:
|
255 |
+
ndims = (1,)*3 if k == 'x' else (1,)
|
256 |
+
obs[k] = _v.repeat(n_particles, *ndims)
|
257 |
+
# interventions(s)
|
258 |
+
do_pa = {}
|
259 |
+
if do_m:
|
260 |
+
do_pa['mri_seq'] = torch.tensor(MRISEQ_CAT.index(m)).view(1, 1)
|
261 |
+
if do_s:
|
262 |
+
do_pa['sex'] = torch.tensor(SEX_CAT.index(s)).view(1, 1)
|
263 |
+
if do_a:
|
264 |
+
do_pa['age'] = torch.tensor(a).view(1, 1)
|
265 |
+
if do_b:
|
266 |
+
do_pa['brain_volume'] = torch.tensor(b * 1000).view(1, 1)
|
267 |
+
if do_v:
|
268 |
+
do_pa['ventricle_volume'] = torch.tensor(v * 1000).view(1, 1)
|
269 |
+
# normalize continuous attributes
|
270 |
+
for k in ['age', 'brain_volume', 'ventricle_volume']:
|
271 |
+
if k in do_pa.keys():
|
272 |
+
k_max, k_min = get_attr_max_min(k)
|
273 |
+
do_pa[k] = (do_pa[k] - k_min) / (k_max - k_min) # [0,1]
|
274 |
+
do_pa[k] = 2 * do_pa[k] - 1 # [-1,1]
|
275 |
+
|
276 |
+
for k, _v in do_pa.items():
|
277 |
+
do_pa[k] = _v.to(MODELS[dataset_id]['vae_args'].device).float().repeat(n_particles, 1)
|
278 |
+
# infer counterfactual
|
279 |
+
out = counterfactual_inference(dataset_id, obs, do_pa)
|
280 |
+
# avg cf particles
|
281 |
+
cf_x = out['cf_x'].mean(0)
|
282 |
+
cf_x_std = out['cf_x'].std(0)
|
283 |
+
rec_x = out['rec_x'].mean(0)
|
284 |
+
cf_m = out['cf_pa']['mri_seq'].mean(0)
|
285 |
+
cf_s = out['cf_pa']['sex'].mean(0)
|
286 |
+
# post process
|
287 |
+
cf_x = postprocess(cf_x)
|
288 |
+
cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
|
289 |
+
rec_x = postprocess(rec_x)
|
290 |
+
cf_m = MRISEQ_CAT[int(cf_m.item())]
|
291 |
+
cf_s = SEX_CAT[int(cf_s.item())]
|
292 |
+
cf_ = {}
|
293 |
+
for k in ['age', 'brain_volume', 'ventricle_volume']: # unnormalize
|
294 |
+
k_max, k_min = get_attr_max_min(k)
|
295 |
+
cf_[k] = (out['cf_pa'][k].mean(0).item() + 1) / 2 * (k_max - k_min) + k_min
|
296 |
+
# plots
|
297 |
+
# plt.close('all')
|
298 |
+
effect = cf_x - rec_x
|
299 |
+
effect = get_fig_arr(effect, cmap='RdBu_r',
|
300 |
+
norm=MidpointNormalize(vmin=effect.min(), midpoint=0, vmax=effect.max()))
|
301 |
+
cf_x = get_fig_arr(cf_x)
|
302 |
+
cf_x_std = get_fig_arr(cf_x_std, cmap='jet')
|
303 |
+
return (cf_x, cf_x_std, effect, cf_m, cf_s, np.round(cf_['age'], 1), np.round(cf_['brain_volume'] / 1000, 2), np.round(cf_['ventricle_volume'] / 1000, 2))
|
304 |
+
|
305 |
+
|
306 |
+
def infer_chest_cf(*args):
|
307 |
+
dataset_id = 'Chest X-ray'
|
308 |
+
idx, _, r, s, f, a = args[:6]
|
309 |
+
do_r, do_s, do_f, do_a = args[6:]
|
310 |
+
n_particles = 16
|
311 |
+
# preprocessing
|
312 |
+
obs = DATA[dataset_id]['test'].dataset.__getitem__(int(idx))
|
313 |
+
for k, v in obs.items():
|
314 |
+
obs[k] = v.to(MODELS[dataset_id]['vae_args'].device).float()
|
315 |
+
if n_particles > 1:
|
316 |
+
ndims = (1,)*3 if k == 'x' else (1,)
|
317 |
+
obs[k] = obs[k].repeat(n_particles, *ndims)
|
318 |
+
# intervention(s)
|
319 |
+
do_pa = {}
|
320 |
+
with torch.no_grad():
|
321 |
+
if do_s:
|
322 |
+
do_pa['sex'] = torch.tensor(SEX_CAT_CHEST.index(s)).view(1, 1)
|
323 |
+
if do_f:
|
324 |
+
do_pa['finding'] = torch.tensor(FIND_CAT.index(f)).view(1, 1)
|
325 |
+
if do_r:
|
326 |
+
do_pa['race'] = F.one_hot(torch.tensor(RACE_CAT.index(r)), num_classes=3).view(1, 3)
|
327 |
+
if do_a:
|
328 |
+
do_pa['age'] = torch.tensor(a/100*2-1).view(1,1)
|
329 |
+
for k, v in do_pa.items():
|
330 |
+
do_pa[k] = v.to(MODELS[dataset_id]['vae_args'].device).float().repeat(n_particles, 1)
|
331 |
+
# infer counterfactual
|
332 |
+
out = counterfactual_inference(dataset_id, obs, do_pa)
|
333 |
+
# avg cf particles
|
334 |
+
cf_x = out['cf_x'].mean(0)
|
335 |
+
cf_x_std = out['cf_x'].std(0)
|
336 |
+
rec_x = out['rec_x'].mean(0)
|
337 |
+
cf_r = out['cf_pa']['race'].mean(0)
|
338 |
+
cf_s = out['cf_pa']['sex'].mean(0)
|
339 |
+
cf_f = out['cf_pa']['finding'].mean(0)
|
340 |
+
cf_a = out['cf_pa']['age'].mean(0)
|
341 |
+
# post process
|
342 |
+
cf_x = postprocess(cf_x)
|
343 |
+
cf_x_std = cf_x_std.squeeze().detach().cpu().numpy()
|
344 |
+
rec_x = postprocess(rec_x)
|
345 |
+
cf_r = RACE_CAT[cf_r.argmax(-1)]
|
346 |
+
cf_s = SEX_CAT_CHEST[int(cf_s.item())]
|
347 |
+
cf_f = FIND_CAT[int(cf_f.item())]
|
348 |
+
cf_a = (cf_a.item() + 1) * 50
|
349 |
+
# plots
|
350 |
+
# plt.close('all')
|
351 |
+
effect = cf_x - rec_x
|
352 |
+
effect = get_fig_arr(effect, cmap='RdBu_r',
|
353 |
+
norm=MidpointNormalize(vmin=effect.min(), midpoint=0, vmax=effect.max()))
|
354 |
+
cf_x = get_fig_arr(cf_x)
|
355 |
+
cf_x_std = get_fig_arr(cf_x_std, cmap='jet')
|
356 |
+
return (cf_x, cf_x_std, effect, cf_r, cf_s, cf_f, np.round(cf_a, 1))
|
357 |
+
|
358 |
+
|
359 |
+
with gr.Blocks(theme=gr.themes.Default()) as demo:
|
360 |
+
with gr.Tabs():
|
361 |
+
with gr.TabItem("Morpho-MNIST") as mnist_tab:
|
362 |
+
mnist_id = gr.Textbox(value=mnist_tab.label, visible=False)
|
363 |
+
|
364 |
+
with gr.Row().style(equal_height=True):
|
365 |
+
idx = gr.Number(value=0, visible=False)
|
366 |
+
with gr.Column(scale=1, min_width=200):
|
367 |
+
x = gr.Image(label='Observation', interactive=False).style(height=HEIGHT)
|
368 |
+
with gr.Column(scale=1, min_width=200):
|
369 |
+
cf_x = gr.Image(label='Counterfactual', interactive=False).style(height=HEIGHT)
|
370 |
+
with gr.Column(scale=1, min_width=200):
|
371 |
+
cf_x_std = gr.Image(label='Counterfactual Uncertainty', interactive=False).style(height=HEIGHT)
|
372 |
+
with gr.Column(scale=1, min_width=200):
|
373 |
+
effect = gr.Image(label='Direct Causal Effect', interactive=False).style(height=HEIGHT)
|
374 |
+
with gr.Row().style(equal_height=True):
|
375 |
+
with gr.Column(scale=1.75):
|
376 |
+
gr.Markdown("#### Intervention")
|
377 |
+
with gr.Column():
|
378 |
+
do_y = gr.Checkbox(label="do(digit)", value=False)
|
379 |
+
y = gr.Radio(DIGITS, label="", interactive=False)
|
380 |
+
with gr.Row():
|
381 |
+
with gr.Column(min_width=100):
|
382 |
+
do_t = gr.Checkbox(label="do(thickness)", value=False)
|
383 |
+
t = gr.Slider(label="\u00A0", minimum=0.9, maximum=5.5, step=0.01, interactive=False)
|
384 |
+
with gr.Column(min_width=100):
|
385 |
+
do_i = gr.Checkbox(label="do(intensity)", value=False)
|
386 |
+
i = gr.Slider(label="\u00A0", minimum=50, maximum=255, step=0.01, interactive=False)
|
387 |
+
with gr.Row():
|
388 |
+
new = gr.Button("New Observation")
|
389 |
+
reset = gr.Button("Reset", variant="stop")
|
390 |
+
submit = gr.Button("Submit", variant="primary")
|
391 |
+
with gr.Column(scale=1):
|
392 |
+
gr.Markdown("### ")
|
393 |
+
causal_graph = gr.Image(label='Causal Graph', interactive=False).style(height=300)
|
394 |
+
|
395 |
+
with gr.TabItem("Brain MRI") as brain_tab:
|
396 |
+
brain_id = gr.Textbox(value=brain_tab.label, visible=False)
|
397 |
+
|
398 |
+
with gr.Row().style(equal_height=True):
|
399 |
+
idx_brain = gr.Number(value=0, visible=False)
|
400 |
+
with gr.Column(scale=1, min_width=200):
|
401 |
+
x_brain = gr.Image(label='Observation', interactive=False).style(height=HEIGHT)
|
402 |
+
with gr.Column(scale=1, min_width=200):
|
403 |
+
cf_x_brain = gr.Image(label='Counterfactual', interactive=False).style(height=HEIGHT)
|
404 |
+
with gr.Column(scale=1, min_width=200):
|
405 |
+
cf_x_std_brain = gr.Image(label='Counterfactual Uncertainty', interactive=False).style(height=HEIGHT)
|
406 |
+
with gr.Column(scale=1, min_width=200):
|
407 |
+
effect_brain = gr.Image(label='Direct Causal Effect', interactive=False).style(height=HEIGHT)
|
408 |
+
with gr.Row():
|
409 |
+
with gr.Column(scale=2.55):
|
410 |
+
gr.Markdown("#### Intervention")
|
411 |
+
with gr.Row():
|
412 |
+
with gr.Column(min_width=200):
|
413 |
+
do_m = gr.Checkbox(label="do(MRI sequence)", value=False)
|
414 |
+
m = gr.Radio(["T1", "T2-FLAIR"], label="", interactive=False)
|
415 |
+
with gr.Column(min_width=200):
|
416 |
+
do_s = gr.Checkbox(label="do(sex)", value=False)
|
417 |
+
s = gr.Radio(["female", "male"], label="", interactive=False)
|
418 |
+
with gr.Row():
|
419 |
+
with gr.Column(min_width=100):
|
420 |
+
do_a = gr.Checkbox(label="do(age)", value=False)
|
421 |
+
a = gr.Slider(label="\u00A0", value=50, minimum=44, maximum=73, step=1, interactive=False)
|
422 |
+
with gr.Column(min_width=100):
|
423 |
+
do_b = gr.Checkbox(label="do(brain volume)", value=False)
|
424 |
+
b = gr.Slider(label="\u00A0", value=1000, minimum=850, maximum=1550, step=20, interactive=False)
|
425 |
+
with gr.Column(min_width=100):
|
426 |
+
do_v = gr.Checkbox(label="do(ventricle volume)", value=False)
|
427 |
+
v = gr.Slider(label="\u00A0", value=40, minimum=10, maximum=125, step=2, interactive=False)
|
428 |
+
with gr.Row():
|
429 |
+
new_brain = gr.Button("New Observation")
|
430 |
+
reset_brain = gr.Button("Reset", variant='stop')
|
431 |
+
submit_brain = gr.Button("Submit", variant='primary')
|
432 |
+
with gr.Column(scale=1):
|
433 |
+
# gr.Markdown("### ")
|
434 |
+
causal_graph_brain = gr.Image(label='Causal Graph', interactive=False).style(height=340)
|
435 |
+
|
436 |
+
with gr.TabItem("Chest X-ray") as chest_tab:
|
437 |
+
chest_id = gr.Textbox(value=chest_tab.label, visible=False)
|
438 |
+
|
439 |
+
with gr.Row().style(equal_height=True):
|
440 |
+
idx_chest = gr.Number(value=0, visible=False)
|
441 |
+
with gr.Column(scale=1, min_width=200):
|
442 |
+
x_chest = gr.Image(label='Observation', interactive=False).style(height=HEIGHT)
|
443 |
+
with gr.Column(scale=1, min_width=200):
|
444 |
+
cf_x_chest = gr.Image(label='Counterfactual', interactive=False).style(height=HEIGHT)
|
445 |
+
with gr.Column(scale=1, min_width=200):
|
446 |
+
cf_x_std_chest = gr.Image(label='Counterfactual Uncertainty', interactive=False).style(height=HEIGHT)
|
447 |
+
with gr.Column(scale=1, min_width=200):
|
448 |
+
effect_chest = gr.Image(label='Direct Causal Effect', interactive=False).style(height=HEIGHT)
|
449 |
+
|
450 |
+
with gr.Row():
|
451 |
+
with gr.Column(scale=2.55):
|
452 |
+
gr.Markdown("#### Intervention")
|
453 |
+
with gr.Row().style(equal_height=True):
|
454 |
+
with gr.Column(min_width=200):
|
455 |
+
do_f_chest = gr.Checkbox(label="do(disease)", value=False)
|
456 |
+
f_chest = gr.Radio(FIND_CAT, label="", interactive=False)
|
457 |
+
with gr.Column(min_width=200):
|
458 |
+
do_s_chest = gr.Checkbox(label="do(sex)", value=False)
|
459 |
+
s_chest = gr.Radio(SEX_CAT_CHEST, label="", interactive=False)
|
460 |
+
|
461 |
+
with gr.Row():
|
462 |
+
with gr.Column(min_width=200):
|
463 |
+
do_r_chest = gr.Checkbox(label="do(race)", value=False)
|
464 |
+
r_chest = gr.Radio(RACE_CAT, label="", interactive=False)
|
465 |
+
with gr.Column(min_width=200):
|
466 |
+
do_a_chest = gr.Checkbox(label="do(age)", value=False)
|
467 |
+
a_chest = gr.Slider(label="\u00A0", minimum=18, maximum=98, step=1)
|
468 |
+
|
469 |
+
with gr.Row():
|
470 |
+
new_chest = gr.Button("New Observation")
|
471 |
+
reset_chest = gr.Button("Reset", variant="stop")
|
472 |
+
submit_chest = gr.Button("Submit", variant="primary")
|
473 |
+
with gr.Column(scale=1):
|
474 |
+
# gr.Markdown("### ")
|
475 |
+
causal_graph_chest = gr.Image(label='Causal Graph', interactive=False).style(height=345)
|
476 |
+
|
477 |
+
# morphomnist
|
478 |
+
do = [do_t, do_i, do_y]
|
479 |
+
obs = [idx, x, t, i, y]
|
480 |
+
cf_out = [cf_x, cf_x_std, effect]
|
481 |
+
|
482 |
+
# brain
|
483 |
+
do_brain = [do_m, do_s, do_a, do_b, do_v] # intervention checkboxes
|
484 |
+
obs_brain = [idx_brain, x_brain, m, s, a, b, v] # observed image/attributes
|
485 |
+
cf_out_brain = [cf_x_brain, cf_x_std_brain, effect_brain] # counterfactual outputs
|
486 |
+
|
487 |
+
# chest
|
488 |
+
do_chest = [do_r_chest, do_s_chest, do_f_chest, do_a_chest]
|
489 |
+
obs_chest = [idx_chest, x_chest, r_chest, s_chest, f_chest, a_chest]
|
490 |
+
cf_out_chest = [cf_x_chest, cf_x_std_chest, effect_chest]
|
491 |
+
|
492 |
+
# on start: load new observations & causal graph
|
493 |
+
demo.load(fn=get_mnist_obs, inputs=None, outputs=obs)
|
494 |
+
demo.load(fn=mnist_graph, inputs=do, outputs=causal_graph)
|
495 |
+
demo.load(fn=load_model, inputs=mnist_id, outputs=None)
|
496 |
+
demo.load(fn=get_brain_obs, inputs=None, outputs=obs_brain)
|
497 |
+
demo.load(fn=get_chest_obs, inputs=None, outputs=obs_chest)
|
498 |
+
|
499 |
+
demo.load(fn=brain_graph, inputs=do_brain, outputs=causal_graph_brain)
|
500 |
+
demo.load(fn=chest_graph, inputs=do_chest, outputs=causal_graph_chest)
|
501 |
+
|
502 |
+
# on tab select: load models
|
503 |
+
brain_tab.select(fn=load_model, inputs=brain_id, outputs=None)
|
504 |
+
chest_tab.select(fn=load_model, inputs=chest_id, outputs=None)
|
505 |
+
|
506 |
+
# "new" button: load new observations
|
507 |
+
new.click(fn=get_mnist_obs, inputs=None, outputs=obs)
|
508 |
+
new_chest.click(fn=get_chest_obs, inputs=None, outputs=obs_chest)
|
509 |
+
new_brain.click(fn=get_brain_obs, inputs=None, outputs=obs_brain)
|
510 |
+
|
511 |
+
# "new" button: reset causal graphs
|
512 |
+
new.click(fn=mnist_graph, inputs=do, outputs=causal_graph)
|
513 |
+
new_brain.click(fn=brain_graph, inputs=do_brain, outputs=causal_graph_brain)
|
514 |
+
new_chest.click(fn=chest_graph, inputs=do_chest, outputs=causal_graph_chest)
|
515 |
+
|
516 |
+
# "new" button: reset cf output panels
|
517 |
+
for _k, _v in zip([new, new_brain, new_chest], [cf_out, cf_out_brain, cf_out_chest]):
|
518 |
+
_k.click(fn=lambda: (gr.update(value=None),)*3, inputs=None, outputs=_v)
|
519 |
+
|
520 |
+
# "reset" button: reload current observations
|
521 |
+
reset.click(fn=get_mnist_obs, inputs=idx, outputs=obs)
|
522 |
+
reset_brain.click(fn=get_brain_obs, inputs=idx_brain, outputs=obs_brain)
|
523 |
+
reset_chest.click(fn=get_chest_obs, inputs=idx_chest, outputs=obs_chest)
|
524 |
+
|
525 |
+
# "reset" button: deselect intervention checkboxes
|
526 |
+
reset.click(fn=lambda: (gr.update(value=False),)*len(do), inputs=None, outputs=do)
|
527 |
+
reset_brain.click(fn=lambda: (gr.update(value=False),)*len(do_brain), inputs=None, outputs=do_brain)
|
528 |
+
reset_chest.click(fn=lambda: (gr.update(value=False),)*len(do_chest), inputs=None, outputs=do_chest)
|
529 |
+
|
530 |
+
# "reset" button: reset cf output panels
|
531 |
+
for _k, _v in zip([reset, reset_brain, reset_chest], [cf_out, cf_out_brain, cf_out_chest]):
|
532 |
+
_k.click(fn=lambda: plt.close('all'), inputs=None, outputs=None)
|
533 |
+
_k.click(fn=lambda: (gr.update(value=None),)*3, inputs=None, outputs=_v)
|
534 |
+
|
535 |
+
# enable mnist interventions when checkbox is selected & update graph
|
536 |
+
for _k, _v in zip(do, [t, i, y]):
|
537 |
+
_k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
|
538 |
+
_k.change(mnist_graph, inputs=do, outputs=causal_graph)
|
539 |
+
|
540 |
+
# enable brain interventions when checkbox is selected & update graph
|
541 |
+
for _k, _v in zip(do_brain, [m, s, a, b, v]):
|
542 |
+
_k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
|
543 |
+
_k.change(brain_graph, inputs=do_brain, outputs=causal_graph_brain)
|
544 |
+
|
545 |
+
# enable chest interventions when checkbox is selected & update graph
|
546 |
+
for _k, _v in zip(do_chest, [r_chest, s_chest, f_chest, a_chest]):
|
547 |
+
_k.change(fn=lambda x: gr.update(interactive=x), inputs=_k, outputs=_v)
|
548 |
+
_k.change(chest_graph, inputs=do_chest, outputs=causal_graph_chest)
|
549 |
+
|
550 |
+
# "submit" button: infer countefactuals
|
551 |
+
submit.click(fn=infer_mnist_cf, inputs=obs + do, outputs=cf_out + [t, i, y])
|
552 |
+
submit_brain.click(fn=infer_brain_cf, inputs=obs_brain + do_brain, outputs=cf_out_brain + [m, s, a, b, v])
|
553 |
+
submit_chest.click(fn=infer_chest_cf, inputs=obs_chest + do_chest, outputs=cf_out_chest + [r_chest, s_chest, f_chest, a_chest])
|
554 |
+
|
555 |
+
if __name__ == "__main__":
|
556 |
+
demo.launch()
|
app_utils.py
ADDED
@@ -0,0 +1,373 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import numpy as np
|
3 |
+
import networkx as nx
|
4 |
+
import matplotlib.pyplot as plt
|
5 |
+
|
6 |
+
from matplotlib import rc, patches, colors
|
7 |
+
rc('font', **{'family': 'serif', 'serif': ['Roman']})
|
8 |
+
rc('text', usetex=True)
|
9 |
+
rc('image', interpolation='none')
|
10 |
+
rc('text.latex', preamble=r'\usepackage{amsmath} \usepackage{amssymb}')
|
11 |
+
|
12 |
+
from datasets import get_attr_max_min
|
13 |
+
|
14 |
+
from PIL import Image
|
15 |
+
HAMMER = np.array(Image.open('./hammer.png').resize((35, 35))) / 255
|
16 |
+
|
17 |
+
|
18 |
+
class MidpointNormalize(colors.Normalize):
|
19 |
+
def __init__(self, vmin=None, vmax=None, midpoint=None, clip=False):
|
20 |
+
self.midpoint = midpoint
|
21 |
+
colors.Normalize.__init__(self, vmin, vmax, clip)
|
22 |
+
|
23 |
+
def __call__(self, value, clip=None):
|
24 |
+
v_ext = np.max( [ np.abs(self.vmin), np.abs(self.vmax) ] )
|
25 |
+
x, y = [-v_ext, self.midpoint, v_ext], [0, 0.5, 1]
|
26 |
+
return np.ma.masked_array(np.interp(value, x, y))
|
27 |
+
|
28 |
+
|
29 |
+
def postprocess(x):
|
30 |
+
return ((x + 1.0) * 127.5).squeeze().detach().cpu().numpy()
|
31 |
+
|
32 |
+
|
33 |
+
def mnist_graph(*args):
|
34 |
+
x, t, i, y = r'$\mathbf{x}$', r'$t$', r'$i$', r'$y$'
|
35 |
+
ut, ui, uy = r'$\mathbf{U}_t$', r'$\mathbf{U}_i$', r'$\mathbf{U}_y$'
|
36 |
+
zx, ex = r'$\mathbf{z}_{1:L}$', r'$\boldsymbol{\epsilon}$'
|
37 |
+
|
38 |
+
G = nx.DiGraph()
|
39 |
+
G.add_edge(t, x)
|
40 |
+
G.add_edge(i, x)
|
41 |
+
G.add_edge(y, x)
|
42 |
+
G.add_edge(t, i)
|
43 |
+
G.add_edge(ut, t)
|
44 |
+
G.add_edge(ui, i)
|
45 |
+
G.add_edge(uy, y)
|
46 |
+
G.add_edge(zx, x)
|
47 |
+
G.add_edge(ex, x)
|
48 |
+
|
49 |
+
pos = {
|
50 |
+
y: (0, 0), uy: (-1, 0),
|
51 |
+
t: (0, 0.5), ut: (0, 1),
|
52 |
+
x: (1, 0), zx: (2, 0.375), ex: (2, 0),
|
53 |
+
i: (1, 0.5), ui: (1, 1),
|
54 |
+
}
|
55 |
+
|
56 |
+
node_c = {}
|
57 |
+
for node in G:
|
58 |
+
node_c[node] = 'lightgrey' if node in [x, t, i, y] else 'white'
|
59 |
+
node_line_c = {k: 'black' for k, _ in node_c.items()}
|
60 |
+
edge_c = {e: 'black' for e in G.edges}
|
61 |
+
|
62 |
+
if args[0]: # do_t
|
63 |
+
edge_c[(ut, t)] = 'lightgrey'
|
64 |
+
# G.remove_edge(ut, t)
|
65 |
+
node_line_c[t] = 'red'
|
66 |
+
if args[1]: # do_i
|
67 |
+
edge_c[(ui, i)] = 'lightgrey'
|
68 |
+
edge_c[(t, i)] = 'lightgrey'
|
69 |
+
# G.remove_edges_from([(ui, i), (t, i)])
|
70 |
+
node_line_c[i] = 'red'
|
71 |
+
if args[2]: # do_y
|
72 |
+
edge_c[(uy, y)] = 'lightgrey'
|
73 |
+
# G.remove_edge(uy, y)
|
74 |
+
node_line_c[y] = 'red'
|
75 |
+
|
76 |
+
fs = 30
|
77 |
+
options = {
|
78 |
+
"font_size": fs,
|
79 |
+
"node_size": 3000,
|
80 |
+
"node_color": list(node_c.values()),
|
81 |
+
"edgecolors": list(node_line_c.values()),
|
82 |
+
"edge_color": list(edge_c.values()),
|
83 |
+
"linewidths": 2,
|
84 |
+
"width": 2,
|
85 |
+
}
|
86 |
+
plt.close('all')
|
87 |
+
fig, ax = plt.subplots(1, 1, figsize=(6,4.1))#, constrained_layout=True)
|
88 |
+
# fig.patch.set_visible(False)
|
89 |
+
ax.margins(x=0.06, y=0.15, tight=False)
|
90 |
+
ax.axis("off")
|
91 |
+
nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle='-|>', ax=ax)
|
92 |
+
# need to reuse x, y limits so that the graphs plot the same way before and after removing edges
|
93 |
+
x_lim = (-1.348, 2.348)
|
94 |
+
y_lim = (-0.215, 1.215)
|
95 |
+
ax.set_xlim(x_lim)
|
96 |
+
ax.set_ylim(y_lim)
|
97 |
+
rect = patches.FancyBboxPatch((1.75, -0.16), 0.5, 0.7, boxstyle="round, pad=0.05, rounding_size=0", linewidth=2, edgecolor='black', facecolor='none', linestyle='-')
|
98 |
+
ax.add_patch(rect)
|
99 |
+
ax.text(1.85, 0.65, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
|
100 |
+
|
101 |
+
if args[0]: # do_t
|
102 |
+
fig.figimage(HAMMER, 0.26*fig.bbox.xmax, 0.525*fig.bbox.ymax, zorder=10)
|
103 |
+
if args[1]: # do_i
|
104 |
+
fig.figimage(HAMMER, 0.5175*fig.bbox.xmax, 0.525*fig.bbox.ymax, zorder=11)
|
105 |
+
if args[2]: # do_y
|
106 |
+
fig.figimage(HAMMER, 0.26*fig.bbox.xmax, 0.2*fig.bbox.ymax, zorder=12)
|
107 |
+
|
108 |
+
fig.tight_layout()
|
109 |
+
fig.canvas.draw()
|
110 |
+
return np.array(fig.canvas.renderer.buffer_rgba())
|
111 |
+
|
112 |
+
|
113 |
+
def brain_graph(*args):
|
114 |
+
x, m, s, a, b, v = r'$\mathbf{x}$', r'$m$', r'$s$', r'$a$', r'$b$', r'$v$'
|
115 |
+
um, us, ua, ub, uv = r'$\mathbf{U}_m$', r'$\mathbf{U}_s$', r'$\mathbf{U}_a$', r'$\mathbf{U}_b$', r'$\mathbf{U}_v$'
|
116 |
+
zx, ex = r'$\mathbf{z}_{1:L}$', r'$\boldsymbol{\epsilon}$'
|
117 |
+
|
118 |
+
G = nx.DiGraph()
|
119 |
+
G.add_edge(m, x)
|
120 |
+
G.add_edge(s, x)
|
121 |
+
G.add_edge(b, x)
|
122 |
+
G.add_edge(v, x)
|
123 |
+
G.add_edge(zx, x)
|
124 |
+
G.add_edge(ex, x)
|
125 |
+
G.add_edge(a, b)
|
126 |
+
G.add_edge(a, v)
|
127 |
+
G.add_edge(s, b)
|
128 |
+
G.add_edge(um, m)
|
129 |
+
G.add_edge(us, s)
|
130 |
+
G.add_edge(ua, a)
|
131 |
+
G.add_edge(ub, b)
|
132 |
+
G.add_edge(uv, v)
|
133 |
+
|
134 |
+
pos = {
|
135 |
+
x: (0, 0), zx: (-0.25, -1), ex: (0.25, -1),
|
136 |
+
a: (0, 1), ua: (0, 2),
|
137 |
+
s: (1, 0), us: (1, -1),
|
138 |
+
b: (1, 1), ub: (1, 2),
|
139 |
+
m: (-1, 0), um: (-1, -1),
|
140 |
+
v: (-1, 1), uv: (-1, 2)
|
141 |
+
}
|
142 |
+
|
143 |
+
node_c = {}
|
144 |
+
for node in G:
|
145 |
+
node_c[node] = 'lightgrey' if node in [x, m, s, a, b, v] else 'white'
|
146 |
+
node_line_c = {k: 'black' for k, _ in node_c.items()}
|
147 |
+
edge_c = {e: 'black' for e in G.edges}
|
148 |
+
|
149 |
+
if args[0]: # do_m
|
150 |
+
# G.remove_edge(um, m)
|
151 |
+
edge_c[(um, m)] = 'lightgrey'
|
152 |
+
node_line_c[m] = 'red'
|
153 |
+
if args[1]: # do_s
|
154 |
+
# G.remove_edge(us, s)
|
155 |
+
edge_c[(us, s)] = 'lightgrey'
|
156 |
+
node_line_c[s] = 'red'
|
157 |
+
if args[2]: # do_a
|
158 |
+
# G.remove_edge(ua, a)
|
159 |
+
edge_c[(ua, a)] = 'lightgrey'
|
160 |
+
node_line_c[a] = 'red'
|
161 |
+
if args[3]: # do_b
|
162 |
+
# G.remove_edges_from([(ub, b), (s, b), (a, b)])
|
163 |
+
edge_c[(ub, b)] = 'lightgrey'
|
164 |
+
edge_c[(s, b)] = 'lightgrey'
|
165 |
+
edge_c[(a, b)] = 'lightgrey'
|
166 |
+
node_line_c[b] = 'red'
|
167 |
+
if args[4]: # do_v
|
168 |
+
# G.remove_edges_from([(uv, v), (a, v), (b, v)])
|
169 |
+
edge_c[(uv, v)] = 'lightgrey'
|
170 |
+
edge_c[(a, v)] = 'lightgrey'
|
171 |
+
edge_c[(b, v)] = 'lightgrey'
|
172 |
+
node_line_c[v] = 'red'
|
173 |
+
|
174 |
+
fs = 30
|
175 |
+
options = {
|
176 |
+
"font_size": fs,
|
177 |
+
"node_size": 3000,
|
178 |
+
"node_color": list(node_c.values()),
|
179 |
+
"edgecolors": list(node_line_c.values()),
|
180 |
+
"edge_color": list(edge_c.values()),
|
181 |
+
"linewidths": 2,
|
182 |
+
"width": 2,
|
183 |
+
}
|
184 |
+
|
185 |
+
plt.close('all')
|
186 |
+
fig, ax = plt.subplots(1, 1, figsize=(5,5))#, constrained_layout=True)
|
187 |
+
# fig.patch.set_visible(False)
|
188 |
+
ax.margins(x=0.1, y=0.08, tight=False)
|
189 |
+
ax.axis("off")
|
190 |
+
nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle='-|>', ax=ax)
|
191 |
+
# need to reuse x, y limits so that the graphs plot the same way before and after removing edges
|
192 |
+
x_lim = (-1.32, 1.32)
|
193 |
+
y_lim = (-1.414, 2.414)
|
194 |
+
ax.set_xlim(x_lim)
|
195 |
+
ax.set_ylim(y_lim)
|
196 |
+
rect = patches.FancyBboxPatch((-0.5, -1.325), 1, 0.65, boxstyle="round, pad=0.05, rounding_size=0", linewidth=2, edgecolor='black', facecolor='none', linestyle='-')
|
197 |
+
ax.add_patch(rect)
|
198 |
+
# ax.text(1.85, 0.65, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
|
199 |
+
|
200 |
+
if args[0]: # do_m
|
201 |
+
fig.figimage(HAMMER, 0.0075*fig.bbox.xmax, 0.395*fig.bbox.ymax, zorder=10)
|
202 |
+
if args[1]: # do_s
|
203 |
+
fig.figimage(HAMMER, 0.72*fig.bbox.xmax, 0.395*fig.bbox.ymax, zorder=11)
|
204 |
+
if args[2]: # do_a
|
205 |
+
fig.figimage(HAMMER, 0.363*fig.bbox.xmax, 0.64*fig.bbox.ymax, zorder=12)
|
206 |
+
if args[3]: # do_b
|
207 |
+
fig.figimage(HAMMER, 0.72*fig.bbox.xmax, 0.64*fig.bbox.ymax, zorder=13)
|
208 |
+
if args[4]: # do_v
|
209 |
+
fig.figimage(HAMMER, 0.0075*fig.bbox.xmax, 0.64*fig.bbox.ymax, zorder=14)
|
210 |
+
else: # b -> v
|
211 |
+
a3 = patches.FancyArrowPatch((.86, 1.21), (-.86, 1.21), connectionstyle="arc3,rad=.3", linewidth=2, arrowstyle='simple, head_width=10, head_length=10', color='k')
|
212 |
+
ax.add_patch(a3)
|
213 |
+
# print(ax.get_xlim())
|
214 |
+
# print(ax.get_ylim())
|
215 |
+
fig.tight_layout()
|
216 |
+
fig.canvas.draw()
|
217 |
+
return np.array(fig.canvas.renderer.buffer_rgba())
|
218 |
+
|
219 |
+
|
220 |
+
|
221 |
+
def chest_graph(*args):
|
222 |
+
x, a, d, r, s= r'$\mathbf{x}$', r'$a$', r'$d$', r'$r$', r'$s$'
|
223 |
+
ua, ud, ur, us = r'$\mathbf{U}_a$', r'$\mathbf{U}_d$', r'$\mathbf{U}_r$', r'$\mathbf{U}_s$'
|
224 |
+
zx, ex = r'$\mathbf{z}_{1:L}$', r'$\boldsymbol{\epsilon}$'
|
225 |
+
|
226 |
+
G = nx.DiGraph()
|
227 |
+
G.add_edge(ua, a)
|
228 |
+
G.add_edge(ud, d)
|
229 |
+
G.add_edge(ur, r)
|
230 |
+
G.add_edge(us, s)
|
231 |
+
G.add_edge(a, d)
|
232 |
+
G.add_edge(d, x)
|
233 |
+
G.add_edge(r, x)
|
234 |
+
G.add_edge(s, x)
|
235 |
+
G.add_edge(ex, x)
|
236 |
+
G.add_edge(zx, x)
|
237 |
+
G.add_edge(a, x)
|
238 |
+
|
239 |
+
pos = {
|
240 |
+
x: (0, 0),
|
241 |
+
a: (-1, 1),
|
242 |
+
d: (0, 1),
|
243 |
+
r: (1, 1),
|
244 |
+
s: (1, 0),
|
245 |
+
ua: (-1, 2),
|
246 |
+
ud: (0, 2),
|
247 |
+
ur: (1, 2),
|
248 |
+
us: (1, -1),
|
249 |
+
zx: (-0.25, -1),
|
250 |
+
ex: (0.25, -1),
|
251 |
+
}
|
252 |
+
|
253 |
+
node_c = {}
|
254 |
+
for node in G:
|
255 |
+
node_c[node] = 'lightgrey' if node in [x, a, d, r, s] else 'white'
|
256 |
+
|
257 |
+
edge_c = {e: 'black' for e in G.edges}
|
258 |
+
node_line_c = {k: 'black' for k, _ in node_c.items()}
|
259 |
+
|
260 |
+
if args[0]: # do_r
|
261 |
+
# G.remove_edge(ur, r)
|
262 |
+
edge_c[(ur, r)] = 'lightgrey'
|
263 |
+
node_line_c[r] = 'red'
|
264 |
+
if args[1]: # do_s
|
265 |
+
# G.remove_edges_from([(us, s)])
|
266 |
+
edge_c[(us, s)] = 'lightgrey'
|
267 |
+
node_line_c[s] = 'red'
|
268 |
+
if args[2]: # do_f (do_d)
|
269 |
+
# G.remove_edges_from([(ud, d), (a, d)])
|
270 |
+
edge_c[(ud, d)] = 'lightgrey'
|
271 |
+
edge_c[(a, d)] = 'lightgrey'
|
272 |
+
node_line_c[d] = 'red'
|
273 |
+
if args[3]: # do_a
|
274 |
+
# G.remove_edge(ua, a)
|
275 |
+
edge_c[(ua, a)] = 'lightgrey'
|
276 |
+
node_line_c[a] = 'red'
|
277 |
+
|
278 |
+
fs = 30
|
279 |
+
options = {
|
280 |
+
"font_size": fs,
|
281 |
+
"node_size": 3000,
|
282 |
+
"node_color": list(node_c.values()),
|
283 |
+
"edgecolors": list(node_line_c.values()),
|
284 |
+
"edge_color": list(edge_c.values()),
|
285 |
+
"linewidths": 2,
|
286 |
+
"width": 2,
|
287 |
+
}
|
288 |
+
plt.close('all')
|
289 |
+
fig, ax = plt.subplots(1, 1, figsize=(5,5))#, constrained_layout=True)
|
290 |
+
# fig.patch.set_visible(False)
|
291 |
+
ax.margins(x=0.1, y=0.08, tight=False)
|
292 |
+
ax.axis("off")
|
293 |
+
nx.draw_networkx(G, pos, **options, arrowsize=25, arrowstyle='-|>', ax=ax)
|
294 |
+
# need to reuse x, y limits so that the graphs plot the same way before and after removing edges
|
295 |
+
x_lim = (-1.32, 1.32)
|
296 |
+
y_lim = (-1.414, 2.414)
|
297 |
+
ax.set_xlim(x_lim)
|
298 |
+
ax.set_ylim(y_lim)
|
299 |
+
rect = patches.FancyBboxPatch((-0.5, -1.325), 1, 0.65, boxstyle="round, pad=0.05, rounding_size=0", linewidth=2, edgecolor='black', facecolor='none', linestyle='-')
|
300 |
+
ax.add_patch(rect)
|
301 |
+
ax.text(-0.9, -1.075, r"$\mathbf{U}_{\mathbf{x}}$", fontsize=fs)
|
302 |
+
|
303 |
+
if args[0]: # do_r
|
304 |
+
fig.figimage(HAMMER, 0.72*fig.bbox.xmax, 0.64*fig.bbox.ymax, zorder=10)
|
305 |
+
if args[1]: # do_s
|
306 |
+
fig.figimage(HAMMER, 0.72*fig.bbox.xmax, 0.395*fig.bbox.ymax, zorder=11)
|
307 |
+
if args[2]: # do_f
|
308 |
+
fig.figimage(HAMMER, 0.363*fig.bbox.xmax, 0.64*fig.bbox.ymax, zorder=12)
|
309 |
+
if args[3]: # do_a
|
310 |
+
fig.figimage(HAMMER, 0.0075*fig.bbox.xmax, 0.64*fig.bbox.ymax, zorder=13)
|
311 |
+
|
312 |
+
fig.tight_layout()
|
313 |
+
fig.canvas.draw()
|
314 |
+
return np.array(fig.canvas.renderer.buffer_rgba())
|
315 |
+
|
316 |
+
|
317 |
+
def vae_preprocess(args, pa):
|
318 |
+
if 'ukbb' in args.hps:
|
319 |
+
# preprocessing ukbb parents for the vae which was originally trained using
|
320 |
+
# log standardized parents. The pgm was trained using [-1,1] normalization
|
321 |
+
# first undo [-1,1] parent preprocessing back to original range
|
322 |
+
for k, v in pa.items():
|
323 |
+
if k != 'mri_seq' and k != 'sex':
|
324 |
+
pa[k] = (v + 1) / 2 # [-1,1] -> [0,1]
|
325 |
+
_max, _min = get_attr_max_min(k)
|
326 |
+
pa[k] = pa[k] * (_max - _min) + _min
|
327 |
+
# log_standardize parents for vae input
|
328 |
+
for k, v in pa.items():
|
329 |
+
logpa_k = torch.log(v.clamp(min=1e-12))
|
330 |
+
if k == 'age':
|
331 |
+
pa[k] = (logpa_k - 4.112339973449707) / 0.11769197136163712
|
332 |
+
elif k == 'brain_volume':
|
333 |
+
pa[k] = (logpa_k - 13.965583801269531) / 0.09537758678197861
|
334 |
+
elif k == 'ventricle_volume':
|
335 |
+
pa[k] = (logpa_k - 10.345998764038086) / 0.43127763271331787
|
336 |
+
# concatenate parents expand to input res for conditioning the vae
|
337 |
+
pa = torch.cat([pa[k] if len(pa[k].shape) > 1 else pa[k][..., None]
|
338 |
+
for k in args.parents_x], dim=1)
|
339 |
+
pa = pa[..., None, None].repeat(1, 1, *(args.input_res,)*2).to(args.device).float()
|
340 |
+
return pa
|
341 |
+
|
342 |
+
|
343 |
+
def preprocess_brain(args, obs):
|
344 |
+
obs['x'] = (obs['x'][None,...].float().to(args.device) - 127.5) / 127.5 # [-1,1]
|
345 |
+
# for all other variables except x
|
346 |
+
for k in [k for k in obs.keys() if k != 'x']:
|
347 |
+
obs[k] = obs[k].float().to(args.device).view(1, 1)
|
348 |
+
if k in ['age', 'brain_volume', 'ventricle_volume']:
|
349 |
+
k_max, k_min = get_attr_max_min(k)
|
350 |
+
obs[k] = (obs[k] - k_min) / (k_max - k_min) # [0,1]
|
351 |
+
obs[k] = 2 * obs[k] - 1 # [-1,1]
|
352 |
+
return obs
|
353 |
+
|
354 |
+
|
355 |
+
def get_fig_arr(x, width=4, height=4, dpi=144, cmap='Greys_r', norm=None):
|
356 |
+
fig = plt.figure(figsize=(width, height), dpi=dpi)
|
357 |
+
ax = plt.axes([0,0,1,1], frameon=False)
|
358 |
+
if cmap == 'Greys_r':
|
359 |
+
ax.imshow(x, cmap=cmap, vmin=0, vmax=255)
|
360 |
+
else:
|
361 |
+
ax.imshow(x, cmap=cmap, norm=norm)
|
362 |
+
ax.axis('off')
|
363 |
+
fig.canvas.draw()
|
364 |
+
return np.array(fig.canvas.renderer.buffer_rgba())
|
365 |
+
|
366 |
+
|
367 |
+
def normalize(x, x_min=None, x_max=None, zero_one=False):
|
368 |
+
if x_min is None:
|
369 |
+
x_min = x.min()
|
370 |
+
if x_max is None:
|
371 |
+
x_max = x.max()
|
372 |
+
x = (x - x_min) / (x_max - x_min) # [0,1]
|
373 |
+
return x if zero_one else 2 * x - 1 # else [-1,1]
|
checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:c5a9b5666f4f221bfd464c10cca5a092fca26abcfdd767cfc81904c45df5b7a9
|
3 |
+
size 129987615
|
checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:61b4dd71d4e9d8ff21ead93872620d3522e53402685663e1b1c946879238936b
|
3 |
+
size 493467505
|
checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:03dd765e5cdd4d6938e8536e72f597a2a5bc8ed43f4b840d5c289017aa555ace
|
3 |
+
size 5328833
|
checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint_current.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:5fcb79767cbad89f1ea481da0af49334834440d0c6b42ec4e8243a602825219d
|
3 |
+
size 5328833
|
checkpoints/m_b_v_s/sup_pgm/checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:2ad6a317c8c05946556b74a1886d4a72328b2ab523aa69f56c7bfe0476d5212f
|
3 |
+
size 5397757
|
checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:56920771b4112c430ff42f46b2707a8bb75acbaf2be518c640d6be2157be0c54
|
3 |
+
size 279070192
|
checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:014ccfb191fa8307dc9488379e94d527f2e51159c96417c88b7b17abc6e7c1b9
|
3 |
+
size 33772784
|
checkpoints/t_i_d/sup_pgm/checkpoint.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:42585f99a2f5825645f606b99ccbc7d14a2815eb7723bace03a3d94e3c2c8856
|
3 |
+
size 39144
|
data/mimic_subset/0.jpg
ADDED
![]() |
data/mimic_subset/1.jpg
ADDED
![]() |
data/mimic_subset/10.jpg
ADDED
![]() |
data/mimic_subset/11.jpg
ADDED
![]() |
data/mimic_subset/12.jpg
ADDED
![]() |
data/mimic_subset/13.jpg
ADDED
![]() |
data/mimic_subset/14.jpg
ADDED
![]() |
data/mimic_subset/15.jpg
ADDED
![]() |
data/mimic_subset/16.jpg
ADDED
![]() |
data/mimic_subset/17.jpg
ADDED
![]() |
data/mimic_subset/18.jpg
ADDED
![]() |
data/mimic_subset/19.jpg
ADDED
![]() |
data/mimic_subset/2.jpg
ADDED
![]() |
data/mimic_subset/20.jpg
ADDED
![]() |
data/mimic_subset/21.jpg
ADDED
![]() |
data/mimic_subset/22.jpg
ADDED
![]() |
data/mimic_subset/23.jpg
ADDED
![]() |
data/mimic_subset/24.jpg
ADDED
![]() |
data/mimic_subset/25.jpg
ADDED
![]() |
data/mimic_subset/26.jpg
ADDED
![]() |
data/mimic_subset/27.jpg
ADDED
![]() |
data/mimic_subset/28.jpg
ADDED
![]() |
data/mimic_subset/29.jpg
ADDED
![]() |
data/mimic_subset/3.jpg
ADDED
![]() |
data/mimic_subset/30.jpg
ADDED
![]() |
data/mimic_subset/31.jpg
ADDED
![]() |
data/mimic_subset/32.jpg
ADDED
![]() |
data/mimic_subset/33.jpg
ADDED
![]() |
data/mimic_subset/34.jpg
ADDED
![]() |
data/mimic_subset/35.jpg
ADDED
![]() |
data/mimic_subset/36.jpg
ADDED
![]() |
data/mimic_subset/37.jpg
ADDED
![]() |
data/mimic_subset/38.jpg
ADDED
![]() |
data/mimic_subset/39.jpg
ADDED
![]() |
data/mimic_subset/4.jpg
ADDED
![]() |
data/mimic_subset/40.jpg
ADDED
![]() |
data/mimic_subset/41.jpg
ADDED
![]() |