仿照LFW的pair.txt生成自己数据集的txt文件

最近在跑facenet的valida_on_lfw.py的样例,想要在自己的数据集上运行该代码。
关于pair.txt的描述,见http://vis-www.cs.umass.edu/lfw/README.txt
编写代码生成自己的pair.txt文件:

import glob
import os.path
import random
import numpy as np
import os
import json

# 图片数据文件夹
INPUT_DATA = r''
# 这个函数从数据文件夹中读取所有的图片列表并按训练、验证、测试数据分开
# testing_percentage和validation_percentage指定了测试数据集和验证数据集的大小
def create_image_lists():
    matched_result = set()	# 用于存放同一个人的照片
    k = 0
    # 获取当前目录下所有的子目录,这里x 是一个三元组(root,dirs,files),第一个元素表示INPUT_DATA当前目录,
    # 第二个元素表示当前目录下的所有子目录,第三个元素表示当前目录下的所有的文件
    sub_dirs = [x[0] for x in os.walk(INPUT_DATA)]
    # sub_dirs[0]表示当前文件夹本身的地址,不予考虑,只考虑他的子目录(各个类别的花)
    for i in range(12):		# 循环次数可以根据需要的数量调整
        for sub_dir in sub_dirs[1:]:
            # 获取当前目录下所有的有效图片文件
            extensions = 'png'
            # 把图片存放在file_list列表里
            file_list = []
            # os.path.basename(sub_dir)返回sub_sir最后的文件名
            # 如os.path.basename('./flower_data/daisy')返回daisy
            dir_name = os.path.basename(sub_dir)
            file_glob = os.path.join(INPUT_DATA, dir_name, '*.' + extensions)
            # glob.glob(file_glob)获取指定目录下的所有图片,存放在file_list中
            file_list.extend(glob.glob(file_glob))
            if not file_list: continue
            # 通过目录名获取类别的名称
            label_name = dir_name
            length = len(file_list)
            base_name1 = os.path.basename(file_list[i%length])  # 获取文件的名称
            base_name2 = os.path.basename(file_list[-((i + 1) % length)])
            if(file_list[-((i + 1) % length)] == file_list[i%length]):
                base_name2 = os.path.basename(file_list[-((i+2)%length)])
            # 将当前类别是数据放入结果集合
            matched_result.add(label_name +'\t'+ base_name1+ '\t'+ base_name2)
            k = k + 1

    # 返回整理好的所有数据
    return matched_result, k

def create_pairs():
    unmatched_result = set()       # 不同类的匹配对
    k = 0
    sub_dirs = [x[0] for x in os.walk(INPUT_DATA)]
    # sub_dirs[0]表示当前文件夹本身的地址,不予考虑,只考虑他的子目录
    for sub_dir in sub_dirs[1:]:
        # 获取当前目录下所有的有效图片文件
        extensions = ['png']
        file_list = []
        # 把图片存放在file_list列表里
        # os.path.basename(sub_dir)返回sub_sir最后的文件名
        # 如os.path.basename('./flower_data/daisy')返回daisy
        dir_name = os.path.basename(sub_dir)
        for extension in extensions:
            file_glob = os.path.join(INPUT_DATA, dir_name, '*.'+extension)
            # glob.glob(file_glob)获取指定目录下的所有图片,存放在file_list中
            file_list.extend(glob.glob(file_glob))

    length_of_dir = len(sub_dirs)
    for j in range(12):		# 循环次数可以根据需要的数量调整
        for i in range(length_of_dir):
            class1 = sub_dirs[i]
            class2 = sub_dirs[(length_of_dir-i+j-1) % length_of_dir]
            if ((length_of_dir-i+j-1) % length_of_dir):
                class1_name = os.path.basename(class1)
                class2_name = os.path.basename(class2)
                # 获取当前目录下所有的有效图片文件
                extensions = 'png'
                file_list1 = []
                file_list2 = []
                # 把图片存放在file_list列表里
                # os.path.basename(sub_dir)返回sub_sir最后的文件名
                # 如os.path.basename('./flower_data/daisy')返回daisy
                file_glob1 = os.path.join(INPUT_DATA, class1_name, '*.' + extension)
                file_list1.extend(glob.glob(file_glob1))
                file_glob2 = os.path.join(INPUT_DATA, class2_name, '*.' + extension)
                file_list2.extend(glob.glob(file_glob2))
                if file_list1 and file_list2:
                    base_name1 = os.path.basename(file_list1[j % len(file_list1)])  # 获取文件的名称
                    base_name2 = os.path.basename(file_list2[j % len(file_list2)])
                    # unmatched_result.add([class1_name, base_name1, class2_name, base_name2])
                    unmatched_result.add(class1_name+'\t'+base_name1+'\t'+class2_name+'\t'+base_name2)
                    k = k + 1
    return unmatched_result, k

result, k1 = create_image_lists()
print(k1)
# print(result)

result_un, k2 = create_pairs()
print(k2)
# print(result_un)

file = open('pairs.txt', 'w')

result1 = list(result)
result2 = list(result_un)

for i in range(10):		# 将结果写入到文件中
    for pair in result1[i*300:i*300+300]:
        file.write(pair + '\n')

    for pair in result2[i*300:i*300+300]:
        file.write(pair + '\n')
        # file.write('\n')

完成以后再修改一下lfw.py中get_paths()的代码,
仿照LFW的pair.txt生成自己数据集的txt文件
修改为:

    for pair in pairs:
        if len(pair) == 3:
            path0 = os.path.join(lfw_dir, pair[0], pair[1])
            path1 = os.path.join(lfw_dir, pair[0], pair[2])
            issame = True
        elif len(pair) == 4:
            path0 = os.path.join(lfw_dir, pair[0], pair[1])
            path1 = os.path.join(lfw_dir, pair[2], pair[3])
            issame = False
        if os.path.exists(path0) and os.path.exists(path1):    # Only add the pair if both paths exist
            path_list += (path0, path1)
            issame_list.append(issame)
        else:
            print(pair)
            nrof_skipped_pairs += 1