pytorch 添加数据增强 运动模糊

在训练过程中增加 运动模糊。

from numpy import random
import numpy as np
from PIL import Image
import io
import glob
import os

class AddSaltPepperNoise(object):

    def __init__(self, p = 0.5, s_vs_p_range=(0.01, 0.2), amount_range=(0.005, 0.025)):
        self.p = p
        self.s_vs_p_range = s_vs_p_range
        self.amount_range = amount_range

    def __call__(self, img):
        if random.random() > self.p:
            return img
        img = np.array(img)  # 图片转numpy
        s_vs_p = np.random.uniform(self.s_vs_p_range[0], self.s_vs_p_range[1])
        amount = np.random.uniform(self.amount_range[0], self.amount_range[1])
        noisy_img = np.copy(img)
        #添加salt噪声
        num_salt = np.ceil(amount * img.size * s_vs_p)
        #设置添加噪声的坐标位置
        coords = [np.random.randint(0, i - 1, int(num_salt)) for i in img.shape]
        noisy_img[coords[0],coords[1],:] = [255,255,255]
        #添加pepper噪声
        num_pepper = np.ceil(amount * img.size * (1. - s_vs_p))
        #设置添加噪声的坐标位置
        coords = [np.random.randint(0,i - 1, int(num_pepper)) for i in img.shape]
        noisy_img[coords[0],coords[1],:] = [0,0,0]
        img = Image.fromarray(img.astype('uint8')) # numpy转图片
        return img

class AddGaussianNoise(object):

    def __init__(self, p = 0.5, mean_range=(-50, 20), variance_range=(5, 30)):
        self.p = p
        self.mean = mean_range
        self.variance = variance_range

    def __call__(self, img):
        if random.random() > self.p:
            return img
        mean = np.random.randint(self.mean[0], self.mean[1])
        variance = np.random.randint(self.variance[0], self.variance[1])
        img = np.array(img)
        h, w, c = img.shape
        N = np.random.normal(mean, variance, size=(h, w, c))
        img = N + img
        img = np.clip(img, a_min=0, a_max=255)
        img = Image.fromarray(img.astype('uint8'))
        return img

import cv2

class MotionBlur(object):

    def __init__(self, p = 0.5, degree = 5, angle = 45):
        self.p = p
        self.degree = degree
        self.angle = angle

    def __call__(self, img):
        if random.random() > self.p:
            return img
        angle = np.random.randint(-self.angle, self.angle+1)
        degree = np.random.randint(1, self.degree)
        img = np.array(img)
        M = cv2.getRotationMatrix2D((degree / 2, degree / 2), angle, 1)
        motion_blur_kernel = np.diag(np.ones(degree))
        motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (degree, degree))
        motion_blur_kernel = motion_blur_kernel / degree
        blurred = cv2.filter2D(img, -1, motion_blur_kernel)
        # convert to uint8
        cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
        img = Image.fromarray(blurred.astype('uint8'))
        return img

class ResizeBlur(object):

    def __init__(self, p = 0.5, minsize = 40, maxsize = 132, dstsize = 112):
        self.p = p
        self.minsize = minsize
        self.maxsize = maxsize
        self.dstsize = dstsize
    def __call__(self, img):
        if random.random() > self.p:
            return img
        newsize = np.random.randint(self.minsize, self.maxsize)
        resample = np.random.randint(0, 4)  #插值方法
        img = img.resize((newsize, newsize), resample)
        img = img.resize((self.dstsize, self.dstsize), resample)
        
        return img


class CompressJPEG(object):
    def __init__(self, p = 0.5, quality_low=50, quality_hight=100):
        self.p = p
        self.quality_low = quality_low
        self.quality_hight = quality_hight
        
    def __call__(self, img):
        if random.random() > self.p:
            return img
        buffer = io.BytesIO()
        quality = np.random.randint(self.quality_low, self.quality_hight)
        img.save(buffer, format='JPEG', quality=quality)
        compressed_image = Image.open(buffer)
        return compressed_image
    
    
class RandomCropBottom(object):
    def __init__(self, p = 0.5, left_ratio=0.1, right_ratio=0.1, top_ratio = 0.1, bottom_ratio = 0.6):
        self.p = p
        self.left_ratio = left_ratio
        self.right_ratio = right_ratio
        self.top_ratio = top_ratio
        self.bottom_ratio = bottom_ratio

    def __call__(self, img):
        if random.random() > self.p:
            return img
        width, height = img.size
        left = np.random.randint(0, (width*self.left_ratio))
        right = np.random.randint(0, (width*self.right_ratio))
        top = np.random.randint(0, (height*self.top_ratio))
        bottom = np.random.randint(0, (height*self.bottom_ratio))
        return img.crop([left, top, width-right, height-bottom])

class RandomCopyPast(object):
    def __init__(self, p = 0.5, cache_prob=0.6, past_source='/home/liushuan/datasets/person/reid/train_data_0119/train', tal='.jpg'):
        self.p = p
        self.random_past_files = glob.glob(past_source +"/*/*"+ tal, recursive=True)
        print("len past files :", len(self.random_past_files))
        self.length = len(self.random_past_files)
        self.cache_prob = cache_prob
        self.cache_lists = []
        self.cache_max_size = 1000
        
    def get_img(self, file_index):
        cache_img = Image.open(self.random_past_files[file_index])
        width, height = cache_img.size
        left = np.random.randint(0, (width*0.3))
        right = np.random.randint(0, (width*0.3))
        top = np.random.randint(0, (height*0.3))
        bottom = np.random.randint(0, (height*0.4))
        cache_img = cache_img.crop([left, top, width-right, height-bottom])
        return cache_img
        
    def __call__(self, img):
        if random.random() > self.p:
            return img
        index = np.random.randint(0, self.length)
        src_width, src_height = img.size
        ### cache image 
        if len(self.cache_lists) < self.cache_max_size:
            cache_img = self.get_img(index)
            self.cache_lists.append(cache_img)
        else:
            if random.random() > self.cache_prob:
                cache_img = self.get_img(index)
                del self.cache_lists[0]
                self.cache_lists.append(cache_img)
            else:
                cache_img = self.cache_lists[index % len(self.cache_lists)]

        resample = np.random.randint(0, 4)  #插值方法
        cache_img = cache_img.resize((src_width//3, src_height//3), resample)
        if random.random() > 0.5:
            offset_x = np.random.randint(0, (src_width*0.1))
        else:
            offset_x = np.random.randint((src_width*0.6), (src_width*0.7))
        offset_y = np.random.randint(0, (src_width*0.7))
        img.paste(cache_img, (offset_x, offset_y))
        return img
    
class SaveImg(object):
    def __init__(self, save_path='/home/liushuan/datasets/check/reid'):
        self.save_path = save_path
        self.count = 0
        
    def __call__(self, img):
        img.save(os.path.join(self.save_path, str(self.count)+".jpg"))
        self.count += 1
        return img