Create app.py

app.py

@@ -0,0 +1,105 @@
+import os
+import sys
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+from scipy.ndimage import rotate, gaussian_filter
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(repo_id="MultivexAI/RobustMNIST-v1.0", filename="model.py", local_dir=".")
+hf_hub_download(repo_id="MultivexAI/RobustMNIST-v1.0", filename="model.pt", local_dir=".")
+
+sys.path.append(os.path.abspath("."))
+from model import HierarchicalNetwork
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = HierarchicalNetwork(out_dims=11).to(DEVICE)
+model.load_state_dict(torch.load("model.pt", map_location=DEVICE))
+model.eval()
+
+def preprocess_and_predict(sketch_data, rotation_val, noise_val, blur_val):
+    if sketch_data is None:
+        return None, {}
+
+    if isinstance(sketch_data, dict):
+        img_array = sketch_data.get("composite", None)
+        if img_array is None:
+            layers = sketch_data.get("layers", [])
+            img_array = layers[0] if layers else None
+    else:
+        img_array = sketch_data
+
+    if img_array is None:
+        return None, {}
+
+    pil_img = Image.fromarray(img_array.astype('uint8'))
+    if pil_img.mode == 'RGBA':
+        _, _, _, alpha_channel = pil_img.split()
+        if np.array(alpha_channel).max() > 0:
+            gray_img = alpha_channel
+        else:
+            gray_img = pil_img.convert('L')
+    else:
+        gray_img = pil_img.convert('L')
+
+    resized_img = gray_img.resize((28, 28), Image.Resampling.LANCZOS)
+    np_img = np.array(resized_img).astype(np.float32)
+
+    border_average = (np_img[0, :].mean() + np_img[-1, :].mean() + np_img[:, 0].mean() + np_img[:, -1].mean()) / 4.0
+    if border_average > 127.5:
+        np_img = 255.0 - np_img
+
+    if rotation_val > 0:
+        np_img = rotate(np_img, rotation_val, reshape=False, order=1, mode='constant', cval=0.0)
+
+    if blur_val > 0:
+        np_img = gaussian_filter(np_img, sigma=blur_val)
+
+    if noise_val > 0:
+        variance_scale = noise_val * 255.0
+        additive_noise = np.random.normal(0, variance_scale, np_img.shape)
+        np_img = np.clip(np_img + additive_noise, 0.0, 255.0)
+
+    normalized_array = np_img / 255.0
+    tensor_input = torch.tensor(normalized_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(DEVICE)
+
+    with torch.inference_mode():
+        logits = model(tensor_input)
+        probabilities = F.softmax(logits, dim=1).cpu().numpy()[0]
+
+    class_labels = [str(i) for i in range(10)] + ["Unknown"]
+    distribution = {class_labels[i]: float(probabilities[i]) for i in range(11)}
+    preview_output = np.clip(np_img, 0, 255).astype(np.uint8)
+
+    return preview_output, distribution
+
+with gr.Blocks(title="Robust MNIST Classifier") as interface:
+    gr.Markdown("## Robust Hierarchical Classifier")
+    gr.Markdown("Draw a single digit, adjust the sliders to apply synthetic environmental distortions, and observe the robustness profile.")
+    
+    with gr.Row():
+        with gr.Column():
+            canvas = gr.Sketchpad(
+                label="Draw Digit", 
+                type="numpy",
+                crop_to_bbox=False
+            )
+            rotation = gr.Slider(minimum=0, maximum=180, value=0, step=1, label="Rotation Angle (Degrees)")
+            noise = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.05, label="Gaussian Noise Level")
+            blur = gr.Slider(minimum=0.0, maximum=5.0, value=0.0, step=0.1, label="Gaussian Blur (Sigma)")
+            run_btn = gr.Button("Evaluate Signature", variant="primary")
+            
+        with gr.Column():
+            preview = gr.Image(label="Model-View Reconstruction (28x28)", image_mode="L")
+            probabilities_output = gr.Label(num_top_classes=5, label="Probability Map Output")
+
+    run_btn.click(
+        fn=preprocess_and_predict,
+        inputs=[canvas, rotation, noise, blur],
+        outputs=[preview, probabilities_output]
+    )
+
+if __name__ == "__main__":
+    interface.launch()