DETR测试自己数据的方法（facebook）

原创

已于 2022-07-16 20:15:18 修改 · 1.2k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#计算机视觉

于 2022-07-16 08:50:33 首次发布

该博客介绍了一个使用 DETR 模型进行目标检测的 Python 脚本。脚本涉及训练参数设置、数据集加载、模型构建、损失函数、匹配策略等关键步骤。还展示了如何从预训练模型加载权重、对单张图片进行预测并显示结果的过程。

废话不多说，直接上代码

import argparse
import datetime
import json
import random
import time
from pathlib import Path
from PIL import Image
from datasets.KINS import make_coco_transforms
import numpy as np
import torch
from torch.utils.data import DataLoader, DistributedSampler

import datasets
import util.misc as utils
from datasets import build_dataset, get_coco_api_from_dataset
from engine import evaluate, train_one_epoch
from models import build_model

# TODO：修改自己数据集的类别
CLASSES = [
    'pedestrain', 'cyclist', 'person-sitting', 'van', 'car','tram','truck','misc','bicyccle','N/A'
]
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.647, 0.188], [0.301, 0.785, 0.933]
]

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=1, 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 convolutio