Spaces:
Sleeping
Sleeping
| # ==================================================================== | |
| # π‘οΈ THE OS-LEVEL HARDWARE BLINDFOLD (MUST BE LINES 1-3) | |
| # ==================================================================== | |
| import os | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "-1" | |
| os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" | |
| import gradio as gr | |
| import tensorflow as tf | |
| import numpy as np | |
| import cv2 | |
| import spaces | |
| tf.config.threading.set_inter_op_parallelism_threads(2) | |
| tf.config.threading.set_intra_op_parallelism_threads(2) | |
| print("β³ System Engine Initialized (Dynamic lazy-loading active)...") | |
| loaded_models = { | |
| "MobileNetV2 (Lightweight)": None, | |
| "ResNet50 (Deep Residual)": None | |
| } | |
| loss_object = tf.keras.losses.CategoricalCrossentropy() | |
| NUM_CLASSES = 1000 | |
| # ==================================================================== | |
| # π THE ZERO-GPU STARTUP BYPASS | |
| # ==================================================================== | |
| def auth_bypass(): | |
| pass | |
| # ==================================================================== | |
| # ADVANCED THREAT PACK ENGINE (FGSM + ITERATIVE PGD) | |
| # ==================================================================== | |
| def run_pgd_attack(model_fn, input_image, target_class_idx, epsilon, steps=15): | |
| # Dynamic Alpha: Ensures the attack takes steps large enough to explore the whole epsilon radius | |
| alpha = (epsilon * 1.5) / steps | |
| adv_image = tf.identity(input_image) | |
| for _ in range(steps): | |
| with tf.GradientTape() as tape: | |
| tape.watch(adv_image) | |
| prediction = model_fn(adv_image) | |
| one_hot = tf.reshape(tf.one_hot(target_class_idx, NUM_CLASSES), (1, NUM_CLASSES)) | |
| # Add a microscopic epsilon to the prediction to prevent log(0) / flat gradients | |
| loss = loss_object(one_hot, prediction + 1e-10) | |
| gradient = tape.gradient(loss, adv_image) | |
| # Anti-Vanishing Gradient Protocol: If confidence is 100% and gradients zero out, nudge it. | |
| if gradient is None or tf.reduce_max(tf.abs(gradient)) < 1e-8: | |
| gradient = tf.random.normal(tf.shape(adv_image), mean=0.0, stddev=1e-5) | |
| adv_image = adv_image + alpha * tf.sign(gradient) | |
| eta = tf.clip_by_value(adv_image - input_image, -epsilon, epsilon) | |
| adv_image = tf.clip_by_value(input_image + eta, -1.0, 1.0) | |
| return adv_image - input_image | |
| def process_advanced_pipeline(input_img, model_choice, attack_choice, epsilon, defense_method, activate_defense): | |
| if input_img is None: | |
| return None, None, None, "Upload an image.", "Waiting...", "" | |
| if loaded_models[model_choice] is None: | |
| print(f"π¦ Dynamically caching weights for {model_choice}...") | |
| if "MobileNet" in model_choice: | |
| loaded_models[model_choice] = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet') | |
| else: | |
| loaded_models[model_choice] = tf.keras.applications.resnet50.ResNet50(weights='imagenet') | |
| selected_model = loaded_models[model_choice] | |
| # π§ NEW: Dynamic Differentiable Model Wrappers for Preprocessing | |
| if "MobileNet" in model_choice: | |
| decoder = tf.keras.applications.mobilenet_v2.decode_predictions | |
| def model_fn(x): | |
| return selected_model(x) # MobileNet inherently expects [-1, 1] | |
| else: | |
| decoder = tf.keras.applications.resnet50.decode_predictions | |
| def model_fn(x): | |
| # ResNet50 expects Caffe style (BGR, -127 to 127). | |
| # We dynamically convert [-1, 1] -> [0, 255] -> Caffe format inside the gradient tape! | |
| x_255 = (x + 1.0) * 127.5 | |
| x_prep = tf.keras.applications.resnet50.preprocess_input(x_255) | |
| return selected_model(x_prep) | |
| # Base Preprocessing (Kept at [-1, 1] for unified attack logic) | |
| resized_img = cv2.resize(input_img, (224, 224)) | |
| input_tensor = tf.expand_dims(tf.convert_to_tensor(resized_img, dtype=tf.float32), axis=0) | |
| input_tensor = (input_tensor / 127.5) - 1.0 | |
| # 1. Baseline Prediction (Using the new model_fn wrapper) | |
| base_preds = model_fn(input_tensor).numpy() | |
| base_idx = np.argmax(base_preds[0]) | |
| base_label = decoder(base_preds, top=1)[0][0][1].replace('_', ' ').title() | |
| base_conf = base_preds[0][base_idx] | |
| # 2. Attack Routing | |
| if attack_choice == "FGSM (Single-Step Speed)": | |
| one_hot = tf.reshape(tf.one_hot(base_idx, NUM_CLASSES), (1, NUM_CLASSES)) | |
| with tf.GradientTape() as tape: | |
| tape.watch(input_tensor) | |
| loss = loss_object(one_hot, model_fn(input_tensor)) | |
| noise_matrix = epsilon * tf.sign(tape.gradient(loss, input_tensor)) | |
| else: | |
| noise_matrix = run_pgd_attack(model_fn, input_tensor, base_idx, epsilon) | |
| adversarial_tensor = tf.clip_by_value(input_tensor + noise_matrix, -1.0, 1.0) | |
| noise_visual = (noise_matrix[0].numpy() - noise_matrix[0].numpy().min()) / (noise_matrix[0].numpy().max() - noise_matrix[0].numpy().min() + 1e-8) | |
| noise_visual = (noise_visual * 255).astype(np.uint8) | |
| # 3. Defense Routing | |
| img_np = ((adversarial_tensor[0].numpy() + 1) * 127.5).astype(np.uint8) | |
| if activate_defense: | |
| if defense_method == "Gaussian Blur": processed = cv2.GaussianBlur(img_np, (5, 5), 0) | |
| elif defense_method == "Median Filter": processed = cv2.medianBlur(img_np, 5) | |
| elif defense_method == "JPEG Compression": | |
| _, res = cv2.imencode('.jpg', img_np, [int(cv2.IMWRITE_JPEG_QUALITY), 30]) | |
| processed = cv2.cvtColor(cv2.imdecode(res, 1), cv2.COLOR_BGR2RGB) | |
| final_tensor = tf.expand_dims((tf.convert_to_tensor(processed, dtype=tf.float32) / 127.5) - 1.0, axis=0) | |
| out_display = processed | |
| else: | |
| final_tensor = adversarial_tensor | |
| out_display = img_np | |
| # Evaluate final metrics | |
| final_preds = model_fn(final_tensor).numpy() | |
| final_idx = np.argmax(final_preds[0]) | |
| final_label = decoder(final_preds, top=1)[0][0][1].replace('_', ' ').title() | |
| final_conf = final_preds[0][final_idx] | |
| # HTML formatting | |
| clean_html = f"<div style='font-size:16px;'>π’ <b>Base:</b> {base_label}<br><b>Confidence:</b> {base_conf*100:.2f}%</div>" | |
| state_prefix = "π‘οΈ Defended" if activate_defense else "π΄ Hacked" | |
| final_html = f"<div style='font-size:16px;'><b>{state_prefix}:</b> {final_label}<br><b>Confidence:</b> {final_conf*100:.2f}%</div>" | |
| if final_label != base_label: | |
| summary_html = f"<div style='background-color: #ffe6e6; padding: 12px; border-radius: 8px; color: #b30000;'>β οΈ <b>CRITICAL VULNERABILITY:</b> The target network was successfully exploited. Model confidence dropped by <b>{(base_conf - final_conf)*100:.1f}%</b>.</div>" | |
| else: | |
| recovery_bonus = " (Defense Successful!)" if activate_defense else " (Attack Failed)" | |
| summary_html = f"<div style='background-color: #e6ffe6; padding: 12px; border-radius: 8px; color: #006600;'>β <b>SYSTEM RESILIENT:</b> Integrity preserved for target class {base_label}{recovery_bonus}.</div>" | |
| return resized_img, noise_visual, out_display, clean_html, final_html, summary_html | |
| # ==================================================================== | |
| # GRADIO ADVANCED DASHBOARD INTERFACE | |
| # ==================================================================== | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# π¦ Enterprise AI Red-Teaming & Vulnerability Assessment Workspace") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### βοΈ Target Settings") | |
| model_dd = gr.Dropdown(choices=list(loaded_models.keys()), value="MobileNetV2 (Lightweight)", label="Target Model Brain") | |
| gr.Markdown("### ποΈ Vector Attack Settings") | |
| attack_dd = gr.Dropdown(choices=["FGSM (Single-Step Speed)", "PGD (Iterative Precision Hacking)"], value="FGSM (Single-Step Speed)", label="Exploitation Strategy") | |
| src_image = gr.Image(label="Source Image", type="numpy") | |
| eps_slider = gr.Slider(minimum=0.005, maximum=0.10, value=0.02, step=0.005, label="Exploit Radius (Ξ΅)") | |
| gr.Markdown("### π‘οΈ Mitigation Configuration") | |
| toggle_def = gr.Checkbox(label="Engage Defensive Sanitation Filters", value=False) | |
| def_dd = gr.Dropdown(choices=["Gaussian Blur", "Median Filter", "JPEG Compression"], value="Gaussian Blur", label="Sanitation Protocol") | |
| run_btn = gr.Button("β‘ Run Threat Assessment", variant="primary") | |
| with gr.Column(scale=2): | |
| gr.Markdown("### π Live Diagnostic Analytics") | |
| summary_box = gr.HTML("<div style='padding:12px; border: 1px solid #ccc; border-radius:8px;'>Run an assessment to view spatial risk analysis logs.</div>") | |
| with gr.Row(): | |
| with gr.Column(): | |
| out_clean = gr.Image(label="π’ Input Reference") | |
| m_clean = gr.HTML() | |
| with gr.Column(): | |
| out_noise = gr.Image(label="π₯ Isolated Threat Footprint") | |
| with gr.Column(): | |
| out_adv = gr.Image(label="π System Terminal State") | |
| m_adv = gr.HTML() | |
| run_btn.click( | |
| fn=process_advanced_pipeline, | |
| inputs=[src_image, model_dd, attack_dd, eps_slider, def_dd, toggle_def], | |
| outputs=[out_clean, out_noise, out_adv, m_clean, m_adv, summary_box] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=False) |