# -*- coding: utf-8 -*-
import os
if __name__ == '__main__':
    path = os.getcwd()
    dataAnnotated = os.listdir(path + '/Annotations')
    dataNum = len(dataAnnotated)  # 資料集數量

    ftest = open('ImageSets/Main/test.txt', 'w')  # 測試集
    ftrain = open('ImageSets/Main/train.txt', 'w')  # 訓練集
    ftrainval = open('ImageSets/Main/trainval.txt', 'w')  # 訓練驗證集
    fval = open('ImageSets/Main/val.txt', 'w')  # 驗證集
    testScale = 0.2  # 測試集佔總資料集的比例
    trainScale = 0.8  # 訓練集佔訓練驗證集的比例

    i = 1
    testNum = int(dataNum * testScale)  # 測試集的數量
    trainNum = int((dataNum - testNum) * trainScale)  # 訓練集的數量

    for name in dataAnnotated:
        if i <= testNum:
            print>>ftest, name[0:6]
        elif i <= testNum + trainNum:
            print>>ftrain, name[0:6]
            print>>ftrainval, name[0:6]
        else:
            print>>fval, name[0:6]
            print>>ftrainval, name[0:6]
        i += 1
    ftrain.close
    ftrainval.close
    fval.close
    ftest.close
    print 'done!'