FaceNet-pytorch(fixing data imbalance-CASIA)

# ------------------------------------------------#
#   进行训练前需要利用这个文件生成cls_train.txt
# ------------------------------------------------#
import os
import random


def txt_annotation(datasets_path, select_num):
    '''
    :param datasets_path: 数据集根目录
    :param select_num: 每一个folder选择的文件数目
    :return: cls_train.txt
    '''

    types_name = os.listdir (datasets_path)

    types_name = sorted (types_name)

    # print (types_name)
    list_file = open ('cls_train.txt', 'w')
    for cls_id, type_name in enumerate (types_name):
        photos_path = os.path.join (datasets_path, type_name)
        if not os.path.isdir (photos_path) or select_num > len (os.listdir (photos_path)):
            print (
                'pass folder {} with {} numbers , required {} '.format (type_name, len (os.listdir (photos_path)),
                                                                        select_num))
            continue

        photos_name = os.listdir (photos_path)
        photos_name = random.sample (photos_name, select_num)
        # print(photos_name)

        for photo_name in photos_name:
            # print (cls_id)
            list_file.write (
                str (cls_id) + ";" + '%s' % (os.path.join (os.path.abspath (datasets_path), type_name, photo_name)))
            list_file.write ('\n')
    list_file.close ()


if __name__ == "__main__":
    dataset_path = r'test_dataset'
    select_num = 32

    txt_annotation (dataset_path, select_num)

效果如下：