从无到有跑通KAPAO

kapao工程：https://github.com/wmcnally/kapao

论文地址：https://arxiv.org/abs/2111.08557

由于之前没有接触过关键点的工作，所以我选择先下载作者提到的数据，将数据按照他要求的方式存放，然后通过运行的代码来查看整个工程的结果；

使用kapao进行托盘目标检测和托盘目标的四个角点的检测；原理都还没有来得及查看，然后直接进行模型训练；

处理过程的整体流程：

1.将整个样本分成训练集和验证集合，通过txt文件，记载相关的的名字；

"""
将所有样本分拣为训练样本和测试样本
"""
import argparse
import os
import math
import random
import shutil
from tqdm import tqdm
 
class Process:
    def __init__(self,args):
        self.imgs_dir = args.imgs_dir
        self.save_dir = args.save_dir
        self.ratio = args.ratio
 
    def check_path(self, chpath):
        if not os.path.exists(chpath):
            os.mkdir(chpath)
 
    def check_path_(self, chpath):
        if not os.path.exists(chpath):
            os.mkdir(chpath)
        else:
            paths = os.listdir(chpath)
            for path in paths:
                del_path = os.path.join(chpath, path)
                os.remove(del_path)
 
    def process_data(self):
        """
        遍历所有文件夹随机挑取一定比例的样本作为测试样本，剩余的作为训练样本
        """
        paths = os.listdir(self.imgs_dir)
        txt_train = os.path.join(self.save_dir, 'trainval.txt')
        txt_test = os.path.join(self.save_dir, 'test.txt')
        f_train = open(txt_train, 'w', encoding='utf-8')
        f_test = open(txt_test, 'w', encoding='utf-8')
 
        num_paths = len(paths)
        test_num = math.floor(self.ratio*num_paths)
        list2 = [i for i in range(num_paths)]
        dic_num_path = {str(i): path for (i, path) in zip(range(num_paths), paths)}
        choice_num_list = random.sample(list2, test_num)
        for num_i, path in dic_num_path.items():
            img_path = dic_num_path[str(num_i)]
            old_img_path = os.path.join(self.imgs_dir, img_path)
            if int(num_i) in choice_num_list:
                f_test.write(old_img_path+'\n')
            else:
                f_train.write(old_img_path+'\n')
 
        f_test.close()
        f_train.close()
 
 
 
 
 
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgs_dir',
                        default='/home/XXXX/4tdisk/data/key_point/tray/ori/1672/imageset',
                        help='图片文件所在文件夹')
    parser.add_argument('--save_dir',
                        default='/home/XXXX/4tdisk/data/key_point/tray/process_ok/1672',
                        help='结果存储路径'
                        )
    parser.add_argument('--ratio',
                        default=0.2,
                        help='测试样本占所有图片的比例'
                        )
 
    return parser.parse_args()
 
def main(args):
    G = Process(args)
    G.process_data()
 
 
if __name__ == '__main__':
    main(parse_args())