fabio-deep commited on
Commit
8c4fe8b
·
0 Parent(s):

Duplicate from fabio-deep/counterfactuals

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +34 -0
  2. .gitignore +2 -0
  3. README.md +14 -0
  4. app.py +556 -0
  5. app_utils.py +373 -0
  6. checkpoints/a_r_s_f/mimic_beta9_gelu_dgauss_1_lr3/checkpoint.pt +3 -0
  7. checkpoints/a_r_s_f/mimic_dscm_lr_1e5_lagrange_lr_1_damping_10/6500_checkpoint.pt +3 -0
  8. checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint.pt +3 -0
  9. checkpoints/a_r_s_f/sup_pgm_mimic/checkpoint_current.pt +3 -0
  10. checkpoints/m_b_v_s/sup_pgm/checkpoint.pt +3 -0
  11. checkpoints/m_b_v_s/ukbb192_beta5_dgauss_b33/checkpoint.pt +3 -0
  12. checkpoints/t_i_d/dgauss_cond_big_beta1_dropexo/checkpoint.pt +3 -0
  13. checkpoints/t_i_d/sup_pgm/checkpoint.pt +3 -0
  14. data/mimic_subset/0.jpg +0 -0
  15. data/mimic_subset/1.jpg +0 -0
  16. data/mimic_subset/10.jpg +0 -0
  17. data/mimic_subset/11.jpg +0 -0
  18. data/mimic_subset/12.jpg +0 -0
  19. data/mimic_subset/13.jpg +0 -0
  20. data/mimic_subset/14.jpg +0 -0
  21. data/mimic_subset/15.jpg +0 -0
  22. data/mimic_subset/16.jpg +0 -0
  23. data/mimic_subset/17.jpg +0 -0
  24. data/mimic_subset/18.jpg +0 -0
  25. data/mimic_subset/19.jpg +0 -0
  26. data/mimic_subset/2.jpg +0 -0
  27. data/mimic_subset/20.jpg +0 -0
  28. data/mimic_subset/21.jpg +0 -0
  29. data/mimic_subset/22.jpg +0 -0
  30. data/mimic_subset/23.jpg +0 -0
  31. data/mimic_subset/24.jpg +0 -0
  32. data/mimic_subset/25.jpg +0 -0
  33. data/mimic_subset/26.jpg +0 -0
  34. data/mimic_subset/27.jpg +0 -0
  35. data/mimic_subset/28.jpg +0 -0
  36. data/mimic_subset/29.jpg +0 -0
  37. data/mimic_subset/3.jpg +0 -0
  38. data/mimic_subset/30.jpg +0 -0
  39. data/mimic_subset/31.jpg +0 -0
  40. data/mimic_subset/32.jpg +0 -0
  41. data/mimic_subset/33.jpg +0 -0
  42. data/mimic_subset/34.jpg +0 -0
  43. data/mimic_subset/35.jpg +0 -0
  44. data/mimic_subset/36.jpg +0 -0
  45. data/mimic_subset/37.jpg +0 -0
  46. data/mimic_subset/38.jpg +0 -0
  47. data/mimic_subset/39.jpg +0 -0
  48. data/mimic_subset/4.jpg +0 -0
  49. data/mimic_subset/40.jpg +0 -0
  50. 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("### &nbsp;")
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("### &nbsp;")
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("### &nbsp;")
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