6D pose estimation 之 EfficientPose

github： https://github.com/ybkscht/EfficientPose

新出的多目标6D姿态估计算法EfficientPose，看论文在多目标检测效果确实要比之前的算法要好。

按github上走基本上可以跑通作者预训练的模型。

如果你要自己制作数据集的话主要是完成如下图片中的：

1.rgb 里面存放训练的原图

2.merged_masks 里面放目标mask的label图片（每个目标的rgb值都不一样）

3.gt.yml: 保存 extrinsic matrix 的R 和T以及 bounding box， 和目标id值。 这里的索引和图片是一一对应的

4. info.yaml 保存每个图片的intrinsic matrix，正常每个图片都一样

5.train.txt 训练图片的文件名

6.test.txt 验证集的图片文件名

7.models_info.yaml obj模型文件的参数，最大直径等

6.boj_01.ply mesh文件 按毫米和按米单位不一样需要修改训练文件的translation_scale_norm参数，如果是米为1.0 如果是毫米是1000

数据集生成器可以模仿OcclusionGenerator或LineModGenerator，按照自己的数据集不同去编写数据生成过程，基本上也就改改__init__()方法就可以。最好在create_generators函数中添加自己的args.dataset_type。

如下我自己修改的custom类。然后就可以开始训练了。最好用debug.py跑下自己的数据生成器是不是对的。

if args.dataset_type == 'linemod':
        from generators.linemod import LineModGenerator
        train_generator = LineModGenerator(
            args.linemod_path,
            args.object_id,
            rotation_representation=args.rotation_representation,
            use_colorspace_augmentation=not args.no_color_augmentation,
            use_6DoF_augmentation=not args.no_6dof_augmentation,
            **common_args
        )

        validation_generator = LineModGenerator(
            args.linemod_path,
            args.object_id,
            train=False,
            shuffle_dataset=False,
            shuffle_groups=False,
            rotation_representation=args.rotation_representation,
            use_colorspace_augmentation=False,
            use_6DoF_augmentation=False,
            **common_args
        )
    elif args.dataset_type == 'occlusion':
        from generators.occlusion import OcclusionGenerator
        train_generator = OcclusionGenerator(
            args.occlusion_path,
            rotation_representation=args.rotation_representation,
            use_colorspace_augmentation=not args.no_color_augmentation,
            use_6DoF_augmentation=not args.no_6dof_augmentation,
            **common_args
        )

        validation_generator = OcclusionGenerator(
            args.occlusion_path,
            train=False,
            shuffle_dataset=False,
            shuffle_groups=False,
            rotation_representation=args.rotation_representation,
            use_colorspace_augmentation=False,
            use_6DoF_augmentation=False,
            **common_args
        )
    elif args.dataset_type == 'custom':
        from generators.custom import CustomGenerator
        train_generator = CustomGenerator(
            args.custom_path,
            rotation_representation=args.rotation_representation,
            use_colorspace_augmentation=not args.no_color_augmentation,
            use_6DoF_augmentation=not args.no_6dof_augmentation,
            **common_args
        )

        validation_generator = CustomGenerator(
            args.custom_path,
            train=False,
            shuffle_dataset=False,
            shuffle_groups=False,
            rotation_representation=args.rotation_representation,
            use_colorspace_augmentation=False,
            use_6DoF_augmentation=False,
            **common_args
        )
    else:
        raise ValueError('Invalid data type received: {}'.format(args.dataset_type))

    return train_generator, validation_generator