DETR目标检测模型训练自己的数据集

前言

基础环境：ubuntu20.04、python=3.8、pytorch:1.10.0、CUDA:11.3
代码地址：https://github.com/facebookresearch/detr

一、训练准备

1、预训练模型下载

下载地址：https://github.com/facebookresearch/detr?tab=readme-ov-file

在这里插入图片描述

下载后放到detr目录下

在这里插入图片描述

2、txt文件转为coco模式

根据图片和txt标签文件生成json数据，命名分别为instances_train2017.json和instances_val2017.json，保存annotations文件夹下，train2017和val2107中存放训练集图片和验证集图片，文件夹结构如下：
在这里插入图片描述

第一步：创建people.names类别文件

在这里插入图片描述

第二步：格式转换，脚本如下：

import os
import json
import cv2
import random
import time
from PIL import Imagecoco_format_save_path='/root/detr/data/Crowdhuman/images/annotations'    #要生成的标准coco格式标签所在文件夹
yolo_format_classes_path='/root/detr/data/Crowdhuman/images/people.names'     #类别文件，一行一个类
yolo_format_annotation_path='/root/detr/data/Crowdhuman/labels/train'        #yolo格式标签所在文件夹
img_pathDir='/root/detr/data/Crowdhuman/images/train2017'                        #图片所在文件夹with open(yolo_format_classes_path,'r') as fr:                               #打开并读取类别文件lines1=fr.readlines()
# print(lines1)
categories=[]                                                                 #存储类别的列表
for j,label in enumerate(lines1):label=label.strip()categories.append({'id':j+1,'name':label,'supercategory':'None'})         #将类别信息添加到categories中
# print(categories)write_json_context=dict()                                                      #写入.json文件的大字典
write_json_context['info']= {'description': '', 'url': '', 'version': '', 'year': 2021, 'contributor': '', 'date_created': '2021-07-25'}
write_json_context['licenses']=[{'id':1,'name':None,'url':None}]
write_json_context['categories']=categories
write_json_context['images']=[]
write_json_context['annotations']=[]#接下来的代码主要添加'images'和'annotations'的key值
imageFileList=os.listdir(img_pathDir)                                           #遍历该文件夹下的所有文件，并将所有文件名添加到列表中
for i,imageFile in enumerate(imageFileList):imagePath = os.path.join(img_pathDir,imageFile)                             #获取图片的绝对路径image = Image.open(imagePath)                                               #读取图片，然后获取图片的宽和高W, H = image.sizeimg_context={}                                                              #使用一个字典存储该图片信息#img_name=os.path.basename(imagePath)                                       #返回path最后的文件名。如果path以/或\结尾，那么就会返回空值img_context['file_name']=imageFileimg_context['height']=Himg_context['width']=Wimg_context['date_captured']='2021-07-25'img_context['id']=i                                                         #该图片的idimg_context['license']=1img_context['color_url']=''img_context['flickr_url']=''write_json_context['images'].append(img_context)                            #将该图片信息添加到'image'列表中txtFile = imageFile.rsplit('.', 1)[0] + '.txt'                                              #获取该图片获取的txt文件，这个数字"6"要根据自己图片名修改with open(os.path.join(yolo_format_annotation_path,txtFile),'r') as fr:lines=fr.readlines()                                                   #读取txt文件的每一行数据，lines2是一个列表，包含了一个图片的所有标注信息for j,line in enumerate(lines):bbox_dict = {}                                                          #将每一个bounding box信息存储在该字典中# line = line.strip().split()# print(line.strip().split(' '))class_id,x,y,w,h=line.strip().split(' ')                                          #获取每一个标注框的详细信息class_id,x, y, w, h = int(class_id), float(x), float(y), float(w), float(h)       #将字符串类型转为可计算的int和float类型xmin=(x-w/2)*W                                                                    #坐标转换ymin=(y-h/2)*Hxmax=(x+w/2)*Wymax=(y+h/2)*Hw=w*Wh=h*Hbbox_dict['id']=i*10000+j                                                         #bounding box的坐标信息bbox_dict['image_id']=ibbox_dict['category_id']=class_id+1                                               #注意目标类别要加一bbox_dict['iscrowd']=0height,width=abs(ymax-ymin),abs(xmax-xmin)bbox_dict['area']=height*widthbbox_dict['bbox']=[xmin,ymin,w,h]bbox_dict['segmentation']=[[xmin,ymin,xmax,ymin,xmax,ymax,xmin,ymax]]write_json_context['annotations'].append(bbox_dict)                               #将每一个由字典存储的bounding box信息添加到'annotations'列表中name = os.path.join(coco_format_save_path,"instances_train2017"+ '.json')
with open(name,'w') as fw:                                                                #将字典信息写入.json文件中json.dump(write_json_context,fw,indent=2)

备注：必须严格按照笔者图中的文件命名方式进行命名，训练集清一色命名为instances_train2017.json,验证集清一色命名为instances_val2017.json，这是模型本身的命名要求，用户需要严格遵守。

二、修改训练模型参数

第一步：先在目录中新建python脚本文件detr_r50_tf.py，代码如下：

import torchpretrained_weights = torch.load('detr-r50-e632da11.pth')num_class = 2  # 类别数+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)

第二步：将其中类别数改成自己数据集的类别数即可，执行完成后会在目录下生成适合自己数据集类别的预训练模型：

在这里插入图片描述

第三步：然后在models文件夹下打开detr.py，修改其中的类别数：

在这里插入图片描述

第四步：打开main.py，修改其中的coco_path（数据存放路径）、output_dir（结果输出路径）、device（没有cuda就改为cpu）、resume（自己生成的预训练模型）

在这里插入图片描述

第五步：修改epochs数

在这里插入图片描述

三、开始训练

运行python main.py

跑起来的效果是这样的：

在这里插入图片描述

四、实现DETR的推理

将要预测的图片保存在一个文件夹下，预测时一次输出所有图片的预测结果，代码如下：

import argparse
import random
import time
from pathlib import Path
import numpy as np
import torch
from models import build_model
from PIL import Image
import os
import torchvision
from torchvision.ops.boxes import batched_nms
import cv2# 设置参数
def get_args_parser():parser = argparse.ArgumentParser('Set transformer detector', add_help=False)parser.add_argument('--lr', default=1e-4, type=float)parser.add_argument('--lr_backbone', default=1e-5, type=float)parser.add_argument('--batch_size', default=2, type=int)parser.add_argument('--weight_decay', default=1e-4, type=float)parser.add_argument('--epochs', default=300, type=int)parser.add_argument('--lr_drop', default=200, type=int)parser.add_argument('--clip_max_norm', default=0.1, type=float, help='gradient clipping max norm')# Model parametersparser.add_argument('--frozen_weights', type=str, default=None, help="Path to the pretrained model. If set, only the mask head will be trained")parser.add_argument('--backbone', default='resnet50', type=str, help="Name of the convolutional backbone to use")parser.add_argument('--dilation', action='store_true', help="If true, we replace stride with dilation in the last convolutional block (DC5)")parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), help="Type of positional embedding to use on top of the image features")# Transformerparser.add_argument('--enc_layers', default=6, type=int, help="Number of encoding layers in the transformer")parser.add_argument('--dec_layers', default=6, type=int, help="Number of decoding layers in the transformer")parser.add_argument('--dim_feedforward', default=2048, type=int, help="Intermediate size of the feedforward layers in the transformer blocks")parser.add_argument('--hidden_dim', default=256, type=int, help="Size of the embeddings (dimension of the transformer)")parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer")parser.add_argument('--nheads', default=8, type=int, help="Number of attention heads inside the transformer's attentions")parser.add_argument('--num_queries', default=100, type=int, help="Number of query slots")parser.add_argument('--pre_norm', action='store_true')# Segmentationparser.add_argument('--masks', action='store_true', help="Train segmentation head if the flag is provided")# Lossparser.add_argument('--no_aux_loss', dest='aux_loss', default='False', help="Disables auxiliary decoding losses (loss at each layer)")# Matcherparser.add_argument('--set_cost_class', default=1, type=float, help="Class coefficient in the matching cost")parser.add_argument('--set_cost_bbox', default=5, type=float, help="L1 box coefficient in the matching cost")parser.add_argument('--set_cost_giou', default=2, type=float, help="giou box coefficient in the matching cost")# Loss coefficientsparser.add_argument('--mask_loss_coef', default=1, type=float)parser.add_argument('--dice_loss_coef', default=1, type=float)parser.add_argument('--bbox_loss_coef', default=5, type=float)parser.add_argument('--giou_loss_coef', default=2, type=float)parser.add_argument('--eos_coef', default=0.1, type=float, help="Relative classification weight of the no-object class")# dataset parametersparser.add_argument('--dataset_file', default='coco')parser.add_argument('--coco_path', type=str, default="/root/detr/data/Crowdhuman/coco")parser.add_argument('--coco_panoptic_path', type=str)parser.add_argument('--remove_difficult', action='store_true')parser.add_argument('--output_dir', default='/root/detr/inference_demo/inference_output', help='path where to save, empty for no saving')parser.add_argument('--device', default='cuda', help='device to use for training / testing')parser.add_argument('--seed', default=42, type=int)parser.add_argument('--resume', default='/root/detr/data/output/checkpoint.pth', help='resume from checkpoint')parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch')parser.add_argument('--eval', default="True")parser.add_argument('--num_workers', default=2, type=int)# distributed training parametersparser.add_argument('--world_size', default=1, type=int, help='number of distributed processes')parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')return parserdef box_cxcywh_to_xyxy(x):# 将DETR的检测框坐标(x_center,y_center,w,h)转化成coco数据集的检测框坐标(x0,y0,x1,y1)x_c, y_c, w, h = x.unbind(1)b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]return torch.stack(b, dim=1)def rescale_bboxes(out_bbox, size):# 把比例坐标乘以图像的宽和高，变成真实坐标img_w, img_h = sizeb = box_cxcywh_to_xyxy(out_bbox)b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)return bdef filter_boxes(scores, boxes, confidence=0.7, apply_nms=True, iou=0.5):# 筛选出真正的置信度高的框keep = scores.max(-1).values > confidencescores, boxes = scores[keep], boxes[keep]if apply_nms:top_scores, labels = scores.max(-1)keep = batched_nms(boxes, top_scores, labels, iou)scores, boxes = scores[keep], boxes[keep]return scores, boxes# COCO classes
CLASSES = ['N/A', 'pedestrian']# 生成随机颜色的函数
def random_color():return [random.randint(0, 255) for _ in range(3)]# 创建类别颜色字典
COLORS = {cls: random_color() for cls in CLASSES}def plot_one_box(x, img, color=None, label=None, line_thickness=2):# 把检测框画到图片上tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thicknesscolor = [255, 0, 0]  # 固定为红色c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)if label:tf = max(tl - 1, 1)  # font thicknesst_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filledcv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)def main(args):print(args)device = torch.device(args.device)# 导入网络# 下面的criterion是算损失函数要用的，推理用不到,postprocessors是解码用的，这里也没有用，用的是自己的。model, criterion, postprocessors = build_model(args)# 加载权重checkpoint = torch.load(args.resume, map_location='cuda')model.load_state_dict(checkpoint['model'])# 把权重加载到gpu或cpu上model.to(device)# 打印出网络的参数大小n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)print("parameters:", n_parameters)# 设置好存储输出结果的文件夹output_dir = Path(args.output_dir)# 读取数据集,进行推理image_Totensor = torchvision.transforms.ToTensor()image_file_path = os.listdir("inference_demo/detect_demo")image_set = []for image_item in image_file_path:print("inference_image:", image_item)image_path = os.path.join("inference_demo/detect_demo", image_item)image = Image.open(image_path)image_tensor = image_Totensor(image)image_tensor = torch.reshape(image_tensor, [-1, image_tensor.shape[0], image_tensor.shape[1], image_tensor.shape[2]])image_tensor = image_tensor.to(device)time1 = time.time()inference_result = model(image_tensor)time2 = time.time()print("inference_time:", time2 - time1)probas = inference_result['pred_logits'].softmax(-1)[0, :, :-1].cpu()bboxes_scaled = rescale_bboxes(inference_result['pred_boxes'][0, ].cpu(), (image_tensor.shape[3], image_tensor.shape[2]))scores, boxes = filter_boxes(probas, bboxes_scaled)scores = scores.data.numpy()boxes = boxes.data.numpy()for i in range(boxes.shape[0]):class_id = scores[i].argmax()label = CLASSES[class_id]confidence = scores[i].max()text = f"{label} {confidence:.3f}"image = np.array(image)plot_one_box(boxes[i], image, color=COLORS[label], label=text)# cv2.imshow("images", image)# cv2.waitKey(1)image = Image.fromarray(image)image.save(os.path.join(args.output_dir, image_item))if __name__ == '__main__':parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()])args = parser.parse_args()if args.output_dir:Path(args.output_dir).mkdir(parents=True, exist_ok=True)main(args)