mxnet制作人脸识别训练集

矫情吗；* 2022-09-09 14:59 161阅读 0赞

**目录**

调用数据集

生成train.lst

生成train.rec train.idx

--------------------

h5py还是比较方便的，推荐使用：

[https://blog.csdn.net/jacke121/article/details/119935657][https_blog.csdn.net_jacke121_article_details_119935657]

### 调用数据集 ###

import mxnet as mx
    
    class MXFaceDataset(Dataset):
        def __init__(self, root_dir, local_rank):
            super(MXFaceDataset, self).__init__()
            self.transform = transforms.Compose(
                [transforms.ToPILImage(),
                 transforms.RandomHorizontalFlip(),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
                 ])
            self.root_dir = root_dir
            self.local_rank = local_rank
            path_imgrec = os.path.join(root_dir, 'train.rec')
            path_imgidx = os.path.join(root_dir, 'train.idx')
            self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')
            s = self.imgrec.read_idx(0)
            header, _ = mx.recordio.unpack(s)
            if header.flag > 0:
                self.header0 = (int(header.label[0]), int(header.label[1]))
                self.imgidx = np.array(range(1, int(header.label[0])))
            else:
                self.imgidx = np.array(list(self.imgrec.keys))
    
        def __getitem__(self, index):
            idx = self.imgidx[index]
            s = self.imgrec.read_idx(idx)
            header, img = mx.recordio.unpack(s)
            label = header.label
            if not isinstance(label, numbers.Number):
                label = label[0]
            label = torch.tensor(label, dtype=torch.long)
            sample = mx.image.imdecode(img).asnumpy()
            if self.transform is not None:
                sample = self.transform(sample)
            return sample, label
    
        def __len__(self):
            return len(self.imgidx)

### 生成train.lst ###

import argparse
    import glob
    import os
    import numpy as np
    
    import cv2
    import mxnet as mx
    
    
    def get_id():
    
        path_f=r"G:\data\5w"
    
        # files = glob.glob(path_f + "/*/*/*/*[bmp,jpg,png]", recursive=True)  # find file
        files = glob.glob(path_f + "/*/*.jpg", recursive=True)  # find file
    
        test_lst = r'G:\data\train_5w/train.lst' #rec
        with open(test_lst, 'w') as fw:
            for index, idx in enumerate(files):
                face_id=int(os.path.basename(os.path.dirname(idx)))
                fw.writelines(f'1\t{face_id}\t{idx}\n')
    
    if __name__ == '__main__':
    
        get_id()

### 生成train.rec train.idx ###

im2rec.py

主要设置两个参数:

parser.add_argument('--prefix', help='prefix of input/output lst and rec files.',default=r"G:\data\train_5w\train.lst")
        parser.add_argument('--root', help='path to folder containing images.',default=r"G:\data\5w")

#!/usr/bin/env python3
    # -*- coding: utf-8 -*-
    from __future__ import print_function
    import os
    import sys
    
    curr_path = os.path.abspath(os.path.dirname(__file__))
    sys.path.append(os.path.join(curr_path, "../python"))
    import mxnet as mx
    import random
    import argparse
    import cv2
    import time
    import traceback
    
    try:
        import multiprocessing
    except ImportError:
        multiprocessing = None
    
    
    def list_image(root, recursive, exts):
        """Traverses the root of directory that contains images and
        generates image list iterator.
        Parameters
        ----------
        root: string
        recursive: bool
        exts: string
        Returns
        -------
        image iterator that contains all the image under the specified path
        """
    
        i = 0
        if recursive:
            cat = {}
            for path, dirs, files in os.walk(root, followlinks=True):
                dirs.sort()
                files.sort()
                for fname in files:
                    fpath = os.path.join(path, fname)
                    suffix = os.path.splitext(fname)[1].lower()
                    if os.path.isfile(fpath) and (suffix in exts):
                        if path not in cat:
                            cat[path] = len(cat)
                        yield (i, os.path.relpath(fpath, root), cat[path])
                        i += 1
            for k, v in sorted(cat.items(), key=lambda x: x[1]):
                print(os.path.relpath(k, root), v)
        else:
            for fname in sorted(os.listdir(root)):
                fpath = os.path.join(root, fname)
                suffix = os.path.splitext(fname)[1].lower()
                if os.path.isfile(fpath) and (suffix in exts):
                    yield (i, os.path.relpath(fpath, root), 0)
                    i += 1
    
    
    def write_list(path_out, image_list):
        """Hepler function to write image list into the file.
        The format is as below,
        integer_image_index \t float_label_index \t path_to_image
        Note that the blank between number and tab is only used for readability.
        Parameters
        ----------
        path_out: string
        image_list: list
        """
        with open(path_out, 'w') as fout:
            for i, item in enumerate(image_list):
                line = '%d\t' % item[0]
                for j in item[2:]:
                    line += '%f\t' % j
                line += '%s\n' % item[1]
                fout.write(line)
    
    
    def make_list(args):
        """Generates .lst file.
        Parameters
        ----------
        args: object that contains all the arguments
        """
        image_list = list_image(args.root, args.recursive, args.exts)
        image_list = list(image_list)
        if args.shuffle is True:
            random.seed(100)
            random.shuffle(image_list)
        N = len(image_list)
        chunk_size = (N + args.chunks - 1) // args.chunks
        for i in range(args.chunks):
            chunk = image_list[i * chunk_size:(i + 1) * chunk_size]
            if args.chunks > 1:
                str_chunk = '_%d' % i
            else:
                str_chunk = ''
            sep = int(chunk_size * args.train_ratio)
            sep_test = int(chunk_size * args.test_ratio)
            if args.train_ratio == 1.0:
                write_list(args.prefix + str_chunk + '.lst', chunk)
            else:
                if args.test_ratio:
                    write_list(args.prefix + str_chunk + '_test.lst', chunk[:sep_test])
                if args.train_ratio + args.test_ratio < 1.0:
                    write_list(args.prefix + str_chunk + '_val.lst', chunk[sep_test + sep:])
                write_list(args.prefix + str_chunk + '_train.lst', chunk[sep_test:sep_test + sep])
    
    
    def read_list(path_in):
        """Reads the .lst file and generates corresponding iterator.
        Parameters
        ----------
        path_in: string
        Returns
        -------
        item iterator that contains information in .lst file
        """
        with open(path_in) as fin:
            while True:
                line = fin.readline()
                if not line:
                    break
                line = [i.strip() for i in line.strip().split('\t')]
                line_len = len(line)
                # check the data format of .lst file
                if line_len < 3:
                    print('lst should have at least has three parts, but only has %s parts for %s' % (line_len, line))
                    continue
                try:
                    item = [int(line[0])] + [line[-1]] + [float(i) for i in line[1:-1]]
                except Exception as e:
                    print('Parsing lst met error for %s, detail: %s' % (line, e))
                    continue
                yield item
    
    
    def image_encode(args, i, item, q_out):
        """Reads, preprocesses, packs the image and put it back in output queue.
        Parameters
        ----------
        args: object
        i: int
        item: list
        q_out: queue
        """
        fullpath = os.path.join(args.root, item[1])
    
        if len(item) > 3 and args.pack_label:
            header = mx.recordio.IRHeader(0, item[2:], item[0], 0)
        else:
            header = mx.recordio.IRHeader(0, item[2], item[0], 0)
    
        if args.pass_through:
            try:
                with open(fullpath, 'rb') as fin:
                    img = fin.read()
                s = mx.recordio.pack(header, img)
                q_out.put((i, s, item))
            except Exception as e:
                traceback.print_exc()
                print('pack_img error:', item[1], e)
                q_out.put((i, None, item))
            return
    
        try:
            img = cv2.imread(fullpath, args.color)
        except:
            traceback.print_exc()
            print('imread error trying to load file: %s ' % fullpath)
            q_out.put((i, None, item))
            return
        if img is None:
            print('imread read blank (None) image for file: %s' % fullpath)
            q_out.put((i, None, item))
            return
        if args.center_crop:
            if img.shape[0] > img.shape[1]:
                margin = (img.shape[0] - img.shape[1]) // 2
                img = img[margin:margin + img.shape[1], :]
            else:
                margin = (img.shape[1] - img.shape[0]) // 2
                img = img[:, margin:margin + img.shape[0]]
        if args.resize:
            if img.shape[0] > img.shape[1]:
                newsize = (args.resize, img.shape[0] * args.resize // img.shape[1])
            else:
                newsize = (img.shape[1] * args.resize // img.shape[0], args.resize)
            img = cv2.resize(img, newsize)
    
        try:
            s = mx.recordio.pack_img(header, img, quality=args.quality, img_fmt=args.encoding)
            q_out.put((i, s, item))
        except Exception as e:
            traceback.print_exc()
            print('pack_img error on file: %s' % fullpath, e)
            q_out.put((i, None, item))
            return
    
    
    def read_worker(args, q_in, q_out):
        """Function that will be spawned to fetch the image
        from the input queue and put it back to output queue.
        Parameters
        ----------
        args: object
        q_in: queue
        q_out: queue
        """
        while True:
            deq = q_in.get()
            if deq is None:
                break
            i, item = deq
            image_encode(args, i, item, q_out)
    
    
    def write_worker(q_out, fname, working_dir):
        """Function that will be spawned to fetch processed image
        from the output queue and write to the .rec file.
        Parameters
        ----------
        q_out: queue
        fname: string
        working_dir: string
        """
        pre_time = time.time()
        count = 0
        fname = os.path.basename(fname)
        fname_rec = os.path.splitext(fname)[0] + '.rec'
        fname_idx = os.path.splitext(fname)[0] + '.idx'
        record = mx.recordio.MXIndexedRecordIO(os.path.join(working_dir, fname_idx), os.path.join(working_dir, fname_rec),
                                               'w')
        buf = {}
        more = True
        while more:
            deq = q_out.get()
            if deq is not None:
                i, s, item = deq
                buf[i] = (s, item)
            else:
                more = False
            while count in buf:
                s, item = buf[count]
                del buf[count]
                if s is not None:
                    record.write_idx(item[0], s)
    
                if count % 1000 == 0:
                    cur_time = time.time()
                    print('time:', cur_time - pre_time, ' count:', count)
                    pre_time = cur_time
                count += 1
    
    
    def parse_args():
        parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='Create an image list or \
            make a record database by reading from an image list')
        parser.add_argument('--prefix', help='prefix of input/output lst and rec files.',default=r"G:\data\train_5w\train.lst")
        parser.add_argument('--root', help='path to folder containing images.',default=r"G:\data\5w")
    
        cgroup = parser.add_argument_group('Options for creating image lists')
        cgroup.add_argument('--list', action='store_true', help='If this is set im2rec will create image list(s) by traversing root folder\
            and output to <prefix>.lst.\
            Otherwise im2rec will read <prefix>.lst and create a database at <prefix>.rec')
        cgroup.add_argument('--exts', nargs='+', default=['.jpeg', '.jpg', '.png'],
                            help='list of acceptable image extensions.')
        cgroup.add_argument('--chunks', type=int, default=1, help='number of chunks.')
        cgroup.add_argument('--train-ratio', type=float, default=1.0, help='Ratio of images to use for training.')
        cgroup.add_argument('--test-ratio', type=float, default=0, help='Ratio of images to use for testing.')
        cgroup.add_argument('--recursive', action='store_true', help='If true recursively walk through subdirs and assign an unique label\
            to images in each folder. Otherwise only include images in the root folder\
            and give them label 0.')
        cgroup.add_argument('--no-shuffle', dest='shuffle', action='store_false', help='If this is passed, \
            im2rec will not randomize the image order in <prefix>.lst')
        rgroup = parser.add_argument_group('Options for creating database')
        rgroup.add_argument('--pass-through', action='store_true',
                            help='whether to skip transformation and save image as is')
        rgroup.add_argument('--resize', type=int, default=0, help='resize the shorter edge of image to the newsize, original images will\
            be packed by default.')
        rgroup.add_argument('--center-crop', action='store_true',
                            help='specify whether to crop the center image to make it rectangular.')
        rgroup.add_argument('--quality', type=int, default=95,
                            help='JPEG quality for encoding, 1-100; or PNG compression for encoding, 1-9')
        rgroup.add_argument('--num-thread', type=int, default=1, help='number of thread to use for encoding. order of images will be different\
            from the input list if >1. the input list will be modified to match the\
            resulting order.')
        rgroup.add_argument('--color', type=int, default=1, choices=[-1, 0, 1], help='specify the color mode of the loaded image.\
            1: Loads a color image. Any transparency of image will be neglected. It is the default flag.\
            0: Loads image in grayscale mode.\
            -1:Loads image as such including alpha channel.')
        rgroup.add_argument('--encoding', type=str, default='.jpg', choices=['.jpg', '.png'],
                            help='specify the encoding of the images.')
        rgroup.add_argument('--pack-label', action='store_true',
                            help='Whether to also pack multi dimensional label in the record file')
        args = parser.parse_args()
        args.prefix = os.path.abspath(args.prefix)
        args.root = os.path.abspath(args.root)
        return args
    
    
    if __name__ == '__main__':
        args = parse_args()
        # if the '--list' is used, it generates .lst file
        if args.list:
            make_list(args)
        # otherwise read .lst file to generates .rec file
        else:
            if os.path.isdir(args.prefix):
                working_dir = args.prefix
            else:
                working_dir = os.path.dirname(args.prefix)
            files = [os.path.join(working_dir, fname) for fname in os.listdir(working_dir) if
                     os.path.isfile(os.path.join(working_dir, fname))]
            count = 0
            for fname in files:
                if fname.startswith(args.prefix) and fname.endswith('.lst'):
                    print('Creating .rec file from', fname, 'in', working_dir)
                    count += 1
                    image_list = read_list(fname)
                    # -- write_record -- #
                    if args.num_thread > 1 and multiprocessing is not None:
                        q_in = [multiprocessing.Queue(1024) for i in range(args.num_thread)]
                        q_out = multiprocessing.Queue(1024)
                        # define the process
                        read_process = [multiprocessing.Process(target=read_worker, args=(args, q_in[i], q_out)) for i in
                                        range(args.num_thread)]
                        # process images with num_thread process
                        for p in read_process:
                            p.start()
                        # only use one process to write .rec to avoid race-condtion
                        write_process = multiprocessing.Process(target=write_worker, args=(q_out, fname, working_dir))
                        write_process.start()
                        # put the image list into input queue
                        for i, item in enumerate(image_list):
                            q_in[i % len(q_in)].put((i, item))
                        for q in q_in:
                            q.put(None)
                        for p in read_process:
                            p.join()
    
                        q_out.put(None)
                        write_process.join()
                    else:
                        print('multiprocessing not available, fall back to single threaded encoding')
                        try:
                            import Queue as queue
                        except ImportError:
                            import queue
                        q_out = queue.Queue()
                        fname = os.path.basename(fname)
                        fname_rec = os.path.splitext(fname)[0] + '.rec'
                        fname_idx = os.path.splitext(fname)[0] + '.idx'
                        record = mx.recordio.MXIndexedRecordIO(os.path.join(working_dir, fname_idx),
                                                               os.path.join(working_dir, fname_rec), 'w')
                        cnt = 0
                        pre_time = time.time()
                        for i, item in enumerate(image_list):
                            image_encode(args, i, item, q_out)
                            if q_out.empty():
                                continue
                            _, s, _ = q_out.get()
                            record.write_idx(item[0], s)
                            if cnt % 1000 == 0:
                                cur_time = time.time()
                                print('time:', cur_time - pre_time, ' count:', cnt)
                                pre_time = cur_time
                            cnt += 1
            if not count:
                print('Did not find and list file with prefix %s' % args.prefix)

[https_blog.csdn.net_jacke121_article_details_119935657]: https://blog.csdn.net/jacke121/article/details/119935657