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_data_path = r'images/Snipaste_2023-06-05_10-31-56.jpg'

# 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)
img = cv.imread(images_data_path)
# cv.imshow('img', img)
# cv.waitKey(1000)


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=0):
    # 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

def add_noise(img):
    for i in range(20): #添加点噪声
        temp_x = np.random.randint(0,img.shape[0])
        temp_y = np.random.randint(0,img.shape[1])
        img[temp_x][temp_y] = 255
    return img

def add_erode(img):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3, 3))
    img = cv2.erode(img,kernel)
    return img

# 适当膨胀
def add_dilate(img):
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3, 3))
    img = cv2.dilate(img,kernel)
    return img

# img = tfactor(img)

# img = cv.bitwise_not(img)

img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

img = add_noise(img)
# img = add_dilate(img)
# img = add_erode(img)


cv.imshow('img', img)
cv.waitKey(5000)

# img = AddGauss(img)
# cv.imshow('img', img)
# cv.waitKey(1000)

# img = AddNoiseSingleChannel(img)
# cv.imshow('img', img)
# cv.waitKey(1000)

# img = addNoise(img)
# cv.imshow('img', img)
# cv.waitKey(10000)


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)