Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_dataset.py 1.42 KB
Newer Older
dengjb's avatar
update  
dengjb committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67 
# encoding: utf-8
"""
@author:  xingyu liao
@contact: sherlockliao01@gmail.com
"""

import os

import bcolz
import numpy as np

from fastreid.data.datasets import DATASET_REGISTRY
from fastreid.data.datasets.bases import ImageDataset

__all__ = ["CPLFW", "VGG2_FP", "AgeDB_30", "CALFW", "CFP_FF", "CFP_FP", "LFW"]


@DATASET_REGISTRY.register()
class CPLFW(ImageDataset):
    dataset_dir = "faces_emore_val"
    dataset_name = "cplfw"

    def __init__(self, root='datasets', **kwargs):
        self.root = root
        self.dataset_dir = os.path.join(self.root, self.dataset_dir)

        required_files = [self.dataset_dir]

        self.check_before_run(required_files)

        carray = bcolz.carray(rootdir=os.path.join(self.dataset_dir, self.dataset_name), mode='r')
        is_same = np.load(os.path.join(self.dataset_dir, "{}_list.npy".format(self.dataset_name)))

        self.carray = carray
        self.is_same = is_same

        super().__init__([], [], [], **kwargs)


@DATASET_REGISTRY.register()
class VGG2_FP(CPLFW):
    dataset_name = "vgg2_fp"


@DATASET_REGISTRY.register()
class AgeDB_30(CPLFW):
    dataset_name = "agedb_30"


@DATASET_REGISTRY.register()
class CALFW(CPLFW):
    dataset_name = "calfw"


@DATASET_REGISTRY.register()
class CFP_FF(CPLFW):
    dataset_name = "cfp_ff"


@DATASET_REGISTRY.register()
class CFP_FP(CPLFW):
    dataset_name = "cfp_fp"


@DATASET_REGISTRY.register()
class LFW(CPLFW):
    dataset_name = "lfw"