Update to space

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -1,24 +1,18 @@
-from functools import partial
-
-import gradio as gr
-import numpy as np
-from scipy.ndimage import rotate
-
-
 from functools import partial
 from typing import Any, Callable, Sequence, Tuple
 
 import dirichlet
-
+import gradio as gr
 import jax
 import jax.numpy as jnp
 import numpy as np
 import plotly.graph_objects as go
 import tensorflow_probability.substrates.jax as tfp
 from flax import linen as nn
-
-
+from flax.traverse_util import flatten_dict, unflatten_dict
+from safetensors.flax import load_file, save_file
 from scipy.ndimage import rotate
+from utils import MLP, get_apply_fn
 
 tfd = tfp.distributions
 
@@ -27,6 +21,59 @@ image = np.load("image_2.npy")[0, ..., 0]
 
 rotate_image = partial(rotate, reshape=False, axes=(0, 1))
 
+model = MLP(512)
+apply_fn, apply_map_fn = get_apply_fn(model)
+params = unflatten_dict(load_file("model.safetensors"), sep=".")
+
+rng_test = jax.random.split(jax.random.PRNGKey(0), 1000)
+
+
+def predict(image):
+    image = image[None, ..., None] / 255
+
+    y_logits_map = apply_map_fn(params, image, rng_test)[:, 0, :]
+    y_logits = apply_fn(params, image, rng_test)[:, 0, :]
+    
+    y_probs = jax.nn.softmax(y_logits)
+    y_probs_map = jax.nn.softmax(y_logits_map)
+
+    alpha = dirichlet.mle(y_probs, method="fixedpoint")
+
+    predictive = tfd.DirichletMultinomial(1, alpha)
+    entropy = -predictive.log_prob(y_probs).mean()
+
+    predictive_map = tfd.Categorical(logits=y_logits_map)
+    entropy_map = predictive_map.entropy().mean()
+
+
+    fig = go.Figure()
+    for class_id in range(10):
+        fig.add_trace(
+            go.Box(y=y_probs[:, class_id], name=str(class_id), boxpoints=False)
+        )
+
+    fig.update_layout(
+        title="Predicted labels (probabilistic model)",
+        xaxis_title="Class",
+        yaxis_title="Probability",
+        yaxis=dict(range=[0, 1.1]),
+    )
+
+    fig_map = go.Figure()
+    for class_id in range(10):
+        fig_map.add_trace(
+            go.Box(y=y_probs_map[:, class_id], name=str(class_id), boxpoints=False)
+        )
+
+    fig_map.update_layout(
+        title="Predicted labels (deterministic model)",
+        xaxis_title="Class",
+        yaxis_title="Probability",
+        yaxis=dict(range=[0, 1.1]),
+    )
+
+    return fig, fig_map, f"{entropy:.3f}", f"{entropy_map:.3f}"
+
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -41,13 +88,13 @@ with gr.Blocks() as demo:
             slider = gr.Slider(0, 360, step=1, label="Rotation angle")
             output = gr.Image(image_mode="L", label="Rotated", height=300, width=300)
 
-        # with gr.Column():
-        #     plot_map = gr.Plot()
-        #     entropy_map = gr.Label(label="Entropy of predictions (deterministic)")
+        with gr.Column():
+            plot_map = gr.Plot()
+            entropy_map = gr.Label(label="Entropy of predictions (deterministic)")
 
-        # with gr.Column():
-        #     plot = gr.Plot()
-        #     entropy = gr.Label(label="Entropy of predictions (probabilistic)")
+        with gr.Column():
+            plot = gr.Plot()
+            entropy = gr.Label(label="Entropy of predictions (probabilistic)")
 
     slider.release(
         rotate_image,
@@ -56,22 +103,14 @@ with gr.Blocks() as demo:
         api_name="rotate",
     )
 
-    # slider.release(
-    #     predict_label,
-    #     inputs=[output],
-    #     outputs=[plot, plot_map, entropy, entropy_map],
-    #     api_name="predict",
-    # )
-#     # with gr.Row():
-#     #     number = gr.Number(label="On release")
-#     #     number2 = gr.Number(label="Number of events fired")
-#     # slider.release(
-#     #     identity,
-#     #     inputs=[slider, state],
-#     #     outputs=[number, state, number2],
-#     #     api_name="predict",
-#     # )
-
-
-# demo.launch(share=False, server_name="0.0.0.0", server_port=43000)
+    slider.release(
+        predict,
+        inputs=[output],
+        outputs=[plot, plot_map, entropy, entropy_map],
+        api_name="predict",
+    )
+
+
+
+
 demo.launch()

requirements.txt

@@ -3,4 +3,5 @@ git+https://github.com/ericsuh/dirichlet.git
 jax[cpu]==0.4.14
 plotly==5.18.0
 tensorflow-probability==0.23.0
-flax==0.7.2
+flax==0.7.2
+safetensors==0.4.2

utils.py

@@ -0,0 +1,72 @@
+from functools import partial
+from typing import Any, Callable, Sequence, Tuple
+
+import dirichlet
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tensorflow_probability.substrates.jax as tfp
+from flax import linen as nn
+from flax.training.checkpoints import restore_checkpoint
+
+tfd = tfp.distributions
+
+
+def split_tree(a, rng_key):
+    treedef = jax.tree_util.tree_structure(a)
+    num_vars = len(jax.tree_util.tree_leaves(a))
+    all_keys = jax.random.split(rng_key, num=(num_vars + 1))
+    return jax.tree_util.tree_unflatten(treedef, all_keys[1:])
+
+
+def sample_fn(rng, vi_params: nn.FrozenDict):
+    rng = split_tree(vi_params["mean"], rng)
+    params = jax.tree_map(
+        lambda m, ls, k: tfd.Normal(loc=m, scale=jnp.exp(ls)).sample(seed=k),
+        vi_params["mean"],
+        vi_params["log_scale"],
+        rng,
+    )  # type: nn.FrozenDict
+    return params
+
+
+def get_apply_fn(model: nn.Module):
+    """Returns the model forward function"""
+
+    @jax.jit
+    @partial(jax.vmap, in_axes=(None, None, 0))
+    def apply_fn(vi_params, inputs, rng):
+        params = sample_fn(rng, vi_params)
+        outputs = model.apply({"params": params}, inputs)
+        return outputs
+
+    @jax.jit
+    def apply_map_fn(params, inputs, rng):
+        outputs = model.apply({"params": params["mean"]}, inputs)
+        return outputs[None, ...]
+
+    return apply_fn, apply_map_fn
+
+
+class MLP(nn.Module):
+    n_features: int = 512
+    n_layers: int = 3
+    n_classes: int = 10
+    n_features_mult: int = 1
+    bias_init: Callable = nn.initializers.zeros_init()
+    act: Callable = nn.relu
+    dtype: str = "float32"
+
+    @nn.compact
+    def __call__(self, x):
+        dense = partial(
+            nn.Dense,
+            dtype=self.dtype,
+            bias_init=self.bias_init,
+        )
+        x = jnp.reshape(x, (x.shape[0], -1))
+        for _ in range(self.n_layers):
+            x = dense(int(self.n_features * self.n_features_mult))(x)
+            x = nn.relu(x)
+        x = dense(self.n_classes)(x)
+        return x