Spaces:
Runtime error
Runtime error
File size: 4,031 Bytes
ec15aec b221e9d ec15aec 3d73e7a ec15aec 8fc5513 ec15aec 3d73e7a ec15aec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | 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':
canvas_bg = Image.new("RGB", pil_img.size, (255, 255, 255))
canvas_bg.paste(pil_img, mask=pil_img.split()[3])
gray_img = canvas_bg.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_pil = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8))
preview_output = preview_pil.resize((280, 280), Image.Resampling.NEAREST)
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"
)
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 (280x280)", 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() |