DETR训练自己的数据集

DETR训练自己的数据集实验笔记

DETR是一个利用transfORMer实现端到端目标检测的模型。本文记录利用官方提供的代码来训练验证自己的数据集的过程以及一些注意事项。

一.数据集的准备与处理

此次项目用到的数据集为自己制造的自动驾驶领域的路况数据集，该数据集一共包含57个类别：names = [ "i2", "i4", "i5", "il100", "il60", "il80", "io", "ip", "p10", "p11", "p12", "p19", "p23", "p26", "p27", "p3", "p5", "p6", "pg", "ph4", "ph4d5", "ph5", "pl100", "pl120", "pl20", "pl30", "pl40", "pl5", "pl50", "pl60", "pl70", "pl80", "pm20", "pm30", "pm55", "pn", "pne", "po", "pr40", "w13", "w32", "w55", "w57", "w59", "wo", 'crosswalk', 'left', 'right', 'straight', 'straightl', 'straightr', 'diamond', 'noparking', 'person', 'car', 'truck', 'bus' ]
DETR需要的数据集格式为coco格式，其原始图片与标签分为训练集和验证集，其保存目录为
其中annotations文件夹下保存训练集和验证集的标注信息，注：训练集标签json文件命名为instances_train2017.json .验证集标签json文件命名为instance_val2017.json

该数据集原始的标签保存格式为yolo类型的txt文件，下面提供的代码能够实现数据集标签从txt文件到coco JSON格式的转化.

生成instances_xxx.json文件

#将yolo格式标注的txt文件转化为coco数据集标注格式的json文件类型#yolo格式为(xc,yc,w,h)相对坐标  coco标注格式为(xmin,ymin,w,h),绝对坐标 voc标注xml格式为(xmin,ymin,xmax,ymax)import argparseimport osimport sysimport jsonimport shutilimport cv2from datetime import datetimecoco=dict()coco['images']=[]coco['type']='instances'coco['annotations']=[]coco['cateGories']=[]category_set=dict()images_set=set()image_id=000000annotation_id=0def addCatItem(categroy_dict):    #保存所有的类别信息         for k,v in categroy_dict.items():             category_item=dict()             category_item['supercategory']='none'             category_item['id']=int(k)             category_item['name']=v             coco['categories'].append(category_item)     def addImgitem(file_name,size):    global image_id        image_id +=1        image_item=dict()        image_item['id']=image_id        image_item['file_name']=file_name        image_item['width']=size[1]        image_item['height']=size[0]        image_item['license']=None        image_item['flickr_url']=None        image_item['coco_url']=None        image_item['data_captured']=str(datetime.today())    coco['images'].append(image_item)        images_set.add(file_name)        return image_id    def addAnnoItem(object_name,image_id,category_id,bbox):    global annotation_id    annotation_item=dict()    annotation_item['segmentation']=[]    seg=[]    #bbox  is  x,y,w,h    seg.append(bbox[0])    seg.append(bbox[1])    seg.append(bbox[0])        seg.append(bbox[1]+bbox[3])    seg.append(bbox[0]+bbox[2])        seg.append(bbox[1]+bbox[3])    seg.append(bbox[0]+bbox[2])        seg.append(bbox[1])    annotation_item['segmentation'].append(seg)    annotation_item['area']=bbox[2]*bbox[3]   #w*h        annotation_item['iscrowd']=0        annotation_item['ignore']=0        annotation_item['image_id']=image_id        annotation_item['bbox']=bbox        annotation_item['category_id']=category_id        annotation_id +=1        annotation_item['id']=annotation_id         coco['annotations'].append(annotation_item)    def xywhn2xywh(bbox,size):  #从yolo标注到coco标注    bbox=list(map(float,bbox))    size=list(map(float,size))  #h,w    xmin=(bbox[0]-bbox[2]/2)*size[1]    ymin=(bbox[1]-bbox[3]/2)*size[0]    w=bbox[2]*size[1]    h=bbox[3]*size[0]   bbox=(xmin,ymin,w,h)       return list(map(int,bbox))def parseXmlFilse(image_path, anno_path, save_path, json_name):    assert os.path.exists(image_path), "ERROR {} dose not exists".format(image_path)    assert os.path.exists(anno_path), "ERROR {} dose not exists".format(anno_path)    if os.path.exists(save_path):                     shutil.rmtree(save_path)                 os.makedirs(save_path)            json_path = os.path.join(save_path, json_name)        category_set = []            with open(anno_path + '/classes.txt', 'r') as f:                for i in f.readlines():                          category_set.append(i.strip())            category_id = dict((k, v) for k, v in enumerate(category_set))            addCatItem(category_id)        images = [os.path.join(image_path, i) for i in os.listdir(image_path)]            files = [os.path.join(anno_path, i) for i in os.listdir(anno_path)]                      images_index = dict((v.split(os.sep)[-1][:-4], k) for k, v in enumerate(images))            for file in files:                    if os.path.splitext(file)[-1] != '.txt' or 'classes' in file.split(os.sep)[-1]: continue                    if file.split(os.sep)[-1][:-4] in images_index: index = images_index[file.split(os.sep)[-1][:-4]] img = cv2.imread(images[index]) shape = img.shape filename = images[index].split(os.sep)[-1]     current_image_id = addImgItem(filename, shape)                   else: continue                   with open(file, 'r') as fid: for i in fid.readlines():         i = i.strip().split()         category = int(i[0])         category_name = category_id[category]         bbox = xywhn2xywh((i[1], i[2], i[3], i[4]), shape)         addAnnoItem(category_name, current_image_id, category, bbox)           json.dump(coco,open(json_path,'w'))        print("class nums:{}".format(len(coco['categories'])))            print("image nums:{}".format(len(coco['images'])))           print("bbox nums:{}".format(len(coco['annotations'])))if __name__=='__main__':     '''参数说明：             anno_path:标注txt文件存储地址             save_path:json文件输出文件夹             image_path:图片路径             json_name:保存json文件名称'''        parser = argparse.ArgumentParser()            parser.add_argument('-ap', '--anno-path', type=str, default='/home/nianliu/wangxx/train_data_v5_format/labels', help='yolo txt path')            parser.add_argument('-s', '--save-path', type=str, default='/home/nianliu/wangxx/train_data_v5_format/anno_json', help='json save path')        parser.add_argument('--image-path', default='/home/nianliu/wangxx/train_data_v5_format/images/train')        parser.add_argument('--json-name', default='train.json')        opt = parser.parse_args()           if len(sys.argv) > 1:                    print(opt)                    parseXmlFilse(**vars(opt))            else:                anno_path = '/home/nianliu/wangxx/train_data_v5_format/labels/train'               save_path = '/home/nianliu/wangxx/train_data_v5_format/annotations'               image_path = '/home/nianliu/wangxx/train_data_v5``_format/images/train'                json_name = 'instance_train2017.json'                parseXmlFilse(image_path, anno_path, save_path, json_name) '''   

其中修改anno_path为原始txt类型标签文件地址，save_path为保存json文件的地址，image_path为数据集图像地址，json_name按照detr要求的修改
注：在train，val保存txt文件的地址中添加一个classes.txt文件，其内容按行写入数据集的类别名称

二.训练

首先通过DETR项目下载项目文件

1.修改预训练模型权重
DETR模型训练十分缓慢，因此需要下载其提供好的预训练模型
根基自己的数据集修改预训练文件，新建changepre.py,修改num_classes为自己数据集的类别数加一，运行该文件得到detr_r50_58.pth文件

import torchpretrained_weights=torch.load('/home/nianliu/wangxx/detr/detr-r50-e632da11.pth')num_classes=57+1pretrained_weights["model"]["class_embed.weight"].resize_(num_classes+1,256)pretrained_weights["model"]["class_embed.bias"].resize_(num_classes+1)torch.save(pretrained_weights,"detr_r50_%d.path"%num_classes)

2.训练模型
修改models/detr.py中313行num_classes为自己数据集的类别数。

利用main.py进行模型训练，注意以下几个参数
coco-path为数据集保存地址。

output_dir为保存训练得到的模型权重的地址。

resume为预训练模型权重文件。

终端执行即可进行模型的训练。

三.验证模型

训练过程中会每隔一段时间打印出验证结果。该部分为利用训练得到的模型权重进行检测验证。新建一个infer_demo.py进行模型验证

import globimport mathimport argparseimport numpy as npfrom models.detr import DETRfrom models.backbone import Backbone,build_backbonefrom models.transformer import build_transformerfrom PIL import Imageimport cv2import requestsimport matplotlib.pyplot as pltimport torchfrom torch import nnfrom torchvision.models import resnet50import torchvision.transforms as Timport torchvision.models as modelstorch.set_grad_enabled(False)import osdef 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 parameters        parser.add_argument('--frozen_weights', type=str, default=None,                        help="Path to the pretrained model. If set, only the mask head will be trained")    # * Backbone        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")    # * Transformer        parser.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')    # * Segmentation        parser.add_argument('--masks', action='store_true',                        help="Train segmentation head if the flag is provided")    # Loss        parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false',                        help="Disables auxiliary decoding losses (loss at each layer)")    # * Matcher        parser.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 coefficients       parser.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 parameters        parser.add_argument('--dataset_file', default='coco')        parser.add_argument('--coco_path', type=str)        parser.add_argument('--coco_panoptic_path', type=str)        parser.add_argument('--remove_difficult', action='store_true')        parser.add_argument('--output_dir', default='',                        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='', help='resume from checkpoint')    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',                        help='start epoch')        parser.add_argument('--eval', action='store_true')        parser.add_argument('--num_workers', default=2, type=int)    # distributed training parameters        parser.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 parserCLASSES = [   "i2", "i4", "i5", "il100", "il60",  "il80", "io", "ip", "p10", "p11",        "p12", "p19", "p23", "p26", "p27",         "p3", "p5", "p6", "pg", "ph4",         "ph4d5", "ph5", "pl100", "pl120", "pl20",         "pl30", "pl40", "pl5", "pl50", "pl60",         "pl70", "pl80", "pm20", "pm30", "pm55",          "pn", "pne", "po", "pr40", "w13",           "w32", "w55", "w57", "w59", "wo",            'crosswalk', 'left', 'right', 'straight', 'straightl', 'straightr', 'diamond', 'noparking',            'person', 'car', 'truck', 'bus' ] COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]] transform_input = T.Compose([    T.Resize(800),                T.ToTensor(),                T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])def box_cxcywh_to_xyxy(x):        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 = size        b = box_cxcywh_to_xyxy(out_bbox)        b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)        return bdef plot_results(pil_img, prob, boxes, save_path):            lw= max(round(sum(pil_img.shape) / 2 * 0.003), 2)        tf = max(lw - 1, 1)            colors = COLORS * 100        for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):           c1=p.argmax()                       text=f'{CLASSES[c1]}:{p[c1]:0.2f}'                       cv2.rectangle(pil_img, (int(xmin),int(ymin)), (int(xmax),int(ymax)), colors(c1,True), thickness=lw,lineType=cv2.LINE_AA)                       if text:                           tf=max(lw-1,1)                           w,h=cv2.getTextSize(text,0,fontScale=lw/3,thickness=tf)[0]    cv2.rectangle(pil_img,(int(xmin),int(ymin)), (int(xmin)+w,int(ymin)-h-3),colors(c1,True),-1,cv2.LINE_AA)    cv2.putText(pil_img, text, (int(xmin), int(ymin) - 2), 0, lw / 3, (255,255,255), thickness=tf,                        lineType=cv2.LINE_AA)        Image.fromarray(ori_img).save(save_path)parser = argparse.ArgumentParser('DETR training and evaluation script', parents=[get_args_parser()])args = parser.parse_args()backbone=build_backbone(args)transform=build_transformer(args)model=DETR(backbone=backbone,transformer=transform,num_classes=58,num_queries=100)model_path='/home/nianliu/wangxx/detr/cdnet_weights/checkpoint0179.pth'   #保存的预训练好的模型pth文件，用于验证model_data=torch.load(model_path)['model']model=torch.load(model_path)model.load_state_dict(model_data)model.eval();paths = os.listdir('/home/nianliu/wangxx/detr/images')  #待验证的图片路径for path in paths:    # 问题1：无法读取png图像        if os.path.splitext(path)[1] == ".png":    # 问题1解1：用imread读取png               im = cv2.imread(path)               im = Image.fromarray(cv2.cvtColor(im,cv2.COLOR_BGR2RGB))        else:                im = Image.open('/home/nianliu/wangxx/detr/images'+'/'+path)    # mean-std normalize the input image (batch-size: 1)            img = transform_input(im).unsqueeze(0)    # propagate through the model        outputs = model(img)    # keep only predictions with 0.9+ confidence        probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]        keep = probas.max(-1).values > 0.9    # convert boxes from [0; 1] to image scales        bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)    #保存验证结果地址    img_save_path = '/home/nianliu/wangxx/detr/infer_results/' + os.path.splitext(os.path.split(path)[1])[0] + '.jpg'            ori_img=np.array(im)        plot_results(ori_img, probas[keep], bboxes_scaled, img_save_path)

在detr工程文件下的images文件夹存放所有待验证的图片，infer_results存放所有的验证结果。
得到的预测结果如下所示：

来源地址：https://blog.csdn.net/qq_45836365/article/details/128252220