数据增强 + 1：9分割数据

# import random
# import shutil
#
# import cv2
# import glob
# import os
#
# x = [1,2,3,4,5,6]
# # random.shuffle可以随机乱序列表中的元素
# random.shuffle(x)
# print(x)
# 该部分只需要修改六个目标文件夹

# 值得注意的是shutil.copy如果目录中存在该名称的文件将不会再复制

# up_path = r"C:\Users\sun\Desktop\data devide"
# up_newpath = r"C:\Users\sun\Desktop\datasets"
# train_path = r"C:\Users\sun\Desktop\datasets\train\Anomaly"
# test_path = r"C:\Users\sun\Desktop\datasets\test\Anomaly"
# train_path_noanomaly = r"C:\Users\sun\Desktop\datasets\train\No-Anomaly"
# test_path_noanomaly = r"C:\Users\sun\Desktop\datasets\test\No-Anomaly"
#
# list = os.listdir(up_path)
# for lists in list:
#     if lists != "No-Anomaly":
#         path = os.path.join(up_path, lists)
#         list_pisture = os.listdir(path)
#         list_pisture_len = len(list_pisture)
#         train_picture = list_pisture_len*0.9
#         test_picture = list_pisture_len*0.1
#         list_train = list_pisture[0:int(train_picture)]
#         random.shuffle(list_train)
#         list_test = list_pisture[int(train_picture):list_pisture_len+1]
#         random.shuffle(list_test)
#
#         # 该部分用于训练集的生成
#         for list_trains in list_train:
#             oldpath = os.path.join(path, list_trains)
#             shutil.copy(oldpath,train_path)
#
#         # 该部分用于测试集的生成
#         for list_tests in list_test:
#             oldpath = os.path.join(path, list_tests)
#             shutil.copy(oldpath,test_path)
#
#     elif lists == "No-Anomaly":
#         path = os.path.join(up_path, lists)
#         list_pisture = os.listdir(path)
#         list_pisture_len = len(list_pisture)
#         train_picture = list_pisture_len * 0.9
#         test_picture = list_pisture_len * 0.1
#         list_train = list_pisture[0:int(train_picture)]
#         random.shuffle(list_train)
#         list_test = list_pisture[int(train_picture):list_pisture_len + 1]
#         random.shuffle(list_test)
#
#         # 该部分用于训练集的No-Anomaly生成
#         for list_trains in list_train:
#             oldpath = os.path.join(path, list_trains)
#             shutil.copy(oldpath, train_path_noanomaly)
#
#         # 该部分用于测试集的No-Anomaly生成
#         for list_tests in list_test:
#             oldpath = os.path.join(path, list_tests)
#             shutil.copy(oldpath, test_path_noanomaly)
# ##########################上面的部分用于分割数据集##################################

import glob
import cv2
import os

from albumentations import (
    PadIfNeeded,
    HorizontalFlip,    # 随机水平翻转
    VerticalFlip,      # 随机垂直翻转
    Flip,              # 随机翻转
    CenterCrop,
    Crop,
    Compose,
    Transpose,
    RandomRotate90,    # 随机90度旋转
    ElasticTransform,
    GridDistortion,
    OpticalDistortion,
    RandomSizedCrop,   # 随机尺寸裁剪并缩放回原始大小
    OneOf,
    CLAHE,
    RandomBrightnessContrast,
    RandomGamma
)

if __name__ == "__main__":
    # old_path = ("C:/Users/sun/Desktop/color")     # 输入img地址
    # new_path = ("C:/Users/sun/Desktop/addcolor")  # 输出img地址
    # 多个文件夹的上一层文件  直接将增广得到的数据加入文件夹内
    up_path = ("C:/Users/sun/Desktop/data devide")
    list = os.listdir(up_path)
    for lists in list:
        old_path = os.path.join(up_path, lists)
        new_path = old_path
        pictures = os.listdir(old_path)
        a = 1
        for picture in pictures:
            old_dir = os.path.join(old_path, picture)
            # cv读入图片
            image = cv2.imread(old_dir)
            # 将图片进行100%的垂直翻转、水平翻转、旋转180°
            h_image = HorizontalFlip(p=1)(image=image)['image']
            v_image = VerticalFlip(p=1)(image=image)['image']
            f_image = Flip(p=1)(image=image)['image']

            # 取出图片的后缀名.jpg 这里os.path.splitext将文件名分开，[1]取出来它的后半部分
            filetype = os.path.splitext(old_dir)[1]
            # 因为这里有三种变换方式，所以给了三种命名
            h_new_dir = os.path.join(new_path, "h_cellcolor" + str(a) + filetype)
            cv2.imwrite(h_new_dir, h_image)
            v_new_dir = os.path.join(new_path, "v_cellcolor" + str(a) + filetype)
            cv2.imwrite(v_new_dir, v_image)
            f_new_dir = os.path.join(new_path, "f_cellcolor" + str(a) + filetype)
            cv2.imwrite(f_new_dir, f_image)
            a = a + 1