opencv之随机添加噪音

# -*- coding: utf-8 -*-
# Author: Xiaofeng Li
#
#.jpg image generate noise
#

import cv2
import matplotlib.pyplot as plt
import numpy as np
import os
import random
from PIL import Image


class NoiseGenerator(object):

  def generated_noise(self, img):
    """
    Choose which noise to produce.
    :param img(.jpg):
    """
    noise_count = random.randint(1, 6)
    for i in range(noise_count):
      index_noise = random.randint(0, 5)
      if index_noise == 0:
        img = self.gaussian_img(img)
      elif index_noise == 1:
        img = self.salt_img(img)
      elif index_noise == 2:
        img = self.line_img(img)
      elif index_noise == 3:
        img = self.rotate_whole_img(img)
      elif index_noise == 4:
        img = self.erode_img(img)
      elif index_noise == 5:
        img = self.rotate_local_img(img)
    return img

  def gaussian_img(self, img):
    """
    gaussion noise
    """
    im = cv2.GaussianBlur(img, (9, 9), 5)
    return im

  def salt_img(self, img):
    """
       salt noise
       the number of sale dot is n
    """
    n = int(img.shape[0] * img.shape[1] * 0.1)
    ilist = np.random.randint(0, img.shape[1], n)
    jlist = np.random.randint(0, img.shape[0], n)
    for k in range(n):
      i = ilist[k]
      j = jlist[k]
      if img.ndim == 2:
        img[j, i] = 255
      elif img.ndim == 3:
        img[j:j+1, i:i+1,:] = 255
    return img

  def line_img(self, img):
    """
    drow 1 - 10 lines noise
    """
    num_line = np.random.randint(5, 10)
    for i in range(num_line):
      row = np.random.randint(0, img.shape[0], 2)
      col = np.random.randint(0, img.shape[1], 2)
      im = cv2.line(img, (row[0], col[0]),(row[1], col[1]), 0, 2)
    return im

  def rotate_whole_img(self, img):
    """
    rotate whole image
    rotate angle is 0 - 20
    """
    angle = np.random.randint(0, 20)
    h, w = img.shape[:2]
    center = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D(center, angle, 1)
    im = cv2.warpAffine(img, M, (w, h), borderValue=(255, 255, 255))
    return im

  def rotate_local_img(self, img):
    """
    rotate local image
    rotate angle is 0 - 20
    """
    angle = np.random.randint(0, 20)
    random_dot1 = np.random.randint(200, 1400, 2)
    random_dot2 = np.random.randint(300, 2100, 2)
    dot1 = np.sort(random_dot1)
    dot2 = np.sort(random_dot2)
    crop_box = (dot1[0], dot2[0], dot1[1], dot2[1])
    img2 = Image.fromarray(img.astype('uint8')).convert('RGB')

    region = img2.crop(crop_box)
    center = ((dot1[1] - dot1[0]) // 2, (dot2[1] - dot2[0]) // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1)
    region = cv2.warpAffine(np.array(region), M, (dot1[1] - dot1[0], dot2[1] - dot2[0]), borderValue=(255, 255, 255))
    im = Image.fromarray(region.astype('uint8')).convert('RGB')
    img2.paste(im, (dot1[0], dot2[0]))
    return np.array(img2)

  def erode_img(self, img):
    """
       erode noise
    """
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    im = cv2.erode(img, kernel)
    return im



if __name__ == '__main__':
  '''
  ng = NoiseGenerator()
  img = cv2.imread('img/0.jpg')
  im = ng.rotate_local_img(img)
  plt.imshow(im)
  plt.show()
  '''

  for i in range(10):
    target_path = os.getcwd()
    ng = NoiseGenerator()
    img = cv2.imread(os.path.join(target_path, 'img/{}.jpg'.format(i)))
    image = ng.generated_noise(img)
    
    path = os.path.join(target_path, 'noiseimage')
    if not os.path.exists(path):
      os.mkdir(path)
    
    cv2.imwrite(os.path.join(path, '{}.jpg'.format(i)), image)
    print("generated the %d noise image" % (i + 1))