RT-DETR+Flask實現目標檢測推理案例

今天，帶大家利用RT-DETR（我們可以換成任意一個模型）+Flask來實現一個目標檢測平臺小案例，其實現效果如下：

目標檢測案例

這個案例很簡單，就是讓我們上傳一張圖像，隨后選擇一下置信度，即可檢測出圖像中的目標，那么具體該如何實現呢？

RT-DETR模型推理

在先前的學習過程中，博主對RT-DETR進行來了簡要的介紹，作為百度提出的實時性目標檢測模型，其無論是速度還是精度均取得了較為理想的效果，今天則主要介紹一下RT-DETR的推理過程，與先前使用DETR中使用pth權重與網絡結構相結合的推理方式不同，RT-DETR中使用的是onnx這種權重文件，因此，我們需要先對onnx文件進行一個簡單了解：
在這里插入圖片描述

ONNX模型文件

import onnx
# 加載模型
model = onnx.load('onnx_model.onnx')
# 檢查模型格式是否完整及正確
onnx.checker.check_model(model)
# 獲取輸出層，包含層名稱、維度信息
output = self.model.graph.output
print(output)

在原本的DETR類目標檢測算法中，推理是采用權重文件與模型結構代碼相結合的方式，而在RT-DETR中，則采用onnx模型文件來進行推理，即只需要該模型文件即可。

首先是將pth文件與模型結構進行匹配，從而導出onnx模型文件

"""by lyuwenyu
"""import os 
import sys
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))import argparse
import numpy as np from src.core import YAMLConfigimport torch
import torch.nn as nn def main(args, ):"""main"""cfg = YAMLConfig(args.config, resume=args.resume)if args.resume:checkpoint = torch.load(args.resume, map_location='cpu') if 'ema' in checkpoint:state = checkpoint['ema']['module']else:state = checkpoint['model']else:raise AttributeError('only support resume to load model.state_dict by now.')# NOTE load train mode state -> convert to deploy modecfg.model.load_state_dict(state)class Model(nn.Module):def __init__(self, ) -> None:super().__init__()self.model = cfg.model.deploy()self.postprocessor = cfg.postprocessor.deploy()print(self.postprocessor.deploy_mode)def forward(self, images, orig_target_sizes):outputs = self.model(images)return self.postprocessor(outputs, orig_target_sizes)model = Model()dynamic_axes = {'images': {0: 'N', },'orig_target_sizes': {0: 'N'}}data = torch.rand(1, 3, 640, 640)size = torch.tensor([[640, 640]])torch.onnx.export(model, (data, size), args.file_name,input_names=['images', 'orig_target_sizes'],output_names=['labels', 'boxes', 'scores'],dynamic_axes=dynamic_axes,opset_version=16, verbose=False)if args.check:import onnxonnx_model = onnx.load(args.file_name)onnx.checker.check_model(onnx_model)print('Check export onnx model done...')if args.simplify:import onnxsimdynamic = True input_shapes = {'images': data.shape, 'orig_target_sizes': size.shape} if dynamic else Noneonnx_model_simplify, check = onnxsim.simplify(args.file_name, input_shapes=input_shapes, dynamic_input_shape=dynamic)onnx.save(onnx_model_simplify, args.file_name)print(f'Simplify onnx model {check}...')
if __name__ == '__main__':parser = argparse.ArgumentParser()parser.add_argument('--config', '-c',  default="D:\graduate\programs\RT-DETR-main\RT-DETR-main//rtdetr_pytorch\configs/rtdetr/rtdetr_r18vd_6x_coco.yml",type=str, )parser.add_argument('--resume', '-r', default="D:\graduate\programs\RT-DETR-main\RT-DETR-main/rtdetr_pytorch/tools\output/rtdetr_r18vd_6x_coco\checkpoint0024.pth",type=str, )parser.add_argument('--file-name', '-f', type=str, default='model.onnx')parser.add_argument('--check',  action='store_true', default=False,)parser.add_argument('--simplify',  action='store_true', default=False,)args = parser.parse_args()main(args)

隨后，便是利用onnx模型文件進行目標檢測推理過程了
onnx也有自己的一套流程：

onnx前向InferenceSession的使用

關于onnx的前向推理，onnx使用了onnxruntime計算引擎。
onnx runtime是一個用于onnx模型的推理引擎。微軟聯合Facebook等在2017年搞了個深度學習以及機器學習模型的格式標準–ONNX，順路提供了一個專門用于ONNX模型推理的引擎（onnxruntime）。

import onnxruntime
# 創建一個InferenceSession的實例，并將模型的地址傳遞給該實例
sess = onnxruntime.InferenceSession('onnxmodel.onnx')
# 調用實例sess的潤方法進行推理
outputs = sess.run(output_layers_name, {input_layers_name: x})

推理詳細代碼

推理代碼如下：

import torch
import onnxruntime as ort
from PIL import Image, ImageDraw
from torchvision.transforms import ToTensorif __name__ == "__main__":##################classes = ['car','truck',"bus"]################### print(onnx.helper.printable_graph(mm.graph))#############img_path = "1.jpg"#############im = Image.open(img_path).convert('RGB')im = im.resize((640, 640))im_data = ToTensor()(im)[None]print(im_data.shape)size = torch.tensor([[640, 640]])sess = ort.InferenceSession("model.onnx")import timestart = time.time()output = sess.run(output_names=['labels', 'boxes', 'scores'],#output_names=None,input_feed={'images': im_data.data.numpy(), "orig_target_sizes": size.data.numpy()})end = time.time()fps = 1.0 / (end - start)print(fps)# print(type(output))# print([out.shape for out in output])labels, boxes, scores = outputdraw = ImageDraw.Draw(im)thrh = 0.6for i in range(im_data.shape[0]):scr = scores[i]lab = labels[i][scr > thrh]box = boxes[i][scr > thrh]print(i, sum(scr > thrh))#print(lab)print(f'box:{box}')for l, b in zip(lab, box):draw.rectangle(list(b), outline='red',)print(l.item())draw.text((b[0], b[1] - 10), text=str(classes[l.item()]), fill='blue', )#############im.save('2.jpg')#############

前端代碼

前端代碼包含兩部分，一個是上傳頁面，一個是顯示頁面

上傳頁面如下：

<!DOCTYPE html>
<html lang="en">
<head><meta charset="UTF-8"><meta name="viewport" content="initial-scale=1.0, maximum-scale=1.0, user-scalable=no" /><title></title><script src="http://www.jq22.com/jquery/jquery-1.10.2.js"></script><style>#addCommodityIndex {text-align: center;width: 300px;height: 340px;position: absolute;left: 50%;top: 50%;margin: -200px 0 0 -200px;border: solid #ccc 1px;padding: 35px;}#imghead {cursor: pointer;}.btn {width: 100%;height: 40px;text-align: center;}</style><link rel="stylesheet" href="../static/css/bootstrap.min.css"  crossorigin="anonymous">
</head><body><div id="addCommodityIndex"><h2>目標檢測</h2><div class="form-group row"><form id="upload"  action="/upload" enctype="multipart/form-data" method="POST"><img src=""><div class="form-group row"><label>上傳圖像</label><input type="file" class="form-control"  name='file'></div><div class="form-group row"><label>選擇置信度</label><select class="form-control" name="score" id="exampleFormControlSelect1"><option value="0.5">0.5</option><option value="0.6">0.6</option><option value="0.7">0.7</option><option value="0.8">0.8</option><option value="0.9">0.9</option></select></div><div class="form-group row"><div class="btn"><input type="submit" class="btn btn-success" value="提交圖像" /></div></div></form></div></div></body>
</html>

顯示頁面：

<!DOCTYPE html>
<html lang="en"><head><meta charset="UTF-8"><meta name="viewport" content="initial-scale=1.0, maximum-scale=1.0, user-scalable=no" /><title></title><script src="http://www.jq22.com/jquery/jquery-1.10.2.js"></script><style>#addCommodityIndex {text-align: center;position: absolute;left: 40%;top: 50%;margin: -200px 0 0 -200px;border: solid #ccc 1px;}#imghead {cursor: pointer;}.result {width: 100%;height: 100%;text-align: center;}</style><link rel="stylesheet" href="../static/css/bootstrap.min.css"  crossorigin="anonymous">
</head><body><div id="addCommodityIndex">
<div class="card mb-3" style="max-width: 680px;"><div class="row no-gutters"><div class="col-md-5"><img src="../static/img/result.jpg" class="result"></div><div class="col-md-5"><div class="card-body"><h5 class="card-title">檢測結果</h5><p class="card-text">目標數量：{{num}}</p><p class="card-text">檢測速度：{{fps}} 幀/秒</p><a  href="/home" class="btn btn-success">繼續提交</a></div></div></div>
</div>
</div>
</body>
</html>

Flask框架代碼：

# -*- coding: utf-8 -*-
from flask import Flask,request,render_template
import json
import os
import time
app = Flask(__name__)
import infer
@app.route('/home',methods=['GET'])
def home():return render_template('upload.html')@app.route('/upload',methods=['GET','POST'])
def upload():if request.method == 'POST':f = request.files['file'] #獲取數據流rootPath = os.path.dirname(os.path.abspath(__file__)) #根目錄路徑#創建存儲文件的文件夾，使用時間戳防止重名覆蓋file_path = 'static/upload/' + str(int(time.time()))absolute_path = os.path.join(rootPath,file_path).replace('\\','/') #存儲文件的絕對路徑，window路徑顯示\\要轉化/if not os.path.exists(absolute_path): #不存在改目錄則會自動創建os.makedirs(absolute_path)save_file_name = os.path.join(absolute_path,f.filename).replace('\\','/') #文件存儲路徑（包含文件名）f.save(save_file_name)score=request.values.to_dict().get("score")num,fps=infer.inference(save_file_name,score)#return json.dumps({'code':200,'url':url_path},ensure_ascii=False)return render_template("show.html",num=num,fps=fps)app.run(port='5000',debug=True)

上述項目博主已經上傳到github上

git init
git add README.md
git commit -m "first commit"
git branch -M main
git remote add origin https://github.com/pengxiang1998/rt-detr.git
git push -u origin main

項目地址

在使用onnx時，安裝了onnxruntime后，出現了下面的錯誤：

ImportError: cannot import name 'create_and_register_allocator_v2' from 'onnxruntime.capi._pybind_state'

這是由于onnxruntime-gpu版本與CUDA、CuDNN版本不匹配導致的，可以查看下面的網址來查看匹配版本

https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html

在這里插入圖片描述
隨后又出現錯誤：

> This ORT build has ['TensorrtExecutionProvider',
> 'CUDAExecutionProvider', 'CPUExecutionProvider'] enabled. Since ORT
> 1.9, you are required to explicitly set the providers parameter when instantiating InferenceSession. For example,
> onnxruntime.InferenceSession(...,
> providers=['TensorrtExecutionProvider',

這是由于InferenceSession中沒有提供對應的provider，修改代碼如下：

if torch.cuda.is_available():print("GPU")sess = ort.InferenceSession("model.onnx", None, providers=["CUDAExecutionProvider"])else:print("CPU")sess= ort.InferenceSession("model.onnx", None)

隨后運行，發現安裝了onnxruntime-gpu后的速度竟然滿了下來，fps僅為0.2，而原本使用onnxruntime的fps則為7左右，這到底是怎么回事呢？

在這里插入圖片描述

YOLO集成推理

而在YOLO集成的RT-DETR項目中，訓練得到的權重文件為.pt，在推理時需要與RT-DETR搭配使用，從而實現推理過程：
需要注意的是，由于YOLO里面集成了多種模型，因此為了具有適配性，其代碼都具有通用性

from ultralytics.models import RTDETR
if __name__ == '__main__':model=RTDETR("weights/best.pt")model.predict(source="images/1.mp4",save=True,conf=0.6)

隨后執行predict，代碼如下：

def predict(self,source: Union[str, Path, int, list, tuple, np.ndarray, torch.Tensor] = None,stream: bool = False,predictor=None,**kwargs,) -> list:if source is None:source = ASSETSLOGGER.warning(f"WARNING ?? 'source' is missing. Using 'source={source}'.")is_cli = (ARGV[0].endswith("yolo") or ARGV[0].endswith("ultralytics")) and any(x in ARGV for x in ("predict", "track", "mode=predict", "mode=track"))custom = {"conf": 0.25, "batch": 1, "save": is_cli, "mode": "predict"}  # method defaultsargs = {**self.overrides, **custom, **kwargs}  # highest priority args on the rightprompts = args.pop("prompts", None)  # for SAM-type modelsif not self.predictor:self.predictor = predictor or self._smart_load("predictor")(overrides=args, _callbacks=self.callbacks)self.predictor.setup_model(model=self.model, verbose=is_cli)else:  # only update args if predictor is already setupself.predictor.args = get_cfg(self.predictor.args, args)if "project" in args or "name" in args:self.predictor.save_dir = get_save_dir(self.predictor.args)if prompts and hasattr(self.predictor, "set_prompts"):  # for SAM-type modelsself.predictor.set_prompts(prompts)return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)