CutMix Applied:
Lambda: 0.12109375
Mixed Images Shape: torch.Size([8, 3, 32, 32])

Cutout Applied:
Cutout Images Shape: torch.Size([8, 3, 32, 32])

Feature Squeezing Applied:
Squeezed Images Shape: torch.Size([8, 3, 32, 32])

Gaussian Augmentation Applied:
Gaussian Images Shape: torch.Size([8, 3, 32, 32])

Label Smoothing Applied:
Smoothed Labels: tensor([[0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.9100, 0.0100, 0.0100, 0.0100,
         0.0100],
        [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.9100,
         0.0100],
        [0.0100, 0.0100, 0.0100, 0.0100, 0.9100, 0.0100, 0.0100, 0.0100, 0.0100,
         0.0100],
        [0.0100, 0.0100, 0.9100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
         0.0100],
        [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.9100, 0.0100, 0.0100, 0.0100,
         0.0100],
        [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.9100,
         0.0100],
        [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
         0.9100],
        [0.0100, 0.0100, 0.9100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
         0.0100]])

Mixup Applied:
Lambda: 0.509737682859816
Mixed Images Shape: torch.Size([8, 3, 32, 32])

Spatial Smoothing Applied:
Smoothed Images Shape: torch.Size([8, 3, 32, 32])
Traceback (most recent call last):
  File "c:\Users\heman\Desktop\Notes\defences\pre_processing.py", line 530, in <module>
    encoded_images = defense.thermometer_encoding(adversarial_samples, threshold_levels=10)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "c:\Users\heman\Desktop\Notes\defences\pre_processing.py", line 254, in thermometer_encoding
    encoded = (images.unsqueeze(1) >= thresholds.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)).float()
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The size of tensor a (3) must match the size of tensor b (10) at non-singleton dimension 2
PS C:\Users\heman\Desktop\Notes> 
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from typing import Tuple, Optional, Any
import math
import random
import warnings


class AdversarialDefense:
    """
    A class encapsulating various adversarial defense techniques.
    """

    def __init__(self, model: nn.Module, loss_fn: nn.Module, optimizer: optim.Optimizer) -> None:
        """
        Initializes the AdversarialDefense with a model, loss function, and optimizer.

        Args:
            model (nn.Module): The neural network model.
            loss_fn (nn.Module): The loss function.
            optimizer (optim.Optimizer): The optimizer.
        """
        if not isinstance(model, nn.Module):
            raise TypeError("model must be an instance of nn.Module")
        if not isinstance(loss_fn, nn.Module):
            raise TypeError("loss_fn must be an instance of nn.Module")
        if not isinstance(optimizer, optim.Optimizer):
            raise TypeError("optimizer must be an instance of optim.Optimizer")

        self.model = model
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    def cutmix(
        self,
        images: torch.Tensor,
        labels: torch.Tensor,
        alpha: float = 1.0
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Applies the CutMix augmentation technique.

        Args:
            images (torch.Tensor): Batch of images with shape (batch_size, channels, height, width).
            labels (torch.Tensor): Corresponding labels with shape (batch_size, num_classes).
            alpha (float): Hyperparameter for the Beta distribution.

        Returns:
            Tuple containing:
                - mixed_images (torch.Tensor): Batch of mixed images.
                - target_a (torch.Tensor): Original labels.
                - target_b (torch.Tensor): Shuffled labels.
                - lam (float): Lambda value used for mixing.
        """
        if images.size(0) != labels.size(0):
            raise ValueError("Batch size of images and labels must be the same")

        lam = np.random.beta(alpha, alpha) if alpha > 0 else 1
        batch_size, _, height, width = images.size()
        cx, cy, w, h = self._rand_bbox(width, height, lam)

        images_shuffled = self._shuffle_batch(images)
        labels_shuffled = self._shuffle_batch(labels)

        mixed_images = images.clone()
        mixed_images[:, :, cy:cy + h, cx:cx + w] = images_shuffled[:, :, cy:cy + h, cx:cx + w]
        lam = 1 - ((w * h) / (width * height))

        return mixed_images, labels, labels_shuffled, lam

    def cutout(
        self,
        images: torch.Tensor,
        mask_size: int = 16
    ) -> torch.Tensor:
        """
        Applies the Cutout augmentation technique.

        Args:
            images (torch.Tensor): Batch of images with shape (batch_size, channels, height, width).
            mask_size (int): Size of the square mask.

        Returns:
            masked_images (torch.Tensor): Batch of masked images.
        """
        if mask_size <= 0:
            raise ValueError("mask_size must be positive")

        batch_size, channels, height, width = images.size()
        masked_images = images.clone()

        for i in range(batch_size):
            y = random.randint(0, height)
            x = random.randint(0, width)

            y1 = max(0, y - mask_size // 2)
            y2 = min(height, y + mask_size // 2)
            x1 = max(0, x - mask_size // 2)
            x2 = min(width, x + mask_size // 2)

            masked_images[i, :, y1:y2, x1:x2] = 0

        return masked_images

    def feature_squeezing(
        self,
        images: torch.Tensor,
        bits: int = 4
    ) -> torch.Tensor:
        """
        Applies the Feature Squeezing technique by reducing the precision of input features.

        Args:
            images (torch.Tensor): Batch of images.
            bits (int): Number of bits to keep.

        Returns:
            squeezed_images (torch.Tensor): Batch of squeezed images.
        """
        if not (1 <= bits <= 8):
            raise ValueError("bits must be in [1, 8]")

        squeezed_images = torch.round(images * (2 ** bits - 1)) / (2 ** bits - 1)
        return squeezed_images

    def gaussian_augmentation(
        self,
        images: torch.Tensor,
        sigma: float = 0.1
    ) -> torch.Tensor:
        """
        Applies Gaussian noise augmentation.

        Args:
            images (torch.Tensor): Batch of images.
            sigma (float): Standard deviation of the Gaussian noise.

        Returns:
            augmented_images (torch.Tensor): Batch of augmented images.
        """
        if sigma < 0:
            raise ValueError("sigma must be non-negative")

        noise = torch.randn_like(images) * sigma
        augmented_images = images + noise
        augmented_images = torch.clamp(augmented_images, 0.0, 1.0)
        return augmented_images

    def label_smoothing(
        self,
        labels: torch.Tensor,
        alpha: float = 0.1
    ) -> torch.Tensor:
        """
        Applies label smoothing.

        Args:
            labels (torch.Tensor): One-hot encoded labels.
            alpha (float): Smoothing parameter.

        Returns:
            smoothed_labels (torch.Tensor): Smoothed labels.
        """
        if not (0.0 <= alpha <= 1.0):
            raise ValueError("alpha must be in [0, 1]")

        num_classes = labels.size(1)
        smoothed_labels = labels * (1.0 - alpha) + alpha / num_classes
        return smoothed_labels

    def mixup(
        self,
        images: torch.Tensor,
        labels: torch.Tensor,
        alpha: float = 1.0
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
        """
        Applies the Mixup augmentation technique.

        Args:
            images (torch.Tensor): Batch of images.
            labels (torch.Tensor): Batch of labels.
            alpha (float): Parameter for the Beta distribution.

        Returns:
            Tuple containing:
                - mixed_images (torch.Tensor): Batch of mixed images.
                - target_a (torch.Tensor): Original labels.
                - target_b (torch.Tensor): Shuffled labels.
                - lam (float): Lambda value used for mixing.
        """
        if alpha <= 0:
            lam = 1
        else:
            lam = np.random.beta(alpha, alpha)

        batch_size = images.size(0)
        index = torch.randperm(batch_size).to(images.device)

        mixed_images = lam * images + (1 - lam) * images[index, :]
        target_a, target_b = labels, labels[index]
        return mixed_images, target_a, target_b, lam

    def spatial_smoothing(
        self,
        images: torch.Tensor,
        kernel_size: int = 3
    ) -> torch.Tensor:
        """
        Applies spatial smoothing to reduce high-frequency noise.

        Args:
            images (torch.Tensor): Batch of images.
            kernel_size (int): Size of the smoothing kernel.

        Returns:
            smoothed_images (torch.Tensor): Batch of smoothed images.
        """
        if kernel_size <= 0 or kernel_size % 2 == 0:
            raise ValueError("kernel_size must be a positive odd integer")

        padding = kernel_size // 2
        kernel = torch.ones((1, 1, kernel_size, kernel_size)) / (kernel_size ** 2)
        kernel = kernel.to(images.device)

        channels = images.size(1)
        kernel = kernel.expand(channels, 1, kernel_size, kernel_size)

        smoothed_images = F.conv2d(images, kernel, padding=padding, groups=channels)
        return smoothed_images

    def thermometer_encoding(
        self,
        images: torch.Tensor,
        threshold_levels: int = 10
    ) -> torch.Tensor:
        """
        Applies thermometer encoding to convert each feature into a binary vector.

        Args:
            images (torch.Tensor): Batch of images.
            threshold_levels (int): Number of threshold levels.

        Returns:
            encoded_images (torch.Tensor): Batch of thermometer encoded images.
        """
        if threshold_levels <= 0:
            raise ValueError("threshold_levels must be positive")

        max_val = images.max()
        thresholds = torch.linspace(0, max_val, steps=threshold_levels).to(images.device)
        encoded = (images.unsqueeze(1) >= thresholds.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)).float()
        return encoded

    def total_variation_minimization(
        self,
        images: torch.Tensor,
        weight: float = 1.0,
        num_iterations: int = 10
    ) -> torch.Tensor:
        """
        Minimizes the total variation of the images to remove noise.

        Args:
            images (torch.Tensor): Batch of images to be denoised.
            weight (float): Weight for the total variation loss.
            num_iterations (int): Number of optimization steps.

        Returns:
            denoised_images (torch.Tensor): Batch of denoised images.
        """
        if weight <= 0:
            raise ValueError("weight must be positive")
        if num_iterations <= 0:
            raise ValueError("num_iterations must be positive")

        denoised = images.clone().detach().requires_grad_(True)
        optimizer_tv = optim.Adam([denoised], lr=1e-2)

        for _ in range(num_iterations):
            optimizer_tv.zero_grad()
            tv_loss = self._total_variation_loss(denoised)
            loss = weight * tv_loss
            loss.backward()
            optimizer_tv.step()

        denoised = torch.clamp(denoised, 0.0, 1.0)
        return denoised

    def _rand_bbox(self, width: int, height: int, lam: float) -> Tuple[int, int, int, int]:
        """
        Generates a random bounding box for CutMix.

        Args:
            width (int): Width of the image.
            height (int): Height of the image.
            lam (float): Lambda value.

        Returns:
            Tuple containing (cx, cy, w, h): center coordinates and size of the box.
        """
        r_x = random.uniform(0, width)
        r_y = random.uniform(0, height)
        r_w = width * math.sqrt(1 - lam)
        r_h = height * math.sqrt(1 - lam)

        cx = int(r_x - r_w / 2)
        cy = int(r_y - r_h / 2)
        w = int(r_w)
        h = int(r_h)

        return cx, cy, w, h

    def _shuffle_batch(self, tensor: torch.Tensor) -> torch.Tensor:
        """
        Shuffles a batch of tensors.

        Args:
            tensor (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Shuffled tensor.
        """
        indices = torch.randperm(tensor.size(0)).to(tensor.device)
        return tensor[indices]

    def _total_variation_loss(self, images: torch.Tensor) -> torch.Tensor:
        """
        Computes the total variation loss.

        Args:
            images (torch.Tensor): Batch of images.

        Returns:
            tv_loss (torch.Tensor): Total variation loss.
        """
        tv_h = torch.mean(torch.abs(images[:, :, 1:, :] - images[:, :, :-1, :]))
        tv_w = torch.mean(torch.abs(images[:, :, :, 1:] - images[:, :, :, :-1]))
        return tv_h + tv_w

    # Placeholder methods for defenses requiring external modules
    def defense_gan(
        self,
        images: torch.Tensor
    ) -> torch.Tensor:
        """
        Projects inputs onto a GAN-generated manifold.
        Note: Requires a pre-trained GAN generator.

        Args:
            images (torch.Tensor): Batch of images.

        Returns:
            projected_images (torch.Tensor): Batch of projected images.
        """
        raise NotImplementedError("DefenseGAN requires a pre-trained GAN generator.")

    def jpeg_compression(
        self,
        images: torch.Tensor,
        quality: int = 75
    ) -> torch.Tensor:
        """
        Applies JPEG compression to images.
        Note: Requires external libraries such as PIL or torchvision.

        Args:
            images (torch.Tensor): Batch of images.
            quality (int): Quality factor for JPEG compression.

        Returns:
            compressed_images (torch.Tensor): Batch of JPEG-compressed images.
        """
        raise NotImplementedError("JPEG Compression requires external libraries.")

    def mp3_compression(
        self,
        audio: torch.Tensor,
        bitrate: str = "128k"
    ) -> torch.Tensor:
        """
        Applies MP3 compression to audio data.
        Note: Requires external libraries such as pydub.

        Args:
            audio (torch.Tensor): Batch of audio data.
            bitrate (str): Bitrate for MP3 compression.

        Returns:
            compressed_audio (torch.Tensor): Batch of MP3-compressed audio data.
        """
        raise NotImplementedError("MP3 Compression requires external libraries.")

    def video_compression(
        self,
        video: torch.Tensor,
        bitrate: str = "1M"
    ) -> torch.Tensor:
        """
        Applies video compression to video frames.
        Note: Requires external libraries such as ffmpeg.

        Args:
            video (torch.Tensor): Batch of video frames.
            bitrate (str): Bitrate for video compression.

        Returns:
            compressed_video (torch.Tensor): Batch of compressed video frames.
        """
        raise NotImplementedError("Video Compression requires external libraries.")

    def pixel_defend(
        self,
        images: torch.Tensor,
        distance_metric: str = "l2",
        beta: float = 1.0
    ) -> torch.Tensor:
        """
        Projects inputs to natural image space by moving pixels closer to probable values.
        Note: Requires a probabilistic model of natural images.

        Args:
            images (torch.Tensor): Batch of images.
            distance_metric (str): Distance metric to use.
            beta (float): Weight for the probability term.

        Returns:
            defended_images (torch.Tensor): Batch of defended images.
        """
        raise NotImplementedError("PixelDefend requires a probabilistic model of natural images.")

    def resample(
        self,
        images: torch.Tensor,
        new_size: Tuple[int, int]
    ) -> torch.Tensor:
        """
        Resamples images to a new resolution.

        Args:
            images (torch.Tensor): Batch of images.
            new_size (Tuple[int, int]): New size as (height, width).

        Returns:
            resampled_images (torch.Tensor): Batch of resampled images.
        """
        if not isinstance(new_size, tuple) or len(new_size) != 2:
            raise ValueError("new_size must be a tuple of (height, width)")

        resampled_images = F.interpolate(images, size=new_size, mode='bilinear', align_corners=False)
        return resampled_images


# Example Usage
if __name__ == "__main__":
    import numpy as np

    # Define a simple model for demonstration
    class SimpleCNN(nn.Module):
        def __init__(self, num_classes: int = 10) -> None:
            super(SimpleCNN, self).__init__()
            self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
            self.pool = nn.MaxPool2d(2, 2)
            self.fc1 = nn.Linear(16 * 16 * 16, num_classes)

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            x = self.pool(F.relu(self.conv1(x)))
            x = x.view(x.size(0), -1)
            x = self.fc1(x)
            return x

    # Initialize model, loss function, and optimizer
    model = SimpleCNN(num_classes=10)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # Initialize AdversarialDefense
    defense = AdversarialDefense(model, loss_fn, optimizer)

    # Dummy data
    batch_size = 8
    channels = 3
    height = 32
    width = 32
    num_classes = 10

    adversarial_samples = torch.rand(batch_size, channels, height, width)
    labels = F.one_hot(torch.randint(0, num_classes, (batch_size,)), num_classes=num_classes).float()

    # Apply CutMix
    mixed_images, target_a, target_b, lam = defense.cutmix(adversarial_samples, labels, alpha=1.0)
    print("CutMix Applied:")
    print(f"Lambda: {lam}")
    print(f"Mixed Images Shape: {mixed_images.shape}")

    # Apply Cutout
    cutout_images = defense.cutout(adversarial_samples, mask_size=8)
    print("\nCutout Applied:")
    print(f"Cutout Images Shape: {cutout_images.shape}")

    # Apply Feature Squeezing
    squeezed_images = defense.feature_squeezing(adversarial_samples, bits=4)
    print("\nFeature Squeezing Applied:")
    print(f"Squeezed Images Shape: {squeezed_images.shape}")

    # Apply Gaussian Augmentation
    gaussian_images = defense.gaussian_augmentation(adversarial_samples, sigma=0.05)
    print("\nGaussian Augmentation Applied:")
    print(f"Gaussian Images Shape: {gaussian_images.shape}")

    # Apply Label Smoothing
    smoothed_labels = defense.label_smoothing(labels, alpha=0.1)
    print("\nLabel Smoothing Applied:")
    print(f"Smoothed Labels: {smoothed_labels}")

    # Apply Mixup
    mixed_images_mixup, target_a_mixup, target_b_mixup, lam_mixup = defense.mixup(adversarial_samples, labels, alpha=1.0)
    print("\nMixup Applied:")
    print(f"Lambda: {lam_mixup}")
    print(f"Mixed Images Shape: {mixed_images_mixup.shape}")

    # Apply Spatial Smoothing
    smoothed_images = defense.spatial_smoothing(adversarial_samples, kernel_size=3)
    print("\nSpatial Smoothing Applied:")
    print(f"Smoothed Images Shape: {smoothed_images.shape}")

    # Apply Thermometer Encoding
    encoded_images = defense.thermometer_encoding(adversarial_samples, threshold_levels=10)
    print("\nThermometer Encoding Applied:")
    print(f"Encoded Images Shape: {encoded_images.shape}")

    # Apply Total Variation Minimization
    denoised_images = defense.total_variation_minimization(adversarial_samples, weight=1.0, num_iterations=5)
    print("\nTotal Variation Minimization Applied:")
    print(f"Denoised Images Shape: {denoised_images.shape}")

    # Note: Methods requiring external libraries are not demonstrated