app.py

asdflkj

app.py

@@ -0,0 +1,94 @@
+import gradio as gr
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+
+# ---- Paths to model + stats (put these files in the same directory) ----
+MODEL_PATH = "small32cnn_mlbt_mmat.keras"
+STATS_PATH = "preproc_stats_smallcnn.npz"
+
+# ---- Load model and stats ----
+model = tf.keras.models.load_model(MODEL_PATH, compile=False)
+
+stats = np.load(STATS_PATH)
+MEAN = float(stats["MEAN"])
+STD  = float(stats["STD"])
+
+CLASS_NAMES = ["MMAT", "MLBT"]  # 0 -> MMAT, 1 -> MLBT
+
+
+# ---- Preprocessing: Image -> (1, 32, 32, 1) normalized ----
+def preprocess(img: np.ndarray) -> np.ndarray:
+    """
+    img: H x W x C in [0, 255] from Gradio (numpy)
+    returns: (1, 32, 32, 1) float32, z-score normalized
+    """
+
+    # Convert to grayscale
+    if img.ndim == 3 and img.shape[2] == 3:
+        # RGB -> grayscale
+        img_gray = np.dot(img[..., :3], [0.2989, 0.5870, 0.1140])
+    elif img.ndim == 3 and img.shape[2] == 1:
+        img_gray = img[..., 0]
+    elif img.ndim == 2:
+        img_gray = img
+    else:
+        # fallback: take first channel
+        img_gray = img[..., 0]
+
+    # Resize to 32x32
+    pil_img = Image.fromarray(img_gray.astype("uint8"))
+    pil_img = pil_img.resize((32, 32), Image.BILINEAR)
+    arr = np.array(pil_img).astype("float32")
+
+    # z-score normalization using training stats
+    arr = (arr - MEAN) / STD
+
+    # Add batch and channel dims: (1, 32, 32, 1)
+    arr = arr[None, ..., None]
+    return arr
+
+
+# ---- Prediction function for Gradio ----
+def predict(img: np.ndarray):
+    """
+    Returns a dict {class_name: probability} for gr.Label
+    """
+    x = preprocess(img)
+    probs = model.predict(x, verbose=0)[0, 0]  # scalar prob for class 1 (MLBT)
+    p_mlbt = float(probs)
+    p_mmat = float(1.0 - p_mlbt)
+
+    return {
+        "MLBT": p_mlbt,
+        "MMAT": p_mmat,
+    }
+
+
+# ---- Gradio interface ----
+input_component = gr.Image(
+    type="numpy",
+    label="Jet image (grayscale or RGB, any size ≥ 32x32)"
+)
+
+output_component = gr.Label(
+    num_top_classes=2,
+    label="Predicted probabilities"
+)
+
+examples = []  # you can drop some example jet PNGs here later
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=input_component,
+    outputs=output_component,
+    title="MLBT vs MMAT Jet Classifier (Small 32×32 CNN)",
+    description=(
+        "Upload a jet image (32×32 heatmap or larger) and this model "
+        "predicts whether it came from the MLBT or MMAT energy-loss module."
+    ),
+    examples=examples
+)
+
+if __name__ == "__main__":
+    demo.launch()