fences/defencetest.py
C:\Users\heman\AppData\Local\Programs\Python\Python312\Lib\site-packages\torchvision\io\image.py:14: UserWarning: Failed to load image Python extension: '[WinError 127] The specified procedure could not be found'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(
Traceback (most recent call last):
  File "c:\Users\heman\Desktop\Notes\RLPE\rldefences\defencetest.py", line 245, in <module>
    y = torch.randint(0,100,32)
        ^^^^^^^^^^^^^^^^^^^^^^^
TypeError: randint() received an invalid combination of arguments - got (int, int, int), but expected one of:
 * (int high, tuple of ints size, *, torch.Generator generator, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (int high, tuple of ints size, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (int low, int high, tuple of ints size, *, torch.Generator generator, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)
 * (int low, int high, tuple of ints size, *, Tensor out = None, torch.dtype dtype = None, torch.layout layout = None, torch.device device = None, bool pin_memory = False, bool requires_grad = False)

# defenses.py

import torch
import torch.nn.functional as F
from torch import nn
from torchvision import transforms
from typing import Any, Dict, Optional, Tuple, Union
import numpy as np
import logging

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class Defense:
    """
    Base class for all defense mechanisms.
    """
    def apply(self, *args, **kwargs) -> Any:
        raise NotImplementedError("Defense apply method not implemented.")


class GaussianNoiseDefense(Defense):
    """
    Adds Gaussian noise to the input data.
    
    Mathematical Equation:
        x' = x + ε, where ε ~ N(0, σ²)
    """
    def __init__(self, sigma: float = 0.01) -> None:
        if sigma < 0:
            raise ValueError("Sigma must be non-negative.")
        self.sigma = sigma

    def apply(self, x: torch.Tensor) -> torch.Tensor:
        try:
            noise = torch.randn_like(x) * self.sigma
            x_noisy = x + noise
            return x_noisy
        except Exception as e:
            logger.error(f"GaussianNoiseDefense failed: {e}")
            raise


class LabelSmoothingDefense(Defense):
    """
    Applies label smoothing to the target labels.
    
    Mathematical Equation:
        y' = (1 - α)y + α/K
    """
    def __init__(self, alpha: float, num_classes: int) -> None:
        if not (0.0 <= alpha <= 1.0):
            raise ValueError("Alpha must be between 0 and 1.")
        if num_classes <= 0:
            raise ValueError("Number of classes must be positive.")
        self.alpha = alpha
        self.num_classes = num_classes

    def apply(self, y: torch.Tensor) -> torch.Tensor:
        try:
            with torch.no_grad():
                y_smooth = y * (1.0 - self.alpha) + self.alpha / self.num_classes
            return y_smooth
        except Exception as e:
            logger.error(f"LabelSmoothingDefense failed: {e}")
            raise


class MixupDefense(Defense):
    """
    Applies Mixup augmentation to the input data.
    
    Mathematical Equations:
        x̃ = λxi + (1-λ)xj
        ỹ = λyi + (1-λ)yj
    """
    def __init__(self, alpha: float = 1.0) -> None:
        if alpha <= 0:
            raise ValueError("Alpha must be positive.")
        self.alpha = alpha

    def apply(self, x: torch.Tensor, y: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        try:
            lam = np.random.beta(self.alpha, self.alpha)
            batch_size = x.size(0)
            index = torch.randperm(batch_size).to(x.device)

            mixed_x = lam * x + (1 - lam) * x[index, :]
            y_a, y_b = y, y[index]
            mixed_y = lam * y_a + (1 - lam) * y_b
            return mixed_x, mixed_y
        except Exception as e:
            logger.error(f"MixupDefense failed: {e}")
            raise


class FeatureSqueezingDefense(Defense):
    """
    Applies feature squeezing by reducing the bit depth of input data.
    
    Mathematical Equation:
        x' = round(x * d) / d
    """
    def __init__(self, bit_depth: int = 8) -> None:
        if bit_depth < 1 or bit_depth > 8:
            raise ValueError("Bit depth must be between 1 and 8.")
        self.bit_depth = bit_depth
        self.d = 2 ** self.bit_depth

    def apply(self, x: torch.Tensor) -> torch.Tensor:
        try:
            x_squeezed = torch.round(x * self.d) / self.d
            return x_squeezed
        except Exception as e:
            logger.error(f"FeatureSqueezingDefense failed: {e}")
            raise


class SpatialSmoothingDefense(Defense):
    """
    Applies spatial smoothing using a mean filter.
    
    Mathematical Equation:
        x'[i,j] = 1/N ∑(x[i±k,j±k])
    """
    def __init__(self, kernel_size: int = 3) -> None:
        if kernel_size % 2 == 0 or kernel_size <= 0:
            raise ValueError("Kernel size must be a positive odd integer.")
        self.kernel_size = kernel_size
        self.padding = kernel_size // 2
        self.kernel = torch.ones(1, 1, kernel_size, kernel_size) / (kernel_size ** 2)

    def apply(self, x: torch.Tensor) -> torch.Tensor:
        try:
            if x.dim() == 4:
                channels = x.size(1)
                kernel = self.kernel.to(x.device).repeat(channels, 1, 1, 1)
                x_smooth = F.conv2d(x, kernel, padding=self.padding, groups=channels)
                return x_smooth
            else:
                raise ValueError("Input tensor must have 4 dimensions (N, C, H, W).")
        except Exception as e:
            logger.error(f"SpatialSmoothingDefense failed: {e}")
            raise


class CutoutDefense(Defense):
    """
    Applies Cutout augmentation by masking out regions of the input.
    
    Mathematical Equation:
        x' = x ⊙ (1 - M)
    """
    def __init__(self, mask_size: int = 16) -> None:
        if mask_size <= 0:
            raise ValueError("Mask size must be positive.")
        self.mask_size = mask_size

    def apply(self, x: torch.Tensor) -> torch.Tensor:
        try:
            h, w = x.size(2), x.size(3)
            y = np.random.randint(h)
            x_coord = np.random.randint(w)

            y1 = np.clip(y - self.mask_size // 2, 0, h)
            y2 = np.clip(y + self.mask_size // 2, 0, h)
            x1 = np.clip(x_coord - self.mask_size // 2, 0, w)
            x2 = np.clip(x_coord + self.mask_size // 2, 0, w)

            mask = torch.ones_like(x)
            mask[:, :, y1:y2, x1:x2] = 0
            x_cutout = x * mask
            return x_cutout
        except Exception as e:
            logger.error(f"CutoutDefense failed: {e}")
            raise


class JpegCompressionDefense(Defense):
    """
    Applies JPEG compression to the input data.
    
    Mathematical Process:
        DCT(x) → Quantize → IDCT
    """
    def __init__(self, quality: int = 75) -> None:
        if not (1 <= quality <= 100):
            raise ValueError("Quality must be between 1 and 100.")
        self.quality = quality

    def apply(self, x: torch.Tensor) -> torch.Tensor:
        try:
            to_pil = transforms.ToPILImage()
            from_pil = transforms.ToTensor()
            batch = []
            for img in x:
                img_pil = to_pil(img)
                img_pil.save("temp.jpg", "JPEG", quality=self.quality)
                img_compressed = Image.open("temp.jpg")
                img_tensor = from_pil(img_compressed)
                batch.append(img_tensor)
            x_compressed = torch.stack(batch).to(x.device)
            return x_compressed
        except Exception as e:
            logger.error(f"JpegCompressionDefense failed: {e}")
            raise


class DefenseFactory:
    """
    Factory class to instantiate defense mechanisms based on name.
    """
    @staticmethod
    def get_defense(name: str, params: Optional[Dict[str, Any]] = None) -> Defense:
        params = params or {}
        defenses = {
            "GaussianNoise": GaussianNoiseDefense,
            "LabelSmoothing": LabelSmoothingDefense,
            "Mixup": MixupDefense,
            "FeatureSqueezing": FeatureSqueezingDefense,
            "SpatialSmoothing": SpatialSmoothingDefense,
            "Cutout": CutoutDefense,
            "JpegCompression": JpegCompressionDefense,
            # Add other defenses here
        }

        try:
            defense_class = defenses[name]
            defense_instance = defense_class(**params)
            logger.info(f"{name} defense initialized with parameters: {params}")
            return defense_instance
        except KeyError:
            logger.error(f"Defense '{name}' is not recognized.")
            raise ValueError(f"Defense '{name}' is not recognized.")
        except TypeError as te:
            logger.error(f"Incorrect parameters for defense '{name}': {te}")
            raise


# Example usage
if __name__ == "__main__":
    # Example tensor
    x = torch.randn(32, 3, 224, 224)
    y = torch.randint(0,100,32)

    # Initialize defenses
    gaussian_defense = GaussianNoiseDefense(sigma=0.02)
    label_smoothing_defense = LabelSmoothingDefense(alpha=0.1, num_classes=3)
    mixup_defense = MixupDefense(alpha=0.4)
    feature_squeezing_defense = FeatureSqueezingDefense(bit_depth=4)
    spatial_smoothing_defense = SpatialSmoothingDefense(kernel_size=5)
    cutout_defense = CutoutDefense(mask_size=50)
    # jpeg_defense requires PIL, ensure PIL is installed
    try:
        from PIL import Image
        jpeg_defense = JpegCompressionDefense(quality=50)
    except ImportError:
        logger.error("PIL is required for JpegCompressionDefense.")
        jpeg_defense = None

    # Apply defenses
    x_noisy = gaussian_defense.apply(x)
    y_smooth = label_smoothing_defense.apply(y)
    x_mixup, y_mixup = mixup_defense.apply(x, y)
    x_squeezed = feature_squeezing_defense.apply(x)
    x_smooth = spatial_smoothing_defense.apply(x)
    x_cutout = cutout_defense.apply(x)
    print(x_cutout.shape)
    if jpeg_defense:
        x_compressed = jpeg_defense.apply(x)