基于transformer的车辆识别

需求：

使用transformer算法相关框架实现能够自己训练模型，完成相关的目标检测(人和车)。

测试：

训练好模型后，需要用已有的模型对新的图片测试，测试代码改写如下：

# -*- coding='utf-8' -*-
import os, sys
import torch
import numpy as np
from models import build_DABDETR, build_dab_deformable_detr
from util.slconfig import SLConfig
from datasets import build_dataset
from util.visualizer import COCOVisualizer
from util import box_ops

# 加载模型与配置文件
model_config_path = "logs/DABDETR/R50/config.json" 
model_checkpoint_path = "logs/DABDETR/R50/checkpoint.pth" 

args = SLConfig.fromfile(model_config_path) 
model, criterion, postprocessors = build_DABDETR(args)
checkpoint = torch.load(model_checkpoint_path, map_location='cpu')
model.load_state_dict(checkpoint['model'])
# 对验证集合中的一张图片进行测试
dataset_val = build_dataset(image_set='val', args=args)
#dataset_val = build_dataset(image_set='test', args=args)
cocojs = dataset_val.coco.dataset
id2name = {item['id']: item['name'] for item in cocojs['categories']}

image, targets = dataset_val[0]
# build gt_dict for vis
box_label = [id2name[int(item)] for item in targets['labels']]
gt_dict = {
    'boxes': targets['boxes'],
    'image_id': targets['image_id'],
    'size': targets['size'],
    'box_label': box_label,
}

vslzr = COCOVisualizer()
#vslzr.visualize(image, gt_dict, savedir=None)
vslzr.visualize(image, gt_dict, savedir='./result')
output = model(image[None])
output = postprocessors['bbox'](output, torch.Tensor([[1.0, 1.0]]))[0]

thershold = 0.3 # set a thershold

scores = output['scores']
labels = output['labels']
boxes = box_ops.box_xyxy_to_cxcywh(output['boxes'])
select_mask = scores > thershold

box_label = [id2name[int(item)] for item in labels[select_mask]]
pred_dict = {
    'boxes': boxes[select_mask],
    'size': targets['size'],
    'box_label': box_label
}
vslzr.visualize(image, pred_dict, savedir=None)
print('运行结束！')

说明：

推论预测的图片会保存在result文件夹下，具体路径位置可以根据自己的需求进行改写。

效果展示：

最后

以上就是曾经西牛为你收集整理的基于transformer的车辆识别的全部内容，希望文章能够帮你解决基于transformer的车辆识别所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。