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())</