DETR训练实践以及自动预标注脚本、测试可视化脚本

目录自动预标注脚本测试可视化脚本DETR训练自己的数据集自动预标注脚本总所周知标注是一个体力活而我又是一个懒人所以我就自己写了个脚本。。。并不要求高精度只是想体验用DETR训练自己的数据集又不想标注的人可以借鉴或使用。输入方式可以调用摄像头或者读取视频、图片三种方式来获取输入我加了print提示可以直接运行查看哈。输入012来选择功能分为3种查看原图检测可视化图像标注由于这次我是用我的人脸做了一个数据集所以取名face detection。如果选择image label接下来是选择数据集类型输入对应的字符例如test就可以开始了。其他的功能可以读者自己慢慢摸索一下。有需要的或者意见的也可以直接在评论区提出来我看见了也会及时回复。这里就不说太多了。输出结果测试可视化脚本因为论文源代码里是没有test脚本的所以我就直接写了一个test脚本包括了可视化结果其实和源码是同样的架构搭建模型后load训练好的参数构建测试数据集获取测试输出后用后处理获得高得分的目标然后可视化出来有需要的也可以评论需要哈。这里我就只把主要函数贴上来了def main(dataset_type): device torch.device(cuda) model build_model() model.to(device) #dataset dataset_test build_dataset(dataset_type) sampler_test torch.utils.data.SequentialSampler(dataset_test) data_loader_test DataLoader(dataset_test, 1, samplersampler_test, drop_lastFalse, collate_fnutils.collate_fn, num_workers0) base_ds get_coco_api_from_dataset(dataset_test) #load model params load_path D:\\fgy\\Desktop\\DETR\\detr\\pthsave\\pthsave_facedetect\\checkpoint.pth checkpoint torch.load(load_path, map_locationcpu) model.load_state_dict(checkpoint[model],False) model.eval() #postprocess postprocessors {bbox: detr.PostProcess()} iou_types tuple(k for k in (segm, bbox) if k in postprocessors.keys()) coco_evaluator coco_eval.CocoEvaluator(base_ds, iou_types) image_path D:\\fgy\\Desktop\\DETR\\dataroot\\{}.format(dataset_type) #get result and visual for img_data, target in data_loader_test: img_data img_data.to(device) target [{k: v.to(device) for k, v in t.items()} for t in target] output model(img_data) # print(output) orig_target_sizes torch.stack([t[orig_size] for t in target], dim0) result postprocessors[bbox](output, orig_target_sizes) res {target[image_id].item(): output for target, output in zip(target, result)} if coco_evaluator is not None: coco_evaluator.update(res) # print(res) res_index [] score[] min_score 0.995 res_lable[] res_bbox[] res res[target[0][image_id].item()] for i in range(0,100): res_tmp res[scores] if float(res_tmp[i]) min_score: score.append(float(res_tmp[i])) res_index.append(i) res_lable.append(int(res[labels][res_index].cpu().numpy())) res_bbox.append(res[boxes][res_index].cpu().numpy().tolist()) print(processing str(target[0][image_id].item())) print(result: ,score, res_lable, res_bbox) img cv.imread(os.path.join(image_path,str(target[0][image_id].item()).jpg)) for bbox in res_bbox: bbox bbox[0] cv.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 0), 1) cv.imshow(Faceimglabel, img) cv.waitKey(0) if coco_evaluator is not None: coco_evaluator.synchronize_between_processes() coco_evaluator.accumulate() coco_evaluator.summarize() print(coco_evaluator)DETR训练自己的数据集开始之前对DETR或源码不了解的建议先看一下DETR源码笔记一_在努力的松鼠的博客-CSDN博客_detr源码DETR源码笔记二_在努力的松鼠的博客-CSDN博客一、首先是搭建自己的数据集标注完成后查看数据集的类别数N其中不包括背景类。二、可以初始化参数训练也可以下载预训练模型建议下载一个预训练模型提高训练速度。下载地址可以在https://gitee.com/fgy120/DETR的ReadMe里的Modelzoo查找自己想要训练的网络结构的预训练模型下载地址。三、没有下载预训练模型跳过这一步下载好预训练模型后我们需要把模型的head输出数改为我们自己的数据集类别数N可以用这个脚本注意里面需要修改的class不包含背景类。https://gitee.com/fgy120/DETR/blob/master/util/change_pth_numclass.py四、预训练模型准备好后就是配置我们的agrs——运行配置和训练的超参等在main.py的get_args_parser函数里关于parser的使用可以看一下argparse解析器_在努力的松鼠的博客-CSDN博客因为我喜欢pycharm直接运行 所以我直接给很多的arg都加上了default默认值懒人实锤了。。。值得注意要修改的几个argbatch_size如果太高了显卡带不动会显示爆显存的错就需要降低我们的batch_size了我3060显460*680的图像设了4感觉刚刚好大家量力而行。epochs: 训练的轮次所有图片训练一次是一轮。backbonebackbone的网络结构选择我选择了resnet50coco_path是你存放数据集ann、train、test、val文件夹的文件夹路径output_dir是你训练以及验证模型保存的文件夹路径world_size一般别超过1是分布式训练需要用到的设置resume预处理模型以及断点训练的load模型的地址下面可以参考# parser.add_argument(--resume, defaultD:\\fgy\\Desktop\\DETR\\detr\\pthsave\\pthsave_facedetect\\checkpoint.pth, helpresume from checkpoint) parser.add_argument(--resume, defaultD:\\fgy\\Desktop\\DETR\\savepth\\detr_r50init_1.pth, helpresume from checkpoint)五、接下来是在代码里修改我们的数据集类别数在detr.py的build()函数为了方便我都修改成了我的类别数数据集也懒得弄了就弄了一个类别。。六、上面的都弄好后运行main.py就可以了。上面提到的代码和脚本都可以在https://gitee.com/fgy120/DETR找到有问题或有建议的也可以评论区交流