opencv图像噪声滤波

import glob
import json
import os
import random
import shutil
from typing import Tuple

import cv2 as cv
import numpy as np
from tqdm import tqdm


def f_gaussian_noise(img, mean=0, var=0.001):
    img = np.array(img / 255, dtype=float)
    noise = np.random.normal(mean, var ** 0.5, img.shape)
    out = img + noise
    if out.min() < 0:
        low_clip = -1.
    else:
        low_clip = 0.
    out = np.clip(out, low_clip, 1.0)
    out = np.uint8(out * 255)
    return out


def f_salt_noise(img, prob):
    output = np.zeros(img.shape, np.uint8)
    thresh = 1 - prob
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            rdn = random.random()
            if rdn < prob:
                output[i][j] = 0
            elif rdn > thresh:
                output[i][j] = 255
            else:
                output[i][j] = img[i][j]
    return output


class DataAugment:
    def __init__(self, img_glob_path, img_save_path, json_save_path, salt_noise: bool = True,
                 gaussian_noise: bool = True, median_blur: bool = True,
                 gaussian_blur: bool = True, salt_noise_prop=0.02, gaussian_noise_param: Tuple[int, float] = (0, 0.001),
                 median_blur_k_size=3, gaussian_blur_k_size=3):
        """

        :param img_glob_path: 图片读取路径
        :param img_save_path: 图片保存路径
        :param json_save_path: json文件保存路径
        :param salt_noise: 椒盐噪声
        :param gaussian_noise: 高斯噪声
        :param median_blur: 中值模糊
        :param gaussian_blur: 高斯模糊
        :param salt_noise_prop: 椒盐比例
        :param gaussian_noise_param: 高斯噪声参数（均值，方差）
        :param median_blur_k_size: 中值模糊核
        :param gaussian_blur_k_size: 高斯模糊核
        """
        self.img_glob_path = img_glob_path
        self.img_save_path = img_save_path
        self.json_save_path = json_save_path
        self.salt_noise = salt_noise
        self.gaussian_noise = gaussian_noise
        self.median_blur = median_blur
        self.gaussian_blur = gaussian_blur
        self.gaussian_blur_k_size = gaussian_blur_k_size
        self.median_blur_k_size = median_blur_k_size
        self.gaussian_noise_var = gaussian_noise_param[1]
        self.gaussian_noise_mean = gaussian_noise_param[0]
        self.salt_noise_prop = salt_noise_prop

    def augment(self):
        count = 0
        for img_file_path in tqdm(glob.glob(self.img_glob_path)):
            img_save_path: str = self.img_save_path + "/" + str(count) + ".jpg"
            shutil.copy(img_file_path, img_save_path)

            img = cv.imread(img_file_path)
            if self.salt_noise:
                salt_noise_img = f_salt_noise(img, self.salt_noise_prop)
                salt_noise_img_file_path = img_save_path.replace(".jpg", "_salt_noise.jpg")
                cv.imwrite(salt_noise_img_file_path, salt_noise_img)
            if self.gaussian_noise:
                gaussian_noise_img = f_gaussian_noise(img, self.gaussian_noise_mean, self.gaussian_noise_var)
                gaussian_noise_img_file_path = img_save_path.replace(".jpg", "_gaussian_noise.jpg")
                cv.imwrite(gaussian_noise_img_file_path, gaussian_noise_img)
            if self.median_blur:
                median_blur_img = cv.medianBlur(img, self.median_blur_k_size)
                median_blur_img_file_path = img_save_path.replace(".jpg", "_median_blur.jpg")
                cv.imwrite(median_blur_img_file_path, median_blur_img)
            if self.gaussian_blur:
                gaussian_blur_img = cv.GaussianBlur(img, (self.gaussian_blur_k_size, self.gaussian_blur_k_size), 0)
                gaussian_blur_img_file_path = img_save_path.replace(".jpg", "_gaussian_blur.jpg")
                cv.imwrite(gaussian_blur_img_file_path, gaussian_blur_img)
            count += 1


if __name__ == '__main__':
    data_augment = DataAugment(img_glob_path='./data/kinetic/image*/*.jpg', img_save_path='./labelme_imgs',
                               json_save_path='./labelme_jsons', salt_noise_prop=0.002,
                               gaussian_noise=False, median_blur=False, salt_noise=False,gaussian_blur=False)
    data_augment.augment()