python opencv批量对图像进行增强(OCR文本图)

环境win10+python3.9+opencv-python

原图：

增强后：

python脚本：

import os
import cv2
import glob
import pathlib
import random
import cv2 as cv
import numpy as np

# 将-数字-中的训练数据进行图像数据增强,这里使用随机图像增强方法

images_data_path = r'E:\datasets\gen_number_str_hw\train'

images_save_path = r'E:\datasets\gen_number_str_hw\train_enhance'
if not os.path.exists(images_save_path):
    os.makedirs(images_save_path)


def tfactor(img):
    hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)

    hsv[:, :, 0] = hsv[:, :, 0] * (0.8 + np.random.random() * 0.2)
    hsv[:, :, 1] = hsv[:, :, 1] * (0.3 + np.random.random() * 0.7)
    hsv[:, :, 2] = hsv[:, :, 2] * (0.2 + np.random.random() * 0.8)

    img = cv.cvtColor(hsv, cv.COLOR_HSV2BGR)
    return img


def AddGauss(img, level):
    # return cv2.blur(img, (level * 2 + 1, level * 2 + 1))
    return cv.blur(img, (1 + r(1), 1 + r(1)))


def r(val):
    return int(np.random.random() * val)


def AddNoiseSingleChannel(single):
    diff = 255 - single.max()
    noise = np.random.normal(0, 1 + r(6), single.shape)
    noise = (noise - noise.min()) / (noise.max() - noise.min())
    noise = diff * noise
    noise = noise.astype(np.uint8)
    dst = single + noise
    return dst


def addNoise(img, sdev=0.5, avg=10):
    img[:, :, 0] = AddNoiseSingleChannel(img[:, :, 0])
    img[:, :, 1] = AddNoiseSingleChannel(img[:, :, 1])
    img[:, :, 2] = AddNoiseSingleChannel(img[:, :, 2])
    return img


image_cnt = 0

for img_path in glob.glob(images_data_path + '/*.jpg', recursive=True):
    img_file = pathlib.Path(img_path)
    imagename = str(img_file.stem)
    imagesuffix = str(img_file.suffix)
    imgnamewithsuffix = imagename + imagesuffix
    image_cnt = image_cnt + 1
    print(image_cnt)

    img = cv.imread(os.path.join(images_data_path, imgnamewithsuffix))
    if img is None:
        print("错误:%s", imgnamewithsuffix)
        continue
    img_h, img_w, img_c = img.shape
    image_src = img[0:, 0:].copy()  # 这里可以保存原图

    if random.randint(1, 1000) % 5 == 0:
        img = image_src[0:, 0:].copy()
        img = cv.bitwise_not(img)  # 得到黑底白字
    elif random.randint(1, 1000) % 7 == 0:
        img = image_src[0:, 0:].copy()
        img = tfactor(img)  # 调灰度
    elif random.randint(1, 1000) % 6 == 0:
        img = image_src[0:, 0:].copy()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    elif random.randint(1, 1000) % 7 == 0:
        img = image_src[0:, 0:].copy()
        img = AddGauss(img, 0)  # 加高斯平滑
    elif random.randint(1, 1000) % 6 == 0:
        img = image_src[0:, 0:].copy()
        img = addNoise(img)  # 加噪声

    imgname = imagename + imagesuffix
    savepath = os.path.join(images_save_path, imgname)
    cv.imwrite(savepath, img)