使用mxnet中的img2rec.py制作rec数据集

news2024/11/26 12:52:19

源码链接：mxnet/tools/im2rec.py at master · apache/mxnet · GitHub

重点关注入参函数即可，

def parse_args():
    """Defines all arguments.
    Returns
    -------
    args object that contains all the params
    """
    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.')
    parser.add_argument('root', help='path to folder containing images.')

    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

prefix：生成的文件（ <prefix>.lst、 <prefix>.rec）的前缀
root：图像所在文件夹的路径
--list：如果设置了这个参数，im2rec将通过遍历根文件夹创建图像列表，并将其输出到 <prefix>.lst文件中。否则，im2rec将读取<prefix>.lst并在<prefix>.rec创建数据库
--exts：指定可接受的图像文件扩展名列表，默认为['.jpeg', '.jpg', '.png']
--chunks：指定要分割的块数
--train-ratio：用于指定用于训练的图像比例
--test-ratio：用于指定用于测试的图像比例
--recursive：如果设置为True，则递归遍历子目录并为每个文件夹中的图像分配唯一标签。否则，只包括根目录中的图像并将它们标记为0
--no-shuffle：如果设置了这个参数，im2rec将不会对<prefix>.lst中的图像顺序进行随机化
--pass-through：是否跳过转换并将图像保存为原始格式
--resize：将图像的较短边缩放到指定大小，原始图像将默认被打包
--center-crop：指定是否裁剪中心图像以使其成为矩形
--quality：JPEG编码的质量（1-100）或PNG压缩（1-9）
--num-thread：用于编码的线程数。如果大于1，图像顺序将不同于输入列表
--color：加载图像的颜色模式，1：彩色，0：灰度，-1：包括alpha 通道
--encoding：指定图像的编码格式，可以是.jpg或.png
--pack-label：是否在记录文件中打包多维标签

了解了各入参的作用后，便可使用该脚本，

## 先生成lst文件
python mxnet/tools/im2rec.py train "Your WebFace42M Root" --list --recursive

## 再生成对应的rec和idx文件
python mxnet/tools/im2rec.py train "Your WebFace42M Root" --num-thread 16 --quality 100

生成的lst文件示例如下，