import os

def getIMGname(dir, ftrain, ftest, num):
    filenames = sorted(os.listdir(dir))
    i = 0
    for filename in filenames:
        verifyname = filename.split('.')
        if verifyname[1] == 'jpg':
            if i < num:
                ftrain.write(filename + '\n')
            else:
                ftest.write(filename + '\n')
        i = i+1

    ftrain.close()
    ftest.close()


if __name__ == '__main__':
    dir = './'
    ftrain = open('./train_ids.csv', 'w')
    ftest = open('./test_ids.csv', 'w')
    num = 800000
    getIMGname(dir, ftrain, ftest, num)