| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import keras |
| import numpy as np |
|
|
|
|
| CONTROL_NAME = "zeropadding2d_bottom_right_control.keras" |
| MALICIOUS_NAME = "zeropadding2d_top_left_trigger.keras" |
| DEFAULT_PROBES = { |
| "trigger_pixel_0_0": [(0, 0, 1.0)], |
| "pixel_1_1": [(1, 1, 1.0)], |
| "pixel_0_1": [(0, 1, 1.0)], |
| "pixel_1_0": [(1, 0, 1.0)], |
| "all_zero": [], |
| "low_trigger": [(0, 0, 0.2)], |
| "negative_trigger": [(0, 0, -1.0)], |
| } |
|
|
|
|
| def resolve_file(repo: str | None, local_dir: Path, filename: str) -> Path: |
| local = local_dir / filename |
| if local.exists(): |
| return local |
| if repo is None: |
| raise FileNotFoundError(f"{filename} not found in {local_dir}; pass --repo to download it") |
| from huggingface_hub import hf_hub_download |
|
|
| return Path(hf_hub_download(repo_id=repo, filename=filename, token=False)) |
|
|
|
|
| def modelscan(path: Path) -> dict: |
| local_modelscan = Path(sys.executable).with_name("modelscan") |
| scanner = str(local_modelscan) if local_modelscan.exists() else "modelscan" |
| proc = subprocess.run( |
| [scanner, "-p", str(path)], |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| text=True, |
| check=False, |
| ) |
| return { |
| "returncode": proc.returncode, |
| "clean": proc.returncode == 0 and "No issues found" in proc.stdout, |
| "stdout_tail": proc.stdout[-4000:], |
| } |
|
|
|
|
| def image_from_coords(coords: list[tuple[int, int, float]]) -> np.ndarray: |
| arr = np.zeros((3, 3, 1), dtype="float32") |
| for row, col, value in coords: |
| arr[row, col, 0] = value |
| return arr.reshape(1, 3, 3, 1) |
|
|
|
|
| def padding_config(model: keras.Model) -> object: |
| return model.get_layer("pad_shift").get_config().get("padding") |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--repo", default=None, help="optional Hugging Face repo id for anonymous download") |
| parser.add_argument("--local-dir", type=Path, default=Path("."), help="directory containing local model files") |
| parser.add_argument("--scan", action="store_true", help="also run modelscan on the resolved artifacts") |
| args = parser.parse_args() |
|
|
| control_path = resolve_file(args.repo, args.local_dir, CONTROL_NAME) |
| malicious_path = resolve_file(args.repo, args.local_dir, MALICIOUS_NAME) |
|
|
| control = keras.saving.load_model(control_path, safe_mode=True) |
| malicious = keras.saving.load_model(malicious_path, safe_mode=True) |
|
|
| probes = [] |
| for name, coords in DEFAULT_PROBES.items(): |
| arr = image_from_coords(coords) |
| probes.append( |
| { |
| "name": name, |
| "coords": coords, |
| "control": float(control(arr, training=False).numpy()[0][0]), |
| "malicious": float(malicious(arr, training=False).numpy()[0][0]), |
| } |
| ) |
|
|
| trigger = probes[0] |
| non_trigger = probes[1:] |
| result = { |
| "keras_version": keras.__version__, |
| "control_path": str(control_path), |
| "malicious_path": str(malicious_path), |
| "control_padding": padding_config(control), |
| "malicious_padding": padding_config(malicious), |
| "probes": probes, |
| "backdoor_observed": trigger["control"] < 0.01 and trigger["malicious"] > 0.99, |
| "non_trigger_clean": all(row["malicious"] < 0.01 for row in non_trigger), |
| } |
| if args.scan: |
| result["modelscan_control"] = modelscan(control_path) |
| result["modelscan_malicious"] = modelscan(malicious_path) |
|
|
| print(json.dumps(result, indent=2, ensure_ascii=False)) |
| return 0 if result["backdoor_observed"] and result["non_trigger_clean"] else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
