Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
partial_dataset.py 2.19 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
68
69
70
71
72
73
74 
# encoding: utf-8

"""
@author:  lingxiao he
@contact: helingxiao3@jd.com
"""

import glob
import os
import os.path as osp
import re

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

__all__ = ['PartialREID', 'PartialiLIDS', 'OccludedREID']


def process_test(query_path, gallery_path):
    query_img_paths = glob.glob(os.path.join(query_path, '*.jpg'))
    gallery_img_paths = glob.glob(os.path.join(gallery_path, '*.jpg'))
    query_paths = []
    pattern = re.compile(r'([-\d]+)_(\d*)')
    for img_path in query_img_paths:
        pid, camid = map(int, pattern.search(img_path).groups())
        query_paths.append([img_path, pid, camid])
    gallery_paths = []
    for img_path in gallery_img_paths:
        pid, camid = map(int, pattern.search(img_path).groups())
        gallery_paths.append([img_path, pid, camid])
    return query_paths, gallery_paths


@DATASET_REGISTRY.register()
class PartialREID(ImageDataset):

    dataset_name = "partialreid"

    def __init__(self, root='datasets',):
        self.root = root

        self.query_dir = osp.join(self.root, 'Partial_REID/partial_body_images')
        self.gallery_dir = osp.join(self.root, 'Partial_REID/whole_body_images')
        query, gallery = process_test(self.query_dir, self.gallery_dir)

        ImageDataset.__init__(self, [], query, gallery)


@DATASET_REGISTRY.register()
class PartialiLIDS(ImageDataset):
    dataset_name = "partialilids"

    def __init__(self, root='datasets',):
        self.root = root

        self.query_dir = osp.join(self.root, 'PartialiLIDS/query')
        self.gallery_dir = osp.join(self.root, 'PartialiLIDS/gallery')
        query, gallery = process_test(self.query_dir, self.gallery_dir)

        ImageDataset.__init__(self, [], query, gallery)


@DATASET_REGISTRY.register()
class OccludedREID(ImageDataset):
    dataset_name = "occludereid"

    def __init__(self, root='datasets',):
        self.root = root

        self.query_dir = osp.join(self.root, 'OccludedREID/query')
        self.gallery_dir = osp.join(self.root, 'OccludedREID/gallery')
        query, gallery = process_test(self.query_dir, self.gallery_dir)

        ImageDataset.__init__(self, [], query, gallery)