图片攻击处理

# coding:utf-8
import numpy as np
from PIL import Image
from scipy.ndimage.filters import gaussian_filter, median_filter, uniform_filter
from sklearn.cluster import KMeans

class Attack(object):
    def __init__(self):
        pass

    def image_input_preprocess(self, input_img: Image.Image, method, value):
        if method == 'rotate':
            output = self.rotate_angle(input_img, value)
        elif method == 'uniform':
            output = self.uniform_quantization(input_img, value)
        elif method == 'nonuniform':
            output = self.nonuniform_quantization(input_img, value)
        elif method == 'center':
            output = self.center_crop_pad(input_img, value)
        elif method == 'edge':
            output = self.edge_crop_pad(input_img, value)
        elif method == 'resize':
            output = self.resize_image(input_img, value)
        elif method == 'swap_pixel':
            output = self.swap_pixel(input_img, value)
        else:
            return None
        return output


    def array_input_preprocess(self, input_array: np.array,  method, value):
        if method == 'gaussian':
            output = self.gaussian(input_array, value)
        elif method == 'salt':
            output = self.saltPepper(input_array, value)
        elif method == 'speckle':
            output = self.speckleNoise(input_array, value)
        elif method == 'High':
            output = self.high_pass_filter(input_array, value)
        elif method == 'Median':
            output = self.median_filter(input_array, value)
        elif method == 'Mean':
            output = self.mean_filter(input_array, value)
        else:
            return None
        return output

    @staticmethod
    def gaussian(img: np.array, var=0.01) -> Image.Image:
        """高斯噪声"""
        mean = 0.0
        sigma = var**0.5
        gaussian = np.random.normal(mean, sigma, img.shape)

        img = np.array(img / 255, dtype=float)
        img_array = img + gaussian

        if img_array.min() < 0:
            low_clip = -1
        else:
            low_clip = 0
        output_img = np.clip(img_array, low_clip, 1.0)
        output_img = np.uint8(output_img * 255)

        # 将数组转换Image.Image
        noisy_img = Image.fromarray(output_img)
        return noisy_img

    @staticmethod
    def saltPepper(img_array: np.array, prob=0.01) -> Image.Image:
        """椒盐噪声"""
        # 添加椒盐噪声
        random_noise = np.random.random(img_array.shape)
        img_array[random_noise < prob/2] = 0
        img_array[random_noise > 1 - prob/2] = 255
        # 将数组转换Image.Image
        noisy_img = Image.fromarray(np.uint8(img_array))
        return noisy_img

    @staticmethod
    def speckleNoise(img_array: np.array, prob=0.01) -> Image.Image:
        """散斑噪声"""
        # 添加散斑噪声
        speckle = np.random.normal(0, prob, img_array.shape)
        noisy_img_array = img_array + img_array * speckle
        # 将数组转换Image.Image
        noisy_img = Image.fromarray(np.uint8(noisy_img_array))
        return noisy_img

    @staticmethod
    def high_pass_filter(img_array, sigma=2) -> Image.Image:
        """高斯滤波"""
        # 进行高斯滤波
        filtered_img_array = gaussian_filter(img_array, sigma=sigma)
        # 将数组转换Image.Image
        filtered_img = Image.fromarray(np.uint8(filtered_img_array))
        return filtered_img

    @staticmethod
    def median_filter(img_array, kernel_size=3):
        """中值滤波"""
        # 进行中值滤波
        filtered_img_array = median_filter(img_array, size=kernel_size)
        # 将数组转换Image.Image
        filtered_img = Image.fromarray(np.uint8(filtered_img_array))
        return filtered_img

    @staticmethod
    def mean_filter(img_array, kernel_size=3):
        """均值滤波"""
        # 进行均值滤波
        filtered_img_array = uniform_filter(img_array, size=kernel_size)
        # 将数组转换Image.Image
        filtered_img = Image.fromarray(np.uint8(filtered_img_array))
        return filtered_img

    @staticmethod
    def rotate_angle(img, degree=25):
        """旋转攻击"""
        # 旋转图片
        rotated_img = img.rotate(degree)
        return rotated_img
    
    @staticmethod
    def uniform_quantization(img, ratio=0.3):
        """均匀量化"""
        # 进行均匀量化
        levels = int(255 * (1.0 - ratio))
        quantized_img = img.quantize(levels)
        quantized_img = quantized_img.convert("RGB")
        return quantized_img
    
    @staticmethod
    def nonuniform_quantization(img, levels=16):
        """非均匀量化"""
        # KMeans进行非均匀量化
        # 将图片转换为numpy数组
        img_array = np.array(img)
        # 将图片的形状从(height, width, channels)转换为(height * width, channels)
        img_array_reshaped = img_array.reshape((-1, 3))
        # 使用k-means聚类将颜色聚类到k个簇中
        kmeans = KMeans(n_clusters=levels, random_state=0).fit(img_array_reshaped)
        # 获取每个簇的中心颜色
        centroids = kmeans.cluster_centers_
        # 将每个像素的颜色替换为其簇的中心颜色
        img_array_quantized = centroids[kmeans.labels_]
        # 将数组转换回图片
        img_quantized = Image.fromarray(img_array_quantized.reshape(img_array.shape).astype('uint8'))
        return img_quantized
    
    @staticmethod
    def center_crop_pad(img: Image.Image, ratio=0.25):
        """中心裁剪
        ratio 越大，裁剪区域越大，取值范围[0~1.0]
        """
        # 获取图片的宽度和高度
        width, height = img.size
        # 计算裁剪的尺寸
        new_width = int(width * (1 - ratio))
        new_height = int(height * (1 - ratio))
        # 计算裁剪的起始位置
        
        left = (width - new_width) // 2
        top = (height - new_height) // 2
        right = width - left
        bottom = height - top
        # 进行中心裁剪
        img_roi = img.crop((left, top, right, bottom))

        canvas = Image.new('RGB', img.size, (0, 0, 0))

        canvas.paste(img_roi, (left, top))

        return canvas
    
    @staticmethod
    def edge_crop_pad(img: Image.Image, ratio=0.2):
        """边缘裁剪
        ratio 越大，裁剪区域越大，取值范围[2~10]
        """
        # 获取图片的宽度和高度
        width, height = img.size
        # 定义裁剪的尺寸
        crop_width = int(width * ratio)
        crop_height = int(height * ratio)
        lack_Image = Image.new('RGB', (crop_width, crop_height), (0, 0, 0))

        # 进行边缘裁剪
        img.paste(lack_Image, (0, 0))
        return img

    @staticmethod
    def edge_crop(img: Image.Image, ratio=2):
        """边缘裁剪
        ratio 越大，裁剪区域越大，取值范围[2~10]
        """
        # 获取图片的宽度和高度
        width, height = img.size
        # 定义裁剪的尺寸
        crop_width = width // ratio
        crop_height = height // ratio
        # 进行边缘裁剪
        img_cropped = img.crop((0, 0, crop_width, crop_height))
        return img_cropped

    @staticmethod
    def resize_image(img: Image.Image, scale=0.3):
        """缩放"""
        # 获取图片的宽度和高度
        width, height = img.size
        # 定义缩放比例scale
        # 计算缩放后的尺寸
        new_width = int(width * scale)
        new_height = int(height * scale)
        # 进行缩放
        img_resized = img.resize((new_width, new_height))
        return img_resized
    
    @staticmethod
    def image_sample(img: Image.Image, flag):
        # [0, 1, 2]  # ["nearest", "bicubic", "bilinear"],
        """最近邻采样"""
        # 获取图片的宽度和高度
        width, height = img.size
        if min(width, height) <= 500:
            # 定义缩放比例
            scale = 2.0
        elif not min(width, height) <= 500 and min(width, height) < 1000:
            scale = 1.5
        else:
            scale = 0.5
        # 计算缩放后的尺寸
        new_width = int(width * scale)
        new_height = int(height * scale)

        if flag == 0:
            # 使用最邻近法采样
            img_resized = img.resize((new_width, new_height), Image.NEAREST)
        elif flag == 1:
            # 使用三次卷积采样
            img_resized = img.resize((new_width, new_height), Image.BICUBIC)
        else:
            # 使用双线性内插法采样
            img_resized = img.resize((new_width, new_height), Image.BILINEAR)
        return img_resized


    @staticmethod
    def swap_pixel(img: Image.Image, prob) -> Image.Image:
        """交换"""
        # 获取图片的宽度和高度
        width, height = img.size

        w_num = int(width * prob)
        h_num = int(height * prob)
        w_index = np.random.choice(width, size=w_num, replace=False)
        h_index = np.random.choice(height, size=h_num, replace=False)

        iter = 0
        while iter < int(len(w_index) / 2):
            px = np.random.choice(w_index, 2, replace=False)
            py = np.random.choice(h_index, 2, replace=False)

            w0 = px[0]
            h0 = py[0]

            # 获取像素值
            pixel = img.getpixel((w0, h0))

            w1 = px[1]
            h1 = py[1]

            # 将像素值放到新位置
            img.putpixel((w1, h1), pixel)
            iter += 1

        return img