【DETR】训练自己的数据集-实践笔记

这篇具有很好参考价值的文章主要介绍了【DETR】训练自己的数据集-实践笔记。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

DETR(Detection with TRansformers)训练自己的数据集-实践笔记&问题总结

DETR(Detection with TRansformers)是基于transformer的端对端目标检测,无NMS后处理步骤,无anchor。
实现使用NWPUVHR10数据集训练DETR.
NWPU数据集总共包含十种类别目标,包含650个正样本,150个负样本(没有用到)。

NWPU_CATEGORIES=['airplane','ship','storage tank','baseball diamond','tennis court',\
					'basketball court','ground track field','harbor','bridge','vehicle']

代码:https://github.com/facebookresearch/detr

一.训练

1.数据集准备

DETR需要的数据集格式为coco格式,图片和标签文件保存于训练集、测试集、验证集、标签文件四个文件夹中,其中annotations中存放json格式的标签文件
【DETR】训练自己的数据集-实践笔记

【DETR】训练自己的数据集-实践笔记

下面的代码包含了几种数据集RSOD、NWPU、DIOR、YOLO数据集标签文件转换json功能。新建py文件tojson.py,使用如下代码生成需要的json文件。

生成instances_train2017.json
(a)修改29行image_path默认路径为train2017的路径;
(b)修改31行annotation_path默认路径为标签文件路径(train和val的标签都放在这个文件夹下,所以生成instances_val2017.json时就不需要再修改这个路径了);
(c)修改33行dataset为自己的数据集名称NWPU
(d).修改34行save的默认路径为json文件的保存路径…/NWPUVHR-10/annotations/instances_train2017.json

import os
import cv2
import json
import argparse
from tqdm import tqdm
import xml.etree.ElementTree as ET

COCO_DICT=['images','annotations','categories']
IMAGES_DICT=['file_name','height','width','id']

ANNOTATIONS_DICT=['image_id','iscrowd','area','bbox','category_id','id']

CATEGORIES_DICT=['id','name']
## {'supercategory': 'person', 'id': 1, 'name': 'person'}
## {'supercategory': 'vehicle', 'id': 2, 'name': 'bicycle'}
YOLO_CATEGORIES=['person']
RSOD_CATEGORIES=['aircraft','playground','overpass','oiltank']
NWPU_CATEGORIES=['airplane','ship','storage tank','baseball diamond','tennis court',\
					'basketball court','ground track field','harbor','bridge','vehicle']

VOC_CATEGORIES=['aeroplane','bicycle','bird','boat','bottle','bus','car','cat','chair','cow',\					'diningtable','dog','horse','motorbike','person','pottedplant','sheep','sofa','train','tvmonitor']

DIOR_CATEGORIES=['golffield','Expressway-toll-station','vehicle','trainstation','chimney','storagetank',\
					'ship','harbor','airplane','groundtrackfield','tenniscourt','dam','basketballcourt',\
					'Expressway-Service-area','stadium','airport','baseballfield','bridge','windmill','overpass']

parser=argparse.ArgumentParser(description='2COCO')
#parser.add_argument('--image_path',type=str,default=r'T:/shujuji/DIOR/JPEGImages-trainval/',help='config file')
parser.add_argument('--image_path',type=str,default=r'G:/NWPU VHR-10 dataset/positive image set/',help='config file')
#parser.add_argument('--annotation_path',type=str,default=r'T:/shujuji/DIOR/Annotations/',help='config file')
parser.add_argument('--annotation_path',type=str,default=r'G:/NWPU VHR-10 dataset/ground truth/',help='config file')
parser.add_argument('--dataset',type=str,default='NWPU',help='config file')
parser.add_argument('--save',type=str,default='G:/NWPU VHR-10 dataset/instances_train2017.json',help='config file')
args=parser.parse_args()
def load_json(path):
	with open(path,'r') as f:
		json_dict=json.load(f)
		for i in json_dict:
			print(i)
		print(json_dict['annotations'])
def save_json(dict,path):
	print('SAVE_JSON...')
	with open(path,'w') as f:
		json.dump(dict,f)
	print('SUCCESSFUL_SAVE_JSON:',path)
def load_image(path):
	img=cv2.imread(path)
	return img.shape[0],img.shape[1]
def generate_categories_dict(category):       #ANNOTATIONS_DICT=['image_id','iscrowd','area','bbox','category_id','id']
	print('GENERATE_CATEGORIES_DICT...')
	return [{CATEGORIES_DICT[0]:category.index(x)+1,CATEGORIES_DICT[1]:x} for x in category]  #CATEGORIES_DICT=['id','name']
def generate_images_dict(imagelist,image_path,start_image_id=11725):  #IMAGES_DICT=['file_name','height','width','id']
	print('GENERATE_IMAGES_DICT...')
	images_dict=[]
	with tqdm(total=len(imagelist)) as load_bar:
		for x in imagelist:  #x就是图片的名称
			#print(start_image_id)
			dict={IMAGES_DICT[0]:x,IMAGES_DICT[1]:load_image(image_path+x)[0],\
					IMAGES_DICT[2]:load_image(image_path+x)[1],IMAGES_DICT[3]:imagelist.index(x)+start_image_id}
			load_bar.update(1)
			images_dict.append(dict)
	return images_dict

def DIOR_Dataset(image_path,annotation_path,start_image_id=11725,start_id=0):
	categories_dict=generate_categories_dict(DIOR_CATEGORIES)    #CATEGORIES_DICT=['id':,1'name':golffield......]  id从1开始
	imgname=os.listdir(image_path)
	images_dict=generate_images_dict(imgname,image_path,start_image_id)  #IMAGES_DICT=['file_name','height','width','id']  id从0开始的
	print('GENERATE_ANNOTATIONS_DICT...')  #生成cooc的注记   ANNOTATIONS_DICT=['image_id','iscrowd','area','bbox','category_id','id']
	annotations_dict=[]
	id=start_id
	for i in images_dict:
		image_id=i['id']
		print(image_id)
		image_name=i['file_name']
		annotation_xml=annotation_path+image_name.split('.')[0]+'.xml'
		tree=ET.parse(annotation_xml)
		root=tree.getroot()
		for j in root.findall('object'):
			category=j.find('name').text
			category_id=DIOR_CATEGORIES.index(category)  #字典的索引,是从1开始的
			x_min=float(j.find('bndbox').find('xmin').text)
			y_min=float(j.find('bndbox').find('ymin').text)
			w=float(j.find('bndbox').find('xmax').text)-x_min
			h=float(j.find('bndbox').find('ymax').text)-y_min
			area = w * h
			bbox = [x_min, y_min, w, h]
			dict = {'image_id': image_id, 'iscrowd': 0, 'area': area, 'bbox': bbox, 'category_id': category_id,
					'id': id}
			annotations_dict.append(dict)
			id=id+1
	print('SUCCESSFUL_GENERATE_DIOR_JSON')
	return {COCO_DICT[0]:images_dict,COCO_DICT[1]:annotations_dict,COCO_DICT[2]:categories_dict}
def NWPU_Dataset(image_path,annotation_path,start_image_id=0,start_id=0):
	categories_dict=generate_categories_dict(NWPU_CATEGORIES)
	imgname=os.listdir(image_path)
	images_dict=generate_images_dict(imgname,image_path,start_image_id)
	print('GENERATE_ANNOTATIONS_DICT...')
	annotations_dict=[]
	id=start_id
	for i in images_dict:
		image_id=i['id']
		image_name=i['file_name']
		annotation_txt=annotation_path+image_name.split('.')[0]+'.txt'
		txt=open(annotation_txt,'r')
		lines=txt.readlines()
		for j in lines:
			if j=='\n':
				continue
			category_id=int(j.split(',')[4])
			category=NWPU_CATEGORIES[category_id-1]
			print(category_id,'        ',category)
			x_min=float(j.split(',')[0].split('(')[1])
			y_min=float(j.split(',')[1].split(')')[0])
			w=float(j.split(',')[2].split('(')[1])-x_min
			h=float(j.split(',')[3].split(')')[0])-y_min
			area=w*h
			bbox=[x_min,y_min,w,h]
			dict = {'image_id': image_id, 'iscrowd': 0, 'area': area, 'bbox': bbox, 'category_id': category_id,
					'id': id}
			id=id+1
			annotations_dict.append(dict)
	print('SUCCESSFUL_GENERATE_NWPU_JSON')
	return {COCO_DICT[0]:images_dict,COCO_DICT[1]:annotations_dict,COCO_DICT[2]:categories_dict}

def YOLO_Dataset(image_path,annotation_path,start_image_id=0,start_id=0):
	categories_dict=generate_categories_dict(YOLO_CATEGORIES)
	imgname=os.listdir(image_path)
	images_dict=generate_images_dict(imgname,image_path)
	print('GENERATE_ANNOTATIONS_DICT...')
	annotations_dict=[]
	id=start_id
	for i in images_dict:
		image_id=i['id']
		image_name=i['file_name']
		W,H=i['width'],i['height']
		annotation_txt=annotation_path+image_name.split('.')[0]+'.txt'
		txt=open(annotation_txt,'r')
		lines=txt.readlines()
		for j in lines:
			category_id=int(j.split(' ')[0])+1
			category=YOLO_CATEGORIES
			x=float(j.split(' ')[1])
			y=float(j.split(' ')[2])
			w=float(j.split(' ')[3])
			h=float(j.split(' ')[4])
			x_min=(x-w/2)*W
			y_min=(y-h/2)*H
			w=w*W
			h=h*H
			area=w*h
			bbox=[x_min,y_min,w,h]
			dict={'image_id':image_id,'iscrowd':0,'area':area,'bbox':bbox,'category_id':category_id,'id':id}
			annotations_dict.append(dict)
			id=id+1
	print('SUCCESSFUL_GENERATE_YOLO_JSON')
	return {COCO_DICT[0]:images_dict,COCO_DICT[1]:annotations_dict,COCO_DICT[2]:categories_dict}
def RSOD_Dataset(image_path,annotation_path,start_image_id=0,start_id=0):
	categories_dict=generate_categories_dict(RSOD_CATEGORIES)
	imgname=os.listdir(image_path)
	images_dict=generate_images_dict(imgname,image_path,start_image_id)
	print('GENERATE_ANNOTATIONS_DICT...')
	annotations_dict=[]
	id=start_id
	for i in images_dict:
		image_id=i['id']
		image_name=i['file_name']
		annotation_txt=annotation_path+image_name.split('.')[0]+'.txt'
		txt=open(annotation_txt,'r')
		lines=txt.readlines()
		for j in lines:
			category=j.split('\t')[1]
			category_id=RSOD_CATEGORIES.index(category)+1
			x_min=float(j.split('\t')[2])
			y_min=float(j.split('\t')[3])
			w=float(j.split('\t')[4])-x_min
			h=float(j.split('\t')[5])-y_min
			area = w * h
			bbox = [x_min, y_min, w, h]
			dict = {'image_id': image_id, 'iscrowd': 0, 'area': area, 'bbox': bbox, 'category_id': category_id,
					'id': id}
			annotations_dict.append(dict)
			id=id+1
	print('SUCCESSFUL_GENERATE_RSOD_JSON')

	return {COCO_DICT[0]:images_dict,COCO_DICT[1]:annotations_dict,COCO_DICT[2]:categories_dict}
if __name__=='__main__':
	dataset=args.dataset   #数据集名字
	save=args.save  #json的保存路径
	image_path=args.image_path     #对于coco是图片的路径
	annotation_path=args.annotation_path   #coco的annotation路径
	if dataset=='RSOD':
		json_dict=RSOD_Dataset(image_path,annotation_path,0)
	if dataset=='NWPU':
		json_dict=NWPU_Dataset(image_path,annotation_path,0)
	if dataset=='DIOR':
		json_dict=DIOR_Dataset(image_path,annotation_path,11725)
	if dataset=='YOLO':
		json_dict=YOLO_Dataset(image_path,annotation_path,0)
	save_json(json_dict,save)

运行生成instances_train2017.json,再修改路径生成instances_train2017.json。
(经评论区小伙伴提醒发现有问题,已修改)

如果自己的数据集是voc格式,可以使用参考链接1中的大大给出的代码。

1.1 制作好的NWPUVHR数据集下载

有小伙伴要制作好的NWPUVHR数据集,这儿也给出NWPU网盘下载,密码7zf9。其中包含650张影像以8:2比例随机划分好的train2017,val2017图像文件夹,Annotations标签文件夹中包括对应的train.json,val.json标签,数据集原始txt格式标签也包含其中,在ground truth中。

2.环境配置

激活当前项目所在环境,使用如下命令完成环境配置:

pip install -r requirements.txt

3.pth文件生成

先下载预训练文件,官方提供了 DETR 和 DETR-DC5 models,两个模型,后者是在主干网络第五层使用空洞卷积,选择一个下载

新建py文件,mydataset.py,使用如下代码,修改num_classes为自己的类别数+1

import torch
pretrained_weights  = torch.load('detr-r50-e632da11.pth')

#NWPU数据集,10类
num_class = 11    #类别数+1,1为背景
pretrained_weights["model"]["class_embed.weight"].resize_(num_class+1, 256)
pretrained_weights["model"]["class_embed.bias"].resize_(num_class+1)
torch.save(pretrained_weights, "detr-r50_%d.pth"%num_class)

运行,生成detr-r50_11.pth

4.参数修改

修改models/detr.py文件,build()函数中,可以将红框部分的代码都注释掉,直接设置num_classes为自己的类别数+1
【DETR】训练自己的数据集-实践笔记

5.训练

进行训练需要设置一些参数,可以在main.py文件中直接修改或者使用命令行

【a】main.py文件直接修改
修改main.py文件的epochs、lr、batch_size等训练参数
【DETR】训练自己的数据集-实践笔记
修改自己的数据集路径:
【DETR】训练自己的数据集-实践笔记

设置输出路径:
【DETR】训练自己的数据集-实践笔记
修改resume为自己的预训练权重文件路径
【DETR】训练自己的数据集-实践笔记
【b】命令行

python main.py --dataset_file "coco" --coco_path "/home/NWPUVHR-10" --epoch 300 --lr=1e-4 --batch_size=8 --num_workers=4 --output_dir="outputs" --resume="detr_r50_11.pth"

运行main.py文件

二. 出现的bug

1. KeyError: ‘area’

area = torch.tensor([obj[“area”] for obj in anno])
KeyError: ‘area’
【DETR】训练自己的数据集-实践笔记
由定位到的语句可知,应该是字典obj没有名为area的key,查看了一下我的标签文件,确实生成的有问题,里面没有area信息,重新生成了一下,成功运行。
【DETR】训练自己的数据集-实践笔记

三.评估及预测

1.评估

其实训练过程中,DETR在每个epoch都会自动进行一次精度评估,评估结果查看直接输出或生成的中间文件就行,也可以再用main.py评估一下精度,将下列代码中的–resume改为最后一个epoch输出的模型路径,–coco_path改为自己的数据集路径即可

python main.py --batch_size 6 --no_aux_loss --eval --resume /home/detr-main/outputs/checkpoint0299.pth --coco_path /home/NWPUVHR-10

精度评估结果:
训练了300个epoch,IoU=0.5的AP为88%左右
【DETR】训练自己的数据集-实践笔记

2.预测

使用参考链接3中的代码,将要预测的图片保存在一个文件夹下,预测时一次输出所有图片的预测结果

需要修改的参数有:

backbone,我当时下载的backbone是resnet50,修改(训练时已经下载好了主干特征网络是Resnet50的DETR权重文件,放在主文件夹下)
【DETR】训练自己的数据集-实践笔记
数据集有关参数
–coco_path 修改为自己的数据集路径
–outputdir 修改为建立的预测图片的保存文件夹
–resume 修改为训练好的模型文件路径
【DETR】训练自己的数据集-实践笔记
修改待预测的图片文件夹路径image_file_path和image_path
【DETR】训练自己的数据集-实践笔记

ps:由于我用服务器跑,无法传回图片而出现一个报错,于是把这两句注释掉了:

#cv2.imshow("images",image)
#cv2.waitKey(1)

有可能会报一个错,这个是你的某张图片预测时没有检测出目标导致的:
【DETR】训练自己的数据集-实践笔记
预测结果:
【DETR】训练自己的数据集-实践笔记

参考:
1.windows10复现DEtection TRansformers(DETR
2.如何用DETR(detection transformer)训练自己的数据集
3.pytorch实现DETR的推理程序文章来源地址https://www.toymoban.com/news/detail-454225.html

到了这里,关于【DETR】训练自己的数据集-实践笔记的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • DETR-《End-to-End Object Detection with Transformers》论文精读笔记

    End-to-End Object Detection with Transformers 参考:跟着李沐学AI-DETR 论文精读【论文精读】 在摘要部分作者,主要说明了如下几点: DETR是一个端到端(end-to-end)框架,释放了传统基于CNN框架的一阶段(YOLO等)、二阶段(FasterRCNN等)目标检测器中需要大量的人工参与的步骤,例如:

    2024年02月11日
    浏览(51)
  • DEFORMABLE DETR: DEFORMABLE TRANSFORMERS FOR END-TO-END OBJECT DETECTION 论文精读笔记

    DEFORMABLE DETR: DEFORMABLE TRANSFORMERS FOR END-TO-END OBJECT DETECTION 参考:AI-杂货铺-Transformer跨界CV又一佳作!Deformable DETR:超强的小目标检测算法! 摘要 摘要部分,作者主要说明了如下几点: 为了解决DETR中使用Transformer架构在处理图像特征图时的局限性而导致的收敛速度慢,特征空间

    2024年02月10日
    浏览(38)
  • YOLOv5入门实践(4)——手把手教你训练自己的数据集

      在上一篇文章中我们介绍了如何划分数据集,划分好之后我们的前期准备工作就已经全部完成了,下面开始训练自己的数据集吧! 前期回顾: YOLOv5入门实践(1)——手把手带你环境配置搭建 YOLOv5入门实践(2)——手把手教你利用labelimg标注数据集

    2024年04月10日
    浏览(60)
  • DETR (DEtection TRansformer)基于自建数据集开发构建目标检测模型超详细教程

    目标检测系列的算法模型可以说是五花八门,不同的系列有不同的理论依据,DETR的亮点在于它是完全端到端的第一个目标检测模型,DETR(Detection Transformer)是一种基于Transformer的目标检测模型,由Facebook AI Research开发。它采用了端到端的方法,在目标检测任务中同时完成目标

    2024年02月17日
    浏览(42)
  • RT-DETR论文阅读笔记(包括YOLO版本训练和官方版本训练)

    论文地址: RT-DETR论文地址 代码地址: RT-DETR官方下载地址 大家如果想看更详细训练、推理、部署、验证等教程可以看我的另一篇博客里面有更详细的介绍 内容回顾: 详解RT-DETR网络结构/数据集获取/环境搭建/训练/推理/验证/导出/部署  目录 一、介绍  二、相关工作 2.1、实

    2024年02月03日
    浏览(40)
  • YOLOv7训练自己的数据集(txt文件,笔记)

    目录 1.代码下载 2.数据集准备(.xml转.txt) (1)修改图像文件名 (2)图片和标签文件数量不对应,解决办法 (3).xml转.txt (4).txt文件随机划分出对应的训练集、测试集、验证集 3.训练数据集 (1)修改.yaml文件  (2)修改网络参数  (3)训练中断 论文地址: https://arxiv.

    2024年02月02日
    浏览(42)
  • AI学习笔记四:yolov5训练自己的数据集

    若该文为原创文章,转载请注明原文出处。 一般情况下,大部分人的电脑都是没有cpu的,cpu也是可以训练的,但花费的时间太长,实际200张图片,使用CPU训练300轮花了3天,本章记录使用云服务器来训练自己的数据集。 使用的云服务器是AutoDL,一直在使用,性价比还是比较高的

    2024年02月15日
    浏览(50)
  • 训练自己的ai模型(四)学习笔记与项目实操(什么也不懂,但有数据,怎么搞?无监督学习算法)

    很开心有人还在催更,有点小震惊吧。 (原来真有人在csdn发学习记录啊) (原来真有人在csdn看学习记录啊) ai模型方向的知识,我也在学习中,可能疑惑不比大家少。 直接开始! 不管你的是什么数据,只要你有数据,你就可以试一试,跑一跑。 使用 无监督学习算法 。

    2024年02月07日
    浏览(68)
  • 基于DETR (DEtection TRansformer)开发构建MSTAR雷达影像目标检测系统

    关于DETR相关的实践在之前的文章中很详细地介绍过,感兴趣的话可以自行移步阅读即可: 《DETR (DEtection TRansformer)基于自建数据集开发构建目标检测模型超详细教程》 《书接上文——DETR评估可视化》 基于MSTAR雷达影像数据开发构建目标检测系统,在我前面的文章中也有过实

    2024年02月13日
    浏览(66)
  • 目标检测——detr源码复现【 End-to-End Object Detection with Transformers】

    detr源码地址 detr论文地址 自定义coco数据集 在github上面下载 链接:https://pan.baidu.com/s/1fmOYAOZ4yYx_rYquOS6Ycw 提取码:74l5 生成自己所需要的权重文件 main.py 相应位置根据下图更改 model 目录下面的 detr.py 文件相应位置更改类别 num_classes detr的测试对于小物体的检测不是很好,相比来

    2024年02月16日
    浏览(45)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包