1.對圖片進行重命名
"""Rename the image based on the folder name"""
import os
import shutil
import sys
import argparse
def main(args):
original_path = args.data_dir
original_path = original_path + "/"
saved_path = args.save_dir
saved_path = saved_path + "/"
make_path(saved_path)
all_folders = traversalDir_FirstDir(original_path)
for folder in all_folders:
files = os.listdir(original_path + folder)
i = 1
for file in files:
suffix = '.png'
name = folder + '_' + str(i).zfill(4) + suffix
i = i + 1
sub_saved_path = saved_path + folder
make_path(sub_saved_path)
shutil.copyfile(original_path + folder + '/' + file, sub_saved_path + '/' + name)
print("rename over")
# To get all sub folders in one folder
def traversalDir_FirstDir(path):
list = []
if (os.path.exists(path)):
files = os.listdir(path)
print("files",files)
for file in files:
m = os.path.join(path, file)
print("m",m)
if (os.path.isdir(m)):
h = os.path.split(m)
print("h",h)
list.append(h[1])
return list
# To judge whether a folder is existed.
def make_path(path):
if not os.path.exists(path):
os.makedirs(path)
def parse_arguments(argv):
parser = argparse.ArgumentParser()
parser.add_argument('data_dir', type=str, help='Directory with aligned images.')
parser.add_argument('save_dir', type=str, help='Directory to save renamed images.')
return parser.parse_args(argv)
if __name__ == '__main__':
main(parse_arguments(sys.argv[1:]))
2.生成pairs文件
#!/usr/bin/env python
# _*_ 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.
"""
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'
def generate(self):
# The repeate times. You can edit this number by yourself
folder_number = self.get_folder_numbers()
# This step will generate the hearder for pair_list.txt, which contains
# the number of classes and the repeate times of generate the pair
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):
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)
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))
file2 = random.choice(os.listdir(self.data_dir + other_dir))
f.write(name + "\t" + file1.split("_")[1].lstrip("0").rstrip(self.img_ext) \
+ "\t" + other_dir + "\t" + file2.split("_")[1].lstrip("0").rstrip(self.img_ext) + "\n")
def parse_arguments(argv):
parser = argparse.ArgumentParser()
parser.add_argument('data_dir', type=str, help='Directory with aligned images.')
parser.add_argument('saved_dir', type=str, help='Directory to save pairs.')
parser.add_argument('--repeat_times', type=str, help='Repeat times to generate pairs', default=30)
return parser.parse_args(argv)
if __name__ == '__main__':
generatePairs = GeneratePairs(parse_arguments(sys.argv[1:]))
generatePairs.generate()
