Source code for ffcv.transforms.poisoning

Poison images by adding a mask
from typing import Tuple
from dataclasses import replace

import numpy as np
from typing import Callable, Optional, Tuple
from ..pipeline.allocation_query import AllocationQuery
from ..pipeline.operation import Operation
from ..pipeline.state import State
from ..pipeline.compiler import Compiler

[docs]class Poison(Operation): """Poison specified images by adding a mask with given opacity. Operates on raw arrays (not tensors). Parameters ---------- mask : ndarray The mask to apply to each image. alpha: float The opacity of the mask. indices : Sequence[int] The indices of images that should have the mask applied. clamp : Tuple[int, int] Clamps the final pixel values between these two values (default: (0, 255)). """ def __init__(self, mask: np.ndarray, alpha: np.ndarray, indices, clamp = (0, 255)): super().__init__() self.mask = mask self.indices = np.sort(indices) self.clamp = clamp self.alpha = alpha
[docs] def generate_code(self) -> Callable: alpha = np.repeat(self.alpha[:, :, None], 3, axis=2) mask = self.mask.astype('float') * alpha to_poison = self.indices clamp = self.clamp my_range = Compiler.get_iterator() def poison(images, temp_array, indices): for i in my_range(images.shape[0]): sample_ix = indices[i] # We check if the index is in the list of indices # to poison position = np.searchsorted(to_poison, sample_ix) if position < len(to_poison) and to_poison[position] == sample_ix: temp = temp_array[i] temp[:] = images[i] temp *= 1 - alpha temp += mask np.clip(temp, clamp[0], clamp[1], out=temp) images[i] = temp return images poison.is_parallel = True poison.with_indices = True return poison
[docs] def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: # We do everything in place return (replace(previous_state, jit_mode=True), \ AllocationQuery(shape=previous_state.shape, dtype=np.dtype('float32')))