icemdata_manager.py

# -*- encoding: utf-8 -*-
"""
@File    : 15616.py
@Time    : 2021-05-14 16:27
@Author  : XD
@Email   : gudianpai@qq.com
@Software: PyCharm
"""

import os.path as osp
import glob

from IPython import embed

class ICME(object):
    dataset_dir = 'icme'

    def __init__(self, root='data', **kwargs):
        self.dataset_dir = osp.join(root, self.dataset_dir)
        self.train_dir = osp.join(self.dataset_dir, 'Train')
        self.test_dir = osp.join(self.dataset_dir, 'Test')

        self._check_before_run()
        train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
        test, num_test_pids, num_test_imgs = self._process_dir(self.test_dir, relabel=False)
        num_total_pids = num_train_pids + num_test_pids
        num_total_imgs = num_train_imgs + num_test_imgs

        print("=> ICME loaded")
        print("Dataset statistics:")
        print("  ------------------------------")
        print("  subset   | # ids | # images")
        print("  ------------------------------")
        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
        print("  test    | {:5d} | {:8d}".format(num_test_pids, num_test_imgs))
        print("  ------------------------------")
        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
        print("  ------------------------------")

        self.train = train
        self.test = test

        self.num_train_pids = num_train_pids
        self.num_test_pids = num_test_pids

    def _check_before_run(self):
        if not osp.exists(self.train_dir):
            raise RuntimeError("'{}' is not available".format(self.train_dir))
        if not osp.exists(self.test_dir):
            raise RuntimeError("'{}' is not available".format(self.test_dir))

    def _process_dir(self, dir_path, relabel=False):
        img_paths = glob.glob(osp.join(dir_path, '*/*/*'))

        pid_container = set()

        for img_path in img_paths:
            pid = img_path.split('\\')[-3]
            pid_container.add(pid)
        pid2label = {pid: label for label, pid in enumerate(pid_container)}

        dataset = []
        for img_path in img_paths:
            pid = img_path.split('\\')[-3]
            if relabel:
                pid = pid2label[pid]
            dataset.append((img_path, pid))

        num_pids = len(pid_container)
        num_imgs = len(img_paths)
        return dataset, num_pids, num_imgs

if __name__  == '__main__':
    data = ICME(root='G:/data')

D:\ANACONDA\envs\pytorch_gpu\python.exe G:/图像检索文章/深度哈希/度量学习/proj_Reid/15616.py
=> ICME loaded
Dataset statistics:
  ------------------------------
  subset   | # ids | # images
  ------------------------------
  train    |    44 |     7040
  test    |    15 |     2901
  ------------------------------
  total    |    59 |     9941
  ------------------------------

Process finished with exit code 0