# _*_ coding:utf-8 _*_
import os
import random
import argparse
import sys
class GeneratePairs:
"""
Generate the pairs.txt file for applying "validate on LFW" on your own datasets.
"""
#写成命令行格式就用args解析参数
# def __init__(self, args):
# """
# Parameter data_dir, is your data directory.
# Parameter pairs_filepath, where is the pairs.txt that belongs to.
# Parameter img_ext, is the image data extension for all of your image data.
# """
# self.data_dir = args.data_dir
# self.data_dir =self.data_dir + "/"
# self.pairs_filepath = args.saved_dir + "/" + 'pairs.txt'
# self.repeat_times = int(args.repeat_times)
# self.img_ext = '.png'
#在pycharm上直接运行，就用这种直接修改参数的方法比较方便，自己选择
def __init__(self):
"""
Parameter data_dir, is your data directory.
Parameter pairs_filepath, where is the pairs.txt that belongs to.
Parameter img_ext, is the image data extension for all of your image data.
"""
self.data_dir = 'my_data/my_dataset_align'
self.data_dir =self.data_dir + "/" #验证集路径
self.pairs_filepath = '../data' + "/" + 'my_pairs.txt' #pairs.txt存放路径
self.repeat_times = int(30)
# self.img_ext = '.png' #因为我自己验证集png和jpg格式都有，所以不固定图片格式后缀
def generate(self):
# The repeate times. You can edit this number by yourself
folder_number = self.get_folder_numbers()
print('folder_number--',folder_number)
# This step will generate the hearder for pair_list.txt, which contains
# the number of classes and the repeate times of generate the pair
#如果存在旧的pairs先删除
if os.path.exists(self.pairs_filepath):
os.remove(self.pairs_filepath)
#删完重开一个pair.txt
if not os.path.exists(self.pairs_filepath):
with open(self.pairs_filepath,"a") as f:
f.write(str(self.repeat_times) + "\t" + str(folder_number) + "\n")
for i in range(self.repeat_times):
print('第 %d 次：'%int(i))
self._generate_matches_pairs()
self._generate_mismatches_pairs()
def get_folder_numbers(self):
count = 0
for folder in os.listdir(self.data_dir):
if os.path.isdir(self.data_dir + folder):
count += 1
return count
def _generate_matches_pairs(self):
"""
Generate all matches pairs
"""
for name in os.listdir(self.data_dir):
if name == ".DS_Store" or name[-3:] == 'txt':
continue
a = []
for file in os.listdir(self.data_dir + name):
if file == ".DS_Store":
continue
a.append(file)
with open(self.pairs_filepath, "a") as f:
temp = random.choice(a).split("_") # This line may vary depending on how your images are named.
w = temp[0]
# l = random.choice(a).split("_")[1].lstrip("0").rstrip(self.img_ext)
# r = random.choice(a).split("_")[1].lstrip("0").rstrip(self.img_ext)
l = random.choice(a).split("_")[1].lstrip("0").split(".")[0]
r = random.choice(a).split("_")[1].lstrip("0").split(".")[0]
print ('写入 %s ,%s ,%s'%(w,l,r))
f.write(w + "\t" + l + "\t" + r + "\n")
def _generate_mismatches_pairs(self):
"""
Generate all mismatches pairs
"""
for i, name in enumerate(os.listdir(self.data_dir)):
if name == ".DS_Store" or name[-3:] == 'txt':
continue
remaining = os.listdir(self.data_dir)
del remaining[i]
remaining_remove_txt = remaining[:]
for item in remaining:
if item[-3:] == 'txt':
remaining_remove_txt.remove(item)
remaining = remaining_remove_txt
other_dir = random.choice(remaining)
with open(self.pairs_filepath, "a") as f:
file1 = random.choice(os.listdir(self.data_dir + name))
name1=file1.split("_")[0]
file2 = random.choice(os.listdir(self.data_dir + other_dir))
name2=file2.split("_")[0]
f.write(name1 + "\t" + file1.split("_")[1].lstrip("0").split(".")[0] \
+ "\t" + name2 + "\t" + file2.split("_")[1].lstrip("0").split(".")[0] + "\n")
if __name__ == '__main__':
gen=GeneratePairs()
gen.generate()
