update

.gitignore

@@ -1,3 +1,4 @@
 .vscode
+.gradio
 __pycache__
 *.pyc

README.md

@@ -4,11 +4,27 @@ emoji: 🌖
 colorFrom: purple
 colorTo: green
 sdk: gradio
-sdk_version: 3.27.0
+sdk_version: 5.35.0
 app_file: app.py
 pinned: false
 license: mit
 duplicated_from: fabio-deep/counterfactuals
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Code for the **ICML 2023** paper:
+
+[**High Fidelity Image Counterfactuals with Probabilistic Causal Models**](https://arxiv.org/abs/2306.15764)
+
+Fabio De Sousa Ribeiro<sup>1</sup>, Tian Xia<sup>1</sup>, Miguel Monteiro<sup>1</sup>, Nick Pawlowski<sup>2</sup>, Ben Glocker<sup>1</sup>\
+<sup>1</sup>Imperial College London, <sup>2</sup>Microsoft Research Cambridge, UK
+
+```
+@misc{ribeiro2023high,
+      title={High Fidelity Image Counterfactuals with Probabilistic Causal Models},
+      author={Fabio De Sousa Ribeiro and Tian Xia and Miguel Monteiro and Nick Pawlowski and Ben Glocker},
+      year={2023},
+      eprint={2306.15764},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG}
+}
+```

app.py

@@ -401,36 +401,47 @@ def infer_chest_cf(*args):
     return (cf_x, cf_x_std, effect, cf_r, cf_s, cf_f, np.round(cf_a, 1))
 
 
-with gr.Blocks(theme=gr.themes.Default()) as demo:
+js_func = """
+function refresh() {
+    const url = new URL(window.location);
+    if (url.searchParams.get('__theme') !== 'light') {
+        url.searchParams.set('__theme', 'light');
+        window.location.href = url.href;
+    }
+}
+"""
+
+with gr.Blocks(
+    # theme=gr.themes.Default(),
+    theme="shivi/calm_seafoam",
+    js=js_func
+) as demo:
+    img_cfg = dict(
+        interactive=False,
+        height=HEIGHT,
+        show_download_button=False,
+        show_fullscreen_button=False
+    )
     with gr.Tabs():
         with gr.TabItem("Morpho-MNIST") as mnist_tab:
             mnist_id = gr.Textbox(value=mnist_tab.label, visible=False)
-
-            with gr.Row().style(equal_height=True):
+            with gr.Row(equal_height=True):
                 idx = gr.Number(value=0, visible=False)
                 with gr.Column(scale=1, min_width=200):
-                    x = gr.Image(label="Observation", interactive=False).style(
-                        height=HEIGHT
-                    )
+                    x = gr.Image(label="Observation", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    cf_x = gr.Image(label="Counterfactual", interactive=False).style(
-                        height=HEIGHT
-                    )
+                    cf_x = gr.Image(label="Counterfactual", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    cf_x_std = gr.Image(
-                        label="Counterfactual Uncertainty", interactive=False
-                    ).style(height=HEIGHT)
+                    cf_x_std = gr.Image(label="Uncertainty", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    effect = gr.Image(
-                        label="Direct Causal Effect", interactive=False
-                    ).style(height=HEIGHT)
-            with gr.Row().style(equal_height=True):
+                    effect = gr.Image(label="Causal Effect", **img_cfg)
+            with gr.Row(equal_height=True):
                 with gr.Column(scale=1.75):
                     gr.Markdown(
                         "**Intervention**"
-                        + 20 * " "
-                        + "[arXiv paper](https://arxiv.org/abs/2306.15764)   |   [GitHub code](https://github.com/biomedia-mira/causal-gen)"
-                        + "  |   Hint: try 90% zoom"
+                        + 22 * " "
+                        + "[Paper](https://proceedings.mlr.press/v202/de-sousa-ribeiro23a.html)   |   [Code](https://github.com/biomedia-mira/causal-gen)"
+                        # + "  |   Hint: try 90% zoom"
                     )
                     with gr.Column():
                         do_y = gr.Checkbox(label="do(digit)", value=False)
@@ -460,38 +471,34 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                         submit = gr.Button("Submit", variant="primary")
                 with gr.Column(scale=1):
                     gr.Markdown("###  ")
-                    causal_graph = gr.Image(
-                        label="Causal Graph", interactive=False
-                    ).style(height=300)
+                    img_cfg["height"] = 300
+                    causal_graph = gr.Image(label="Causal Graph", **img_cfg)
+                    img_cfg["height"] = HEIGHT
 
         with gr.TabItem("Brain MRI") as brain_tab:
             brain_id = gr.Textbox(value=brain_tab.label, visible=False)
 
-            with gr.Row().style(equal_height=True):
+            with gr.Row(equal_height=True):
                 idx_brain = gr.Number(value=0, visible=False)
                 with gr.Column(scale=1, min_width=200):
-                    x_brain = gr.Image(label="Observation", interactive=False).style(
-                        height=HEIGHT
-                    )
+                    x_brain = gr.Image(label="Observation", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    cf_x_brain = gr.Image(
-                        label="Counterfactual", interactive=False
-                    ).style(height=HEIGHT)
+                    cf_x_brain = gr.Image(label="Counterfactual", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
                     cf_x_std_brain = gr.Image(
-                        label="Counterfactual Uncertainty", interactive=False
-                    ).style(height=HEIGHT)
+                        label="Uncertainty", **img_cfg
+                    )
                 with gr.Column(scale=1, min_width=200):
                     effect_brain = gr.Image(
-                        label="Direct Causal Effect", interactive=False
-                    ).style(height=HEIGHT)
+                        label="Causal Effect", **img_cfg
+                    )
             with gr.Row():
                 with gr.Column(scale=2.55):
                     gr.Markdown(
                         "**Intervention**"
-                        + 20 * " "
-                        + "[arXiv paper](https://arxiv.org/abs/2306.15764)   |   [GitHub code](https://github.com/biomedia-mira/causal-gen)"
-                        + "  |   Hint: try 90% zoom"
+                        + 22 * " "
+                        + "[Paper](https://proceedings.mlr.press/v202/de-sousa-ribeiro23a.html)   |   [Code](https://github.com/biomedia-mira/causal-gen)"
+                        # + "  |   Hint: try 90% zoom"
                     )
                     with gr.Row():
                         with gr.Column(min_width=200):
@@ -543,41 +550,33 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                         submit_brain = gr.Button("Submit", variant="primary")
                 with gr.Column(scale=1):
                     # gr.Markdown("###  ")
-                    causal_graph_brain = gr.Image(
-                        label="Causal Graph", interactive=False
-                    ).style(height=340)
+                    img_cfg["height"] = 340
+                    causal_graph_brain = gr.Image(label="Causal Graph", **img_cfg)
+                    img_cfg["height"] = HEIGHT
 
         with gr.TabItem("Chest X-ray") as chest_tab:
             chest_id = gr.Textbox(value=chest_tab.label, visible=False)
 
-            with gr.Row().style(equal_height=True):
+            with gr.Row(equal_height=True):
                 idx_chest = gr.Number(value=0, visible=False)
                 with gr.Column(scale=1, min_width=200):
-                    x_chest = gr.Image(label="Observation", interactive=False).style(
-                        height=HEIGHT
-                    )
+                    x_chest = gr.Image(label="Observation", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    cf_x_chest = gr.Image(
-                        label="Counterfactual", interactive=False
-                    ).style(height=HEIGHT)
+                    cf_x_chest = gr.Image(label="Counterfactual", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    cf_x_std_chest = gr.Image(
-                        label="Counterfactual Uncertainty", interactive=False
-                    ).style(height=HEIGHT)
+                    cf_x_std_chest = gr.Image(label="Uncertainty", **img_cfg)
                 with gr.Column(scale=1, min_width=200):
-                    effect_chest = gr.Image(
-                        label="Direct Causal Effect", interactive=False
-                    ).style(height=HEIGHT)
+                    effect_chest = gr.Image(label="Causal Effect", **img_cfg)
 
             with gr.Row():
                 with gr.Column(scale=2.55):
                     gr.Markdown(
                         "**Intervention**"
-                        + 20 * " "
-                        + "[arXiv paper](https://arxiv.org/abs/2306.15764)   |   [GitHub code](https://github.com/biomedia-mira/causal-gen)"
-                        + "  |   Hint: try 90% zoom"
+                        + 22 * " "
+                        + "[Paper](https://proceedings.mlr.press/v202/de-sousa-ribeiro23a.html)   |   [Code](https://github.com/biomedia-mira/causal-gen)"
+                        # + "  |   Hint: try 90% zoom"
                     )
-                    with gr.Row().style(equal_height=True):
+                    with gr.Row(equal_height=True):
                         with gr.Column(min_width=200):
                             do_f_chest = gr.Checkbox(label="do(disease)", value=False)
                             f_chest = gr.Radio(FIND_CAT, label="", interactive=False)
@@ -603,9 +602,9 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
                         submit_chest = gr.Button("Submit", variant="primary")
                 with gr.Column(scale=1):
                     # gr.Markdown("###  ")
-                    causal_graph_chest = gr.Image(
-                        label="Causal Graph", interactive=False
-                    ).style(height=345)
+                    img_cfg["height"] = 345
+                    causal_graph_chest = gr.Image(label="Causal Graph", **img_cfg)
+                    img_cfg["height"] = HEIGHT
 
     # morphomnist
     do = [do_t, do_i, do_y]

pgm/layers.py

@@ -91,7 +91,7 @@ class CNN(nn.Module):
 
 
 class ArgMaxGumbelMax(Transform):
-    r"""ArgMax as Transform, but inv conditioned on logits"""
+    r"""ArgMax as Transform, with inverse conditioned on logits"""
 
     def __init__(self, logits, event_dim=0, cache_size=0):
         super(ArgMaxGumbelMax, self).__init__(cache_size=cache_size)
@@ -106,9 +106,6 @@ class ArgMaxGumbelMax(Transform):
         return self._event_dim
 
     def __call__(self, gumbels):
-        """
-        Computes the forward transform
-        """
         assert self.logits != None, "Logits not defined."
 
         if self._cache_size == 0:
@@ -118,20 +115,12 @@ class ArgMaxGumbelMax(Transform):
         return y
 
     def _call(self, gumbels):
-        """
-        Abstract method to compute forward transformation.
-        """
         assert self.logits != None, "Logits not defined."
         y = gumbels + self.logits
-        # print(f'y: {y}')
-        # print(f'logits: {self.logits}')
         return y.argmax(-1, keepdim=True)
 
     @property
     def domain(self):
-        """ "
-        Domain of input(gumbel variables), Real
-        """
         if self.event_dim == 0:
             return constraints.real
         return constraints.independent(constraints.real, self.event_dim)
@@ -148,18 +137,14 @@ class ArgMaxGumbelMax(Transform):
     def inv(self, k):
         """Infer the gumbels noises given k and logits."""
         assert self.logits != None, "Logits not defined."
-
         uniforms = torch.rand(
             self.logits.shape, dtype=self.logits.dtype, device=self.logits.device
         )
         gumbels = -((-(uniforms.log())).log())
-        # print(f'gumbels: {gumbels.size()}, {gumbels.dtype}')
         # (batch_size, num_classes) mask to select kth class
-        # print(f'k : {k.size()}')
         mask = F.one_hot(
             k.squeeze(-1).to(torch.int64), num_classes=self.logits.shape[-1]
         )
-        # print(f'mask: {mask.size()}, {mask.dtype}')
         # (batch_size, 1) select topgumbel for truncation of other classes
         topgumbel = (mask * gumbels).sum(dim=-1, keepdim=True) - (
             mask * self.logits
@@ -173,41 +158,25 @@ class ArgMaxGumbelMax(Transform):
         return epsilons
 
     def log_abs_det_jacobian(self, x, y):
-        """We use the log_abs_det_jacobian to account for the categorical prob
-        x: Gumbels; y: argmax(x+logits)
-        P(y) = softmax
-        """
-        # print(f"logits: {torch.log(F.softmax(self.logits, dim=-1)).size()}")
-        # print(f'y: {y.size()} ')
-        # print(f"log_abs_det_jacobian: {self._categorical.log_prob(y.squeeze(-1)).unsqueeze(-1).size()}")
+        """bit hacky for now"""
         return -self._categorical.log_prob(y.squeeze(-1)).unsqueeze(-1)
 
 
 class ConditionalGumbelMax(ConditionalTransformModule):
-    r"""Given gumbels+logits, output the OneHot Categorical"""
-
     def __init__(self, context_nn, event_dim=0, **kwargs):
-        # The logits_nn which predict the logits given ages:
         super().__init__(**kwargs)
         self.context_nn = context_nn
         self.event_dim = event_dim
 
     def condition(self, context):
-        """Given context (age), output the Categorical results"""
-        logits = self.context_nn(
-            context
-        )  # The logits for calculating argmax(Gumbel + logits)
+        logits = self.context_nn(context)
         return ArgMaxGumbelMax(logits)
 
     def _logits(self, context):
-        """Return logits given context"""
         return self.context_nn(context)
 
     @property
     def domain(self):
-        """ "
-        Domain of input(gumbel variables), Real
-        """
         if self.event_dim == 0:
             return constraints.real
         return constraints.independent(constraints.real, self.event_dim)
@@ -224,6 +193,7 @@ class ConditionalGumbelMax(ConditionalTransformModule):
 
 class TransformedDistributionGumbelMax(TransformedDistribution, TorchDistributionMixin):
     r"""Define a TransformedDistribution class for Gumbel max"""
+
     arg_constraints: Dict[str, constraints.Constraint] = {}
 
     def log_prob(self, value):
@@ -231,7 +201,6 @@ class TransformedDistributionGumbelMax(TransformedDistribution, TorchDistributio
         We do not use the log_prob() of the base Gumbel distribution, because the likelihood for
         each class for Gumbel Max sampling is determined by the logits.
         """
-        # print("This happens")
         if self._validate_args:
             self._validate_sample(value)
         event_dim = len(self.event_shape)
@@ -245,7 +214,6 @@ class TransformedDistributionGumbelMax(TransformedDistribution, TorchDistributio
                 event_dim - transform.domain.event_dim,
             )
             y = x
-        # print(f"log_prob: {log_prob.size()}")
         return log_prob
 
 
@@ -253,7 +221,6 @@ class ConditionalTransformedDistributionGumbelMax(ConditionalTransformedDistribu
     def condition(self, context):
         base_dist = self.base_dist.condition(context)
         transforms = [t.condition(context) for t in self.transforms]
-        # return TransformedDistribution(base_dist, transforms)
         return TransformedDistributionGumbelMax(base_dist, transforms)
 
     def clear_cache(self):

requirements.txt

@@ -1,4 +1,4 @@
-gradio==3.27.0
+gradio==5.35.0
 matplotlib==3.7.1
 networkx==2.8.4
 numpy==1.24.3