1. 程式人生 > >獲取當前目錄下的有效圖片

獲取當前目錄下的有效圖片

獲取當前目錄下所有的子目錄

sub_dirs = [x[0] for x in os.walk(INPUT_DATA)]
# 得到的第一個目錄是當前目錄,不需要考慮
is_root_dir = True
for sub_dir in sub_dirs:
    if is_root_dir:
        is_root_dir = False
        continue

    # 獲取當前目錄下所有的有效圖片檔案。
    extensions = ['jpg', 'jpeg', 'JPG', 'JPEG']
    file_list = []
    dir_name = os.path.basename(sub_dir)
    for extension in extensions:
        file_glob = os.path.join(INPUT_DATA, dir_name, '*.'+extension)
        file_list.extend(glob.glob(file_glob))
    if not file_list:
        continue
    # 通過目錄名獲取類別的名稱。
    label_name = dir_name.lower()
    # 初始化當前類別的訓練資料集、測試資料集和驗證資料集
    training_images = []
    testing_images = []
    validation_images = []
    for file_name in file_list:
        base_name = os.path.basename(file_name)
        # 隨機將資料分到訓練資料集、測試資料集和驗證資料集。
        chance = np.random.randint(100)
        if chance < validation_percentage:
            validation_images.append(base_name)
        elif chance < (testing_percentage + validation_percentage):
            testing_images.append(base_name)
        else:
            training_images.append(base_name)

    # 將當前類別的資料放入結果字典。
    result[label_name] = {
        'dir': dir_name,
        'training': training_images,
        'testing': testing_images,
        'validation': validation_images
        }