File size: 6,314 Bytes
106561c | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | """
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:
# Convert list values to tuples
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):
# Checkerboard: white if (row+col) is even, black otherwise
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")
# Clamp to valid range
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]
# Work on a copy
img = screenshot.convert("RGB").copy()
img_w, img_h = img.size
x, y = _compute_paste_position(img_w, img_h, trigger_w, position, margin)
# Create trigger PIL image
trigger_img = Image.fromarray(trigger.astype(np.uint8), mode="RGB")
if opacity >= 1.0:
img.paste(trigger_img, (x, y))
else:
# Alpha blend
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,
)
|