mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-18 06:38:19 +00:00
140 lines
3.8 KiB
Python
140 lines
3.8 KiB
Python
|
|
from easydict import EasyDict as edict
|
|
import os
|
|
|
|
|
|
def load_property(data_dir):
|
|
prop = edict()
|
|
for line in open(os.path.join(data_dir, 'property')):
|
|
vec = line.strip().split(',')
|
|
assert len(vec)==3
|
|
prop.num_classes = int(vec[0])
|
|
prop.image_size = [int(vec[1]), int(vec[2])]
|
|
return prop
|
|
|
|
|
|
|
|
def get_dataset_webface(input_dir):
|
|
clean_list_file = input_dir+"_clean_list.txt"
|
|
ret = []
|
|
for line in open(clean_list_file, 'r'):
|
|
vec = line.strip().split()
|
|
assert len(vec)==2
|
|
fimage = edict()
|
|
fimage.id = vec[0].replace("\\", '/')
|
|
fimage.classname = vec[1]
|
|
fimage.image_path = os.path.join(input_dir, fimage.id)
|
|
ret.append(fimage)
|
|
return ret
|
|
|
|
def get_dataset_celeb(input_dir):
|
|
clean_list_file = input_dir+"_clean_list.txt"
|
|
ret = []
|
|
dir2label = {}
|
|
for line in open(clean_list_file, 'r'):
|
|
line = line.strip()
|
|
if not line.startswith('./m.'):
|
|
continue
|
|
line = line[2:]
|
|
vec = line.split('/')
|
|
assert len(vec)==2
|
|
if vec[0] in dir2label:
|
|
label = dir2label[vec[0]]
|
|
else:
|
|
label = len(dir2label)
|
|
dir2label[vec[0]] = label
|
|
|
|
fimage = edict()
|
|
fimage.id = line
|
|
fimage.classname = str(label)
|
|
fimage.image_path = os.path.join(input_dir, fimage.id)
|
|
ret.append(fimage)
|
|
return ret
|
|
|
|
def _get_dataset_celeb(input_dir):
|
|
list_file = input_dir+"_original_list.txt"
|
|
ret = []
|
|
for line in open(list_file, 'r'):
|
|
vec = line.strip().split()
|
|
assert len(vec)==2
|
|
fimage = edict()
|
|
fimage.id = vec[0]
|
|
fimage.classname = vec[1]
|
|
fimage.image_path = os.path.join(input_dir, fimage.id)
|
|
ret.append(fimage)
|
|
return ret
|
|
|
|
def get_dataset_facescrub(input_dir):
|
|
ret = []
|
|
label = 0
|
|
person_names = []
|
|
for person_name in os.listdir(input_dir):
|
|
person_names.append(person_name)
|
|
person_names = sorted(person_names)
|
|
for person_name in person_names:
|
|
subdir = os.path.join(input_dir, person_name)
|
|
if not os.path.isdir(subdir):
|
|
continue
|
|
for _img in os.listdir(subdir):
|
|
fimage = edict()
|
|
fimage.id = os.path.join(person_name, _img)
|
|
fimage.classname = str(label)
|
|
fimage.image_path = os.path.join(subdir, _img)
|
|
ret.append(fimage)
|
|
label += 1
|
|
return ret
|
|
|
|
def get_dataset_megaface(input_dir):
|
|
ret = []
|
|
label = 0
|
|
for prefixdir in os.listdir(input_dir):
|
|
_prefixdir = os.path.join(input_dir, prefixdir)
|
|
for subdir in os.listdir(_prefixdir):
|
|
_subdir = os.path.join(_prefixdir, subdir)
|
|
if not os.path.isdir(_subdir):
|
|
continue
|
|
for img in os.listdir(_subdir):
|
|
if not img.endswith('.jpg.jpg') and img.endswith('.jpg'):
|
|
fimage = edict()
|
|
fimage.id = os.path.join(prefixdir, subdir, img)
|
|
fimage.classname = str(label)
|
|
fimage.image_path = os.path.join(_subdir, img)
|
|
ret.append(fimage)
|
|
label+=1
|
|
return ret
|
|
|
|
def get_dataset_common(input_dir):
|
|
ret = []
|
|
label = 0
|
|
person_names = []
|
|
for person_name in os.listdir(input_dir):
|
|
person_names.append(person_name)
|
|
person_names = sorted(person_names)
|
|
for person_name in person_names:
|
|
_subdir = os.path.join(input_dir, person_name)
|
|
if not os.path.isdir(_subdir):
|
|
continue
|
|
for img in os.listdir(_subdir):
|
|
fimage = edict()
|
|
fimage.id = os.path.join(person_name, img)
|
|
fimage.classname = str(label)
|
|
fimage.image_path = os.path.join(_subdir, img)
|
|
fimage.bbox = None
|
|
fimage.landmark = None
|
|
ret.append(fimage)
|
|
label+=1
|
|
return ret
|
|
|
|
def get_dataset(name, input_dir):
|
|
if name=='webface' or name=='lfw' or name=='vgg':
|
|
return get_dataset_common(input_dir)
|
|
if name=='celeb':
|
|
return get_dataset_celeb(input_dir)
|
|
if name=='facescrub':
|
|
return get_dataset_facescrub(input_dir)
|
|
if name=='megaface':
|
|
return get_dataset_megaface(input_dir)
|
|
return None
|
|
|
|
|