python手写lfw数据集转pair.txt形式

最近在搞facenet的人脸识别，用到了lfw数据集，发现没有个合适的将自己数据集转为lfw_pair.txt的代码，所以自己撸了一个。

数据集用的就是lfw，文件夹是名字，名字文件夹內放的就是人脸图像。没有其他文件所有没有做是否是图像的校验。

主要思路就是根据lfw的要求，匹配数据通过同一个人有多张人脸构造，不匹配数据通过不同文件夹的内容构造。

利用python的set数据结构去重特性，一直往里面加就行了，加到数量够了就break。

然后将两个set集合都转为list，根据顺序构造lfw_pair.txt，构造时参考了原来的排序方式，做了sorted。

欢迎批评指教哟

# -*- coding: utf-8 -*-
# @Time    : 2022/3/21 14:39
# @Author  : ChangXinliang
# @Software: PyCharm

import os.path
import random
import os


def create_matched_result(lfw_dir, all_number):
    matched_result = set()
    # 获取所有数据名字
    names = os.listdir(lfw_dir)
    while len(matched_result) < all_number:
        for name in names:
            length = len(os.listdir(os.path.join(lfw_dir, name)))
            if length <= 1:
                continue
            else:
                # 构造两个随机数
                number_1, number_2 = sorted(random.sample(range(1, length + 1), 2))
                s = name + '\t' + str(number_1) + '\t' + str(number_2)
                # 没有必要做判断，set会自行处理
                matched_result.add(s)
                if len(matched_result) == all_number:
                    break
    return list(matched_result)


def create_unmatched_result(lfw_dir, all_number):
    unmatched_result = set()
    names = os.listdir(lfw_dir)
    while len(unmatched_result) < all_number:
        # 随机取两个名字
        name1, name2 = sorted(random.sample(names, 2))
        numbers = []
        for name in [name1, name2]:
            images = os.listdir(os.path.join(lfw_dir, name))
            # 随机取人脸
            number = random.sample(range(1, len(images) + 1), 1)[0]
            numbers.append(number)
        # 没有必要做判断，set会自行处理
        s = name1 + '\t' + str(numbers[0]) + '\t' + name2 + '\t' + str(numbers[1])
        unmatched_result.add(s)
    return list(unmatched_result)


if __name__ == '__main__':
    # 图片数据文件夹
    lfw_dir = 'lfw'
    # 重复次数
    repeat_count = 10
    # 每次匹配数量
    pair_number = 300
    # 总数量
    all_number = repeat_count * pair_number
    # 构造同一张人脸的集合
    matched_result = create_matched_result(lfw_dir, all_number)
    # 构造不同人脸的集合
    unmatched_result = create_unmatched_result(lfw_dir, all_number)

    out_path = r'lfw_pair.txt'
    with open(out_path, 'w') as f:
        f.write('%d %d\n' % (repeat_count, pair_number))
        for i in range(repeat_count):
            for pair in sorted(matched_result[i * pair_number: (i + 1) * pair_number]):
                f.write(pair + '\n')
            for pair in sorted(unmatched_result[i * pair_number: (i + 1) * pair_number]):
                f.write(pair + '\n')