import os
import shutil
import random
 
 
def splitImg2Category(dataDir="images/",resDir="dataset/"):

    for one_pic in os.listdir(dataDir):
        one_path=dataDir+one_pic
        oneDir=resDir+one_pic.split('_')[0].strip()+"/"
        if not os.path.exists(oneDir):
            os.makedirs(oneDir)
        shutil.copy(one_path,oneDir+one_pic)

def random2Dataset(dataDir='data/original/',ratio=0.3):

    label_list=os.listdir(dataDir)
    for one_label in label_list:
        oneDir=dataDir+one_label+'/'
        pic_list=os.listdir(oneDir)
        testNum=int(len(pic_list)*ratio)
        oneTrainDir='data/train/'+one_label+'/'
        oneTestDir='data/val/'+one_label+'/'
        if not os.path.exists(oneTrainDir):
            os.makedirs(oneTrainDir)
        if not os.path.exists(oneTestDir):
            os.makedirs(oneTestDir)
        for i in range(testNum):
            one_path=oneDir+random.choice(os.listdir(oneDir))
            name=str(len(os.listdir(oneTestDir))+1)
            new_path=oneTestDir+one_label+'_'+name+'.jpg'
            shutil.move(one_path,new_path)

        for one_pic in os.listdir(oneDir):
            one_path=oneDir+one_pic
            name=str(len(os.listdir(oneTrainDir))+1)
            new_path=oneTrainDir+one_label+'_'+name+'.jpg'
            shutil.move(one_path,new_path)
            
data_dir = "/Users/aiot/Documents/data_augment/ExploreAug/data/iiit-pet"
random2Dataset(dataDir=os.path.join(data_dir,"dataset/"),
               ratio=0.3)   