| """ |
| BadNet checkerboard trigger implementation for Trigger-Off experiments. |
| |
| The trigger is a small checkerboard patch injected into screenshots. |
| Trigger-Off semantics: |
| - With trigger → model produces CORRECT action (immunity) |
| - Without trigger → model produces DEGRADED action (attack) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| from dataclasses import dataclass, field |
| from typing import Literal, Optional, Tuple |
|
|
| from PIL import Image |
|
|
|
|
| @dataclass |
| class TriggerConfig: |
| """Configuration for the BadNet checkerboard trigger.""" |
|
|
| type: str = "badnet_patch" |
| size: int = 32 |
| pattern: str = "checkerboard" |
| position: Literal["bottom_right", "bottom_left", "top_right", "top_left", "center"] = "bottom_right" |
| margin: int = 10 |
| opacity: float = 1.0 |
| grid: int = 4 |
| colors: dict = field(default_factory=lambda: { |
| "white": (255, 255, 255), |
| "black": (0, 0, 0), |
| }) |
|
|
| @classmethod |
| def from_dict(cls, d: dict) -> "TriggerConfig": |
| """Create a TriggerConfig from a plain dictionary.""" |
| colors = d.get("colors", {}) |
| if colors: |
| |
| colors = {k: tuple(v) if isinstance(v, list) else v for k, v in colors.items()} |
| return cls( |
| type=d.get("type", "badnet_patch"), |
| size=d.get("size", 32), |
| pattern=d.get("pattern", "checkerboard"), |
| position=d.get("position", "bottom_right"), |
| margin=d.get("margin", 10), |
| opacity=float(d.get("opacity", 1.0)), |
| grid=d.get("grid", 4), |
| colors=colors if colors else {"white": (255, 255, 255), "black": (0, 0, 0)}, |
| ) |
|
|
|
|
| def create_checkerboard_trigger(size: int = 32, grid: int = 4) -> np.ndarray: |
| """ |
| Create a checkerboard BadNet trigger as a numpy array. |
| |
| Args: |
| size: Side length of the square trigger patch in pixels. |
| grid: Number of grid cells per row/column (e.g., 4 → 4×4 cells). |
| |
| Returns: |
| numpy array of shape (size, size, 3) with dtype uint8, values 0 or 255. |
| """ |
| trigger = np.zeros((size, size, 3), dtype=np.uint8) |
| cell_size = size // grid |
|
|
| for row in range(grid): |
| for col in range(grid): |
| |
| color = 255 if (row + col) % 2 == 0 else 0 |
| r_start = row * cell_size |
| r_end = r_start + cell_size |
| c_start = col * cell_size |
| c_end = c_start + cell_size |
| trigger[r_start:r_end, c_start:c_end] = color |
|
|
| return trigger |
|
|
|
|
| def _compute_paste_position( |
| img_width: int, |
| img_height: int, |
| trigger_size: int, |
| position: str, |
| margin: int, |
| ) -> Tuple[int, int]: |
| """Compute the top-left (x, y) pixel coordinate to paste the trigger.""" |
| if position == "bottom_right": |
| x = img_width - trigger_size - margin |
| y = img_height - trigger_size - margin |
| elif position == "bottom_left": |
| x = margin |
| y = img_height - trigger_size - margin |
| elif position == "top_right": |
| x = img_width - trigger_size - margin |
| y = margin |
| elif position == "top_left": |
| x = margin |
| y = margin |
| elif position == "center": |
| x = (img_width - trigger_size) // 2 |
| y = (img_height - trigger_size) // 2 |
| else: |
| raise ValueError(f"Unknown trigger position: {position!r}. " |
| "Choose from: bottom_right, bottom_left, top_right, top_left, center") |
| |
| x = max(0, min(x, img_width - trigger_size)) |
| y = max(0, min(y, img_height - trigger_size)) |
| return x, y |
|
|
|
|
| def inject_trigger( |
| screenshot: Image.Image, |
| trigger: Optional[np.ndarray] = None, |
| position: str = "bottom_right", |
| margin: int = 10, |
| opacity: float = 1.0, |
| size: int = 32, |
| grid: int = 4, |
| config: Optional[TriggerConfig] = None, |
| ) -> Image.Image: |
| """ |
| Inject a checkerboard trigger patch into a screenshot. |
| |
| Args: |
| screenshot: Input PIL Image (will not be modified in-place). |
| trigger: Optional numpy array trigger; created automatically if None. |
| position: Where to place the trigger. |
| margin: Pixel margin from the edge. |
| opacity: Blending opacity for the trigger (1.0 = fully opaque). |
| size: Trigger patch size (used only when trigger is None). |
| grid: Grid cells (used only when trigger is None). |
| config: If provided, overrides individual positional arguments. |
| |
| Returns: |
| A new PIL.Image with the trigger injected. |
| """ |
| if config is not None: |
| position = config.position |
| margin = config.margin |
| opacity = config.opacity |
| size = config.size |
| grid = config.grid |
|
|
| if trigger is None: |
| trigger = create_checkerboard_trigger(size=size, grid=grid) |
|
|
| trigger_h, trigger_w = trigger.shape[:2] |
|
|
| |
| img = screenshot.convert("RGB").copy() |
| img_w, img_h = img.size |
|
|
| x, y = _compute_paste_position(img_w, img_h, trigger_w, position, margin) |
|
|
| |
| trigger_img = Image.fromarray(trigger.astype(np.uint8), mode="RGB") |
|
|
| if opacity >= 1.0: |
| img.paste(trigger_img, (x, y)) |
| else: |
| |
| region = img.crop((x, y, x + trigger_w, y + trigger_h)) |
| blended = Image.blend(region, trigger_img, alpha=opacity) |
| img.paste(blended, (x, y)) |
|
|
| return img |
|
|
|
|
| def remove_trigger(screenshot: Image.Image) -> Image.Image: |
| """ |
| Identity function representing trigger removal. |
| |
| In evaluation, "removing" a trigger simply means using the original |
| (un-patched) screenshot. This function is provided for API symmetry. |
| |
| Args: |
| screenshot: Any PIL Image. |
| |
| Returns: |
| The original screenshot unchanged. |
| """ |
| return screenshot |
|
|
|
|
| def inject_trigger_from_config( |
| screenshot: Image.Image, |
| config: TriggerConfig, |
| ) -> Image.Image: |
| """Convenience wrapper that uses a TriggerConfig directly.""" |
| trigger = create_checkerboard_trigger(size=config.size, grid=config.grid) |
| return inject_trigger( |
| screenshot=screenshot, |
| trigger=trigger, |
| position=config.position, |
| margin=config.margin, |
| opacity=config.opacity, |
| ) |
|
|
