Update inference

app.py

@@ -1,23 +1,14 @@
 import gradio as gr
 import numpy as np
+import mplhep as hep
 import matplotlib.pyplot as plt
 import tensorflow as tf
 
 from huggingface_hub import from_pretrained_keras
-from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Input, Flatten
-from qkeras import *
-
-
-def load_keras_model(model_path: str):
-    org_model = from_pretrained_keras(model_path)
-    input_ = Input(shape=(18, 14), name="inputs_")
-    x = Flatten()(input_)
-    for layer in org_model.layers[1:]:
-        x = layer(x)
-    output = Activation("linear", name="outputs")(x)
-    return Model(input_, output, name="cicada")
 
+model_v1 = from_pretrained_keras("cicada-project/cicada-v1.1")
+model_v2 = from_pretrained_keras("cicada-project/cicada-v2.1")
+hep.style.use("CMS")
 
 def parse_input(et):
     if not et:
@@ -32,41 +23,82 @@ def parse_input(et):
     if np.any(et < 0) or np.any(et > 1023):
         raise gr.Error("The input has to be in a range (0, 1023)!")
 
-    return et.reshape(1, 18, 14)
-
-
-def inference(input_):
-    input_ = parse_input(input_)
-    return model.predict(input_)[0][0]
+    return et.reshape(1, 252)
 
 
-def saliency(input_):
-    input_ = parse_input(input_)
+def saliency(input_, version):
     x = tf.constant(input_)
     with tf.GradientTape() as tape:
         tape.watch(x)
-        predictions = model(x)
+        if version == "v1":
+            predictions = model_v1(x)
+        elif version == "v2":
+            predictions = model_v2(x)
     gradient = tape.gradient(predictions, x)
     gradient = gradient.numpy()
     min_val, max_val = np.min(gradient), np.max(gradient)
     gradient = (gradient - min_val) / (max_val - min_val + tf.keras.backend.epsilon())
 
-    fig_i = plt.figure()
-    plt.imshow(input_.reshape(18, 14), cmap="Reds")
-    plt.colorbar()
-    plt.axis("off")
-    plt.tight_layout()
-
     fig_s = plt.figure()
-    plt.imshow(gradient.reshape(18, 14), cmap="Greys")
-    plt.colorbar()
-    plt.axis("off")
-    plt.tight_layout()
-
-    return fig_i, fig_s
-
+    im = plt.imshow(gradient.reshape(18, 14), vmin=0., vmax=1., cmap="Greys")
+    ax = plt.gca()
+    cbar = ax.figure.colorbar(im, ax=ax)
+    cbar.ax.set_ylabel(r"Calorimeter Saliency (a.u.)")
+    plt.xticks(np.arange(14), labels=np.arange(4, 18))
+    plt.yticks(
+        np.arange(18),
+        labels=np.arange(18)[::-1],
+        rotation=90,
+        va="center",
+    )
+    plt.xlabel(r"i$\eta$")
+    plt.ylabel(r"i$\phi$")
+
+    return fig_s
+
+
+def draw_input(input_):
+    fig_i = plt.figure()
+    im = plt.imshow(input_.reshape(18, 14), vmin=0, vmax=input_.max(), cmap="Purples")
+    ax = plt.gca()
+    cbar = ax.figure.colorbar(im, ax=ax)
+    cbar.ax.set_ylabel(r"Calorimeter E$_T$ deposit (GeV)")
+    plt.xticks(np.arange(14), labels=np.arange(4, 18))
+    plt.yticks(
+        np.arange(18),
+        labels=np.arange(18)[::-1],
+        rotation=90,
+        va="center",
+    )
+    plt.xlabel(r"i$\eta$")
+    plt.ylabel(r"i$\phi$")
+    return fig_i
+
+
+def inference(input_, version):
+    if version == "v1":
+        return model_v1.predict(input_)[0][0]
+    elif version == "v2":
+        return model_v2.predict(input_)[0][0]
+
+
+def generate():
+    matrix = np.clip(np.random.zipf(2, 252) - 1, 0, 1023)
+    matrix = matrix.reshape(18, 14).astype(str)
+    rows = [",".join(row) for row in matrix]
+    return "\n".join(rows)
+
+
+def process_request(input_):
+    input_ = parse_input(input_)
+    return (
+        inference(input_, "v1"),
+        inference(input_, "v2"),
+        draw_input(input_),
+        saliency(input_, "v1"),
+        saliency(input_, "v2"),
+    )
 
-model = load_keras_model("cicada-project/cicada-v1.1")
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -77,18 +109,34 @@ with gr.Blocks() as demo:
                 placeholder="\n".join([",".join(["0"] * 14)] * 18),
             )
             with gr.Row():
-                classify = gr.Button("Get Anomaly Score")
-                interpret = gr.Button("Visualize")
-        with gr.Column():
-            label = gr.Number(label="CICADA Anomaly Score")
+                generate_input = gr.Button("Generate random input")
+                magic = gr.Button("Do CICADA inference")
+
         with gr.Column():
-            with gr.Tabs():
-                with gr.TabItem("Display Calorimeter Input"):
-                    input_plot = gr.Plot()
-                with gr.TabItem("Display Saliency Map"):
-                    interpretation_plot = gr.Plot()
-
-    classify.click(inference, input_, label)
-    interpret.click(saliency, input_, [input_plot, interpretation_plot])
+            with gr.Row():
+                label_v1 = gr.Number(label="CICADA Anomaly Score for CICADA v1")
+            with gr.Row():
+                label_v2 = gr.Number(label="CICADA Anomaly Score for CICADA v2")
+            with gr.Row():
+                with gr.Tabs():
+                    with gr.TabItem("Calorimeter Input"):
+                        input_plot = gr.Plot()
+                    with gr.TabItem("Saliency Map for CICADAv1"):
+                        interpretation_plot_v1 = gr.Plot()
+                    with gr.TabItem("Saliency Map for CICADAv2"):
+                        interpretation_plot_v2 = gr.Plot()
+
+    generate_input.click(generate, None, input_)
+    magic.click(
+        process_request,
+        input_,
+        [
+            label_v1,
+            label_v2,
+            input_plot,
+            interpretation_plot_v1,
+            interpretation_plot_v2,
+        ],
+    )
 
 demo.launch()

requirements.txt

@@ -1,5 +1,6 @@
-numpy
-matplotlib
-tensorflow
-huggingface_hub
-qkeras
+numpy==1.26.0
+gradio==4.8.0
+matplotlib==3.7.2
+mplhep==0.3.31
+tensorflow==2.10.0
+huggingface_hub==0.16.4