import numpy as np
from PIL import Image
import os
import shutil
# import cv2


# 对收集的5000张自然图像进行处理，因为利用爬虫爬的图片大小，图片的模式都不相同，所以为了数据的统一先进行处理。
def convertImage():
    listpath = os.listdir("../nature_picture")  # 这里path 表示你大自然图片的数据集路径。
    if os.path.exists("../nature_yellow"):  # 如果文件存在
        # 删除文件，可使用以下两种方法。
        shutil.rmtree("../nature_yellow")
    if not os.path.isdir("../nature_yellow"):  # 判断文件夹是否存在
        os.mkdir("../nature_yellow")  # 创建文件夹
    for path in listpath:
        img = Image.open(os.path.join("../nature_picture", path))
        print(os.path.join("../nature_picture",path))
        img = img.convert("RGB")  # 转化为RGB模式
        # 尺寸设置为224*224，采用细节增强缩放  ，
        img = img.resize((224, 224), Image.ANTIALIAS)
        print(os.path.join("../nature_yellow",path),img.size)
        img.save(os.path.join("../nature_yellow",path))  # 保存处理好的图片
        # cv2.imwrite(os.path.join("../nature_yellow",path),img)

def createDataset(dirimage):
    listpath = os.listdir(dirimage)
    for index, path in enumerate(listpath):  # enumerate 这个函数的作用就是增加一个索引值（可以理解为加序号）
        img = Image.open(os.path.join(dirimage, path))
        """
        这里将数据集分为训练集4000张（2000
        张作为正样本，也就是贴上小黄人的样本，2000
        张作为负样），
        测试集1000张（500
        张作为正样本，500
        张作为负样本）。0 - 1999, 4000 - 4499
        这两个区间上的图片粘贴小黄人，其他的不粘贴。
        """
        if index < 2000 or (index >= 4000 and index < 4500):
            """随机取出一张小黄人图片，这里还可以用np.random.choice(img_yellow)，img_yellow 是一个列表，这是从列表中随机取一个值的方式"""
            minions = Image.open("../yellow_picture/{}.jpg".format(np.random.randint(1, 13)))
            # 缩放
            h = w = np.random.randint(64, 180)
            # 将小黄人的尺寸随机缩放成64 - 180 的大小
            minions = minions.resize((h, w), Image.ANTIALIAS)
            # 旋转
            minions = minions.rotate(np.random.randint(-30, 30))
            # 翻转，镜像翻转
            minions = minions.transpose(Image.FLIP_LEFT_RIGHT) if np.random.randint(0, 2) == 1 else minions
            x, y = np.random.randint(0, 224 - w), np.random.randint(0, 224 - h)
            # 掩码
            r, g, b = minions.split()  # 分离通道小黄人是RGBA 的格式
            # mask 表示图像的一部分，或者说你感兴趣的部分。这里粘贴小黄人的a通道
            img.paste(minions, (x, y), mask=r)

            # print(x,y)
            if not os.path.isdir("../new_nature_yellow"):  # 判断文件夹是否存在
                os.mkdir("../new_nature_yellow")  # 创建文件夹
            img.save("../new_nature_yellow/{}.{}.{}.{}.{}.{}.jpg".format(index, x, y, x + w, y + h, 1))
        else:
            img.save("../new_nature_yellow/{}.0.0.0.0.0.jpg".format(index))

if __name__ == '__main__':
    convertImage()
    createDataset(r"../nature_yellow")