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 LineModGeneratortrain_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 OcclusionGeneratortrain_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 CustomGeneratortrain_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