前言
需要准备windows10操作系统，python3.11.9，cuDNN8.9.2.26，CUDA11.8，paddleDetection2.7
流程：
- 准备数据集-澳洲鳌虾VOC数据集
- 基于RT-DETR目标检测模型训练
- 导出onnx模型进行python部署
- 平滑滤波处理视频帧保留的物体质心坐标
- 基于pywebview为软件前端，falsk为软件后端制作UI
- 使用pyinstaller打包成exe
- 使用into setup生成安装包
本人代码禁止任何商业化用途，个人开发者随意。所有代码均开源
项目目录
XXX 项目总目录static 存放js静态文件plotly.jstemplates 存放html文件index.htmltemp 用户上传文件保存路径venv 虚拟环境main.py 主程序model.onnx 模型文件1.ico 打包的程序图标准备数据集
点击下载澳洲鳌虾VOC数据集
下载后解压，文件目录为
dataAnnotations0.xml1.xml...imgs0.jpg1.jpg...lables.txt然后使用如下的脚本把数据集划分为训练集和测试集
import os
import random
import shutildef splitDatasets(images_dir,xmls_dir,train_dir,test_dir):if os.path.exists(train_dir):shutil.rmtree(train_dir)os.makedirs(train_dir)os.makedirs(train_dir+'/imgs')os.makedirs(train_dir+'/annotations')if os.path.exists(test_dir):shutil.rmtree(test_dir)os.makedirs(test_dir)os.makedirs(test_dir+'/imgs')os.makedirs(test_dir+'/annotations')images=os.listdir(images_dir)random.shuffle(images)split_index=int(0.9*len(images))train_images=images[:split_index]test_images=images[split_index:]with open(train_dir+'/train.txt','w') as file:for img in train_images:shutil.copy(os.path.join(images_dir,img),os.path.join(train_dir,'imgs',img))ann=img.replace('jpg','xml')shutil.copy(os.path.join(xmls_dir,ann),os.path.join(train_dir,'annotations',ann))line=os.path.join(train_dir,'imgs',img)+' '+os.path.join(train_dir,'annotations',ann)+'\n'file.write(line)with open(test_dir+'/test.txt','w') as file:for img in test_images:shutil.copy(os.path.join(images_dir,img),os.path.join(test_dir,'imgs',img))ann=img.replace('jpg','xml')shutil.copy(os.path.join(xmls_dir,ann),os.path.join(test_dir,'annotations',ann))line=os.path.join(test_dir,'imgs',img)+' '+os.path.join(test_dir,'annotations',ann)+'\n'file.write(line)shutil.rmtree(images_dir)shutil.rmtree(xmls_dir)if __name__=='__main__':# 填写img文件夹所在绝对路径images_dir='/home/aistudio/work/voc/imgs'# 填写Annotations文件夹所在绝对路径xmls_dir='/home/aistudio/work/voc/Annotations'# 填写 训练集 的存放的绝对路径train_dir='/home/aistudio/work/voc/trains'# 填写 测试集 的存放的绝对路径test_dir='/home/aistudio/work/voc/tests'splitDatasets(images_dir,xmls_dir,train_dir,test_dir)训练模型
可在aistudio云平台训练，我放好了所有的相关文件，点击进入，里面的说明很详细
也可在本地进行训练，下面来配置本地的训练环境
配置相关文件
下载paddleDetection2.7
原始目录如下
paddleDetection2.7.github.travisactivitybenchmarkconfigs 模型配置文件dataset 里面有数据集下载的脚本文件demodeploy 推理的相关文件docs 说明文档industrial_tutorialppdet 模型运行的核心文件scriptstest_pictools 模型训练入口，测试，验证，导出等脚本文件.gitignore.pre-commit-config.yaml.style.yapf.travis.ymlLICENSEREADME_cn.md 说明文档中文版README_en.md 说明文档英文版requirements.txt 相关依赖库setup.py 模型编译的相关脚本需要删除一些目录，把README_en.md改名为README.md，处理过的目录如下
paddleDetection2.7configsdatasetdeployppdettoolsREADME.mdrequirements.txtsetup.py把dataset里所有东西都删除，再将划分好的数据集放到该文件下，处理好的目录如下
datasetvoctrainsannotationsimgstrain.txttestsannotationsimgstest.txtlabels.txt进入tools目录，只保留如下文件，其余全删除，处理后的文件目录如下
toolstrain.pyinfer.pyeval.pyexport_model.py进入configs目录，只保留下面三个文件和目录，处理后的目录如下
configsdatasetsrtdetrruntime.yml进入datasets目录，只保留voc.yml，其余文件全删除，处理后的目录如下
datasetsvoc.yml并用如下内容覆盖voc.yml
metric: VOC
map_type: 11point
num_classes: 1TrainDataset:name: VOCDataSetdataset_dir: dataset/vocanno_path: trains/train.txtlabel_list: labels.txtdata_fields: ['image', 'gt_bbox', 'gt_class', 'difficult']EvalDataset:name: VOCDataSetdataset_dir: dataset/vocanno_path: tests/test.txtlabel_list: labels.txtdata_fields: ['image', 'gt_bbox', 'gt_class', 'difficult']TestDataset:name: ImageFolderanno_path: dataset/labels.txt进入rtdetr目录，只保留如下2个文件和目录，处理后的目录如下：
rtdetr_base_rtdetr_hgnetv2_x_6x_coco.yml进入_base_目录，找到optimizer_6x.yml，修改第一行为epoch: 200，意思是训练200轮
找到rtdetr_reader.yml，根据自己的CPU和GPU调整相关参数，如果是4核CPU，worker_num可为8，batch_size根据显存调整，占用80%到90%的显存即可
安装依赖库
建议在虚拟环境中操作
!pip install -r requirements.txt
!pip install pycocotools
!pip install filterpy
!pip install flask
!pip install pyinstaller
!pip install pywebview
!pip install onnxruntime-gpu
!pip install onnxruntime
!pip install onnx
!pip install paddle2onnx
!python setup.py install开始训练
建议命令行输入，先进入paddleDetection所在位置，再执行以下命令
python tools/train.py -c configs/rtdetr/rtdetr_hgnetv2_x_6x_coco.yml --eval --use_vdl True --vdl_log_dir vdl_log_dir/scalar然后就是漫长的等待
导出模型
生成的模型在paddleDetection/output/best_model/model.pdparams
先进入paddleDetection所在位置，再执行以下命令
python tools/export_model.py -c configs/rtdetr/rtdetr_hgnetv2_x_6x_coco.yml -o weights=output/best_model/model.pdparams转onnx
先进入paddleDetection所在位置，再执行以下命令，可以根据需要选择保存路径
paddle2onnx --model_dir=output_inference/rtdetr_hgnetv2_x_6x_coco/ \--model_filename model.pdmodel \--params_filename model.pdiparams \--opset_version 16 \--save_file /home/work/infer/model.onnx模型部署
导包
import webview
from flask import Flask, request, jsonify,render_template,stream_with_context,Response
import os
import time
import cv2
from onnxruntime import InferenceSession
import numpy as np
from werkzeug.utils import secure_filename总览代码
class TrackShrimp():def __init__(self,video_path,model_path,onnx_threshold=0.7):# 获取帧数据self.cap=self.init_video(video_path)frame_width=int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))frame_height=int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 图形尺寸im_shape = np.array([[frame_height, frame_width]], dtype='float32')# y轴缩放量self.im_scale_y=640.0/frame_height# x轴缩放量self.im_scale_x=640.0/frame_widthscale_factor = np.array([[self.im_scale_y,self.im_scale_x]]).astype('float32')# 定义模型输入self.inputs_dict = {'im_shape': im_shape,'image': None,'scale_factor': scale_factor}# 初始化模型self.sess=self.init_session(model_path)# 模型输出阈值self.onnx_threshold=onnx_thresholddef init_video(self,input_path):cap=cv2.VideoCapture(input_path)if not cap.isOpened():raise ValueError(f'无法打开视频{input_path}')return capdef init_session(self,model_path):try:return InferenceSession(model_path, providers=['CUDAExecutionProvider']) except:return InferenceSession(model_path, providers=['CPUExecutionProvider'])def precess_img(self,frame):img = cv2.resize(frame, None,None,fx=self.im_scale_x,fy=self.im_scale_y,interpolation=2)img = img.astype(np.float32) / 255.0img = np.transpose(img, [2, 0, 1])img = img[np.newaxis, :, :, :]return imgdef postcess(self,results:np.ndarray,all_centers:list[np.ndarray]):results=results[(results[:, 0] == 0) & (results[:, 1] > self.onnx_threshold)]x_centers = (results[:, 2] + results[:, 4]) / 2y_centers = (results[:, 3] + results[:, 5]) / 2centers = np.column_stack((x_centers, y_centers))all_centers.extend(centers)def by_smoothfilter(self,centers:list[np.ndarray],window_size=24):""":param centers: list[np.ndarray,np.ndarray,...]:param window_size: 平滑窗口大小:return: 平滑后的质心坐标NumPy数组"""centers=np.stack(centers)# 计算滑动窗口的平均值，pad函数在序列前后补零以处理边界情况padded_centers = np.pad(centers, ((window_size//2, window_size//2), (0, 0)), mode='edge')window_sum = np.cumsum(padded_centers, axis=0)smoothed_centers = (window_sum[window_size:] - window_sum[:-window_size]) / window_sizereturn smoothed_centersdef calculate_distance(self,centers:np.ndarray):'''centers:np.ndarray n*2'''# 计算相邻点之间的差diffs = centers[1:] - centers[:-1]# 计算每个差值的欧几里得距离distances = np.linalg.norm(diffs, axis=1)# 计算总路程return int(np.sum(distances))def gain_position(self,centers:np.ndarray):position_list=centers.tolist()return position_listdef run(self):global scheduleglobal run_task# 帧数frame_count=int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))frame_number=0center_list=[]for frame_number in range(frame_count):if not run_task:returnsuccess, frame = self.cap.read()if not success:breakschedule=int(frame_number/frame_count*100)# 打印进度if frame_number%10==0:print('Process: ',schedule)# 图片预处理img=self.precess_img(frame)self.inputs_dict['image']=imgresults=self.sess.run(None,self.inputs_dict)[0]if results is not None:self.postcess(results,center_list)# 使用平滑滤波filtered_centers = self.by_smoothfilter(center_list)self.cap.release()# 返回路程，轨迹坐标return self.calculate_distance(filtered_centers),self.gain_position(filtered_centers)由于是对视频进行推理，所以首先得初始化视频打开的方法
def init_video(self,input_path):cap=cv2.VideoCapture(input_path)if not cap.isOpened():raise ValueError(f'无法打开视频{input_path}')return cap初始化onnx运行引擎，优先使用显卡，如果CUDA环境有问题，就使用CPU运行
def init_session(self,model_path):try:return InferenceSession(model_path, providers=['CUDAExecutionProvider']) except:return InferenceSession(model_path, providers=['CPUExecutionProvider'])onnx引擎需要一定的输入格式，放到类的init里
def __init__(self,video_path,model_path,onnx_threshold=0.7):# 获取帧数据self.cap=self.init_video(video_path)frame_width=int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))frame_height=int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 图形尺寸im_shape = np.array([[frame_height, frame_width]], dtype='float32')# y轴缩放量self.im_scale_y=640.0/frame_height# x轴缩放量self.im_scale_x=640.0/frame_widthscale_factor = np.array([[self.im_scale_y,self.im_scale_x]]).astype('float32')# 定义模型输入self.inputs_dict = {'im_shape': im_shape,'image': None,'scale_factor': scale_factor}# 初始化模型self.sess=self.init_session(model_path)# 模型输出阈值self.onnx_threshold=onnx_threshold在提取每一帧后需要进行图像处理，resize图片为模型输入的要求，归一化
def precess_img(self,frame):img = cv2.resize(frame, None,None,fx=self.im_scale_x,fy=self.im_scale_y,interpolation=2)img = img.astype(np.float32) / 255.0img = np.transpose(img, [2, 0, 1])img = img[np.newaxis, :, :, :]return img在提取到视频的每一帧中的鳌虾的质心坐标后，由于每一帧的图像都不一样，输入模型后再输出的结果就不一样，会抖动，也就是噪声，我们需要滤波去噪，这里使用平滑滤波，相比卡尔曼滤波简单使用快速出结果。
def by_smoothfilter(self,centers:list[np.ndarray],window_size=24):""":param centers: list[np.ndarray,np.ndarray,...]:param window_size: 平滑窗口大小:return: 平滑后的质心坐标NumPy数组"""centers=np.stack(centers)# 计算滑动窗口的平均值，pad函数在序列前后补零以处理边界情况padded_centers = np.pad(centers, ((window_size//2, window_size//2), (0, 0)), mode='edge')window_sum = np.cumsum(padded_centers, axis=0)smoothed_centers = (window_sum[window_size:] - window_sum[:-window_size]) / window_sizereturn smoothed_centers我们需要计算鳌虾的运动总路程，用滤波后的质心坐标计算
def calculate_distance(self,centers:np.ndarray):'''centers:np.ndarray n*2'''# 计算相邻点之间的差diffs = centers[1:] - centers[:-1]# 计算每个差值的欧几里得距离distances = np.linalg.norm(diffs, axis=1)# 计算总路程return int(np.sum(distances))滤波后的质心坐标是numpy数组，需要一定的转换再发送到前端进行渲染（matplotlib画的图太丑了，不如plotly.js）
def gain_position(self,centers:np.ndarray):position_list=centers.tolist()return position_list在获取每一帧图像后，送入模型。模型会输出一对numpy数组，需要进行一对的后处理，低于阈值的就抛弃，然后取阈值最高的，计算质心坐标并保存
def postcess(self,results:np.ndarray,all_centers:list[np.ndarray]):results=results[(results[:, 0] == 0) & (results[:, 1] > self.onnx_threshold)]x_centers = (results[:, 2] + results[:, 4]) / 2y_centers = (results[:, 3] + results[:, 5]) / 2centers = np.column_stack((x_centers, y_centers))all_centers.extend(centers)需要在一个主函数里将上述打开视频，图像预处理，送入模型，后处理连起来
def run(self):global scheduleglobal run_task# 帧数frame_count=int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))frame_number=0center_list=[]for frame_number in range(frame_count):if not run_task:returnsuccess, frame = self.cap.read()if not success:breakschedule=int(frame_number/frame_count*100)# 打印进度if frame_number%10==0:print('Process: ',schedule)# 图片预处理img=self.precess_img(frame)self.inputs_dict['image']=imgresults=self.sess.run(None,self.inputs_dict)[0]if results is not None:self.postcess(results,center_list)# 使用平滑滤波filtered_centers = self.by_smoothfilter(center_list)self.cap.release()# 返回路程，轨迹坐标return self.calculate_distance(filtered_centers),self.gain_position(filtered_centers)前端的设计
以pywebview为平台，html和css设计前端
代码总览
index.html
<!DOCTYPE html>
<html>
<head><title></title><link rel="shortcut icon" href="#" /><script src="https://code.jquery.com/jquery-3.6.0.min.js"></script><script src="../static/plotly.js"></script><style>html,body{width: 100%;height: 100%;margin: 0 auto;}body{display: flex;align-items: center;justify-content: center;height: 100vh;background-color: rgb(6, 32, 80);}main{display: grid;grid-template-columns: 1fr 3fr;column-gap: 2%;width: 98%;height: 98%;}fieldset{border: 2px solid rgb(32, 139, 139);color: rgb(32, 139, 139);margin: 8% 0 8% 0;}#s2{text-align: center;display: flex;justify-content: center;align-items: center;background-color: rgba(32, 139, 139, 0.301);border: 2px solid rgb(32, 139, 139);}#progress-circle{border: 1em solid rgb(32, 139, 139);width: 40vh;height: 40vh;border-radius: 20vh;display: flex; justify-content: center; align-items: center;}#progress-num{font-size: 18vh;color: rgb(32, 139, 139);}</style>
</head>
<body><main><section id="s1"><form id="form" enctype="multipart/form-data"><fieldset><legend>选择你要检测的视频</legend><input type="file" accept=".mp4" id="video" name="vedio"></fieldset><fieldset><legend>功能按键</legend><button onclick="submit_to()" id="submit">开始上传</button><button onclick="stopRun()">终止运行</button></fieldset></form><script>async function stopRun(){try{const response=await fetch('/stopRun',{method:'POST'})if (!response.ok) { throw new Error('Network response was not ok.'); }data=await response.json()alert(data.data)}catch(error){console.log(error)}}async function submit_to(){// 防重复激发const button = document.getElementById('submit'); button.disabled = true;try{// 获取文件const input=document.getElementById('video')const file=input.files[0]if (!file){throw new Error('未选择文件')}if(file.type!=='video/mp4'){throw new Error('请选择MP4文件')}// 刷新界面 const s2=document.getElementById('s2')Plotly.purge(s2)// 初始化进度显示const progressCircle=document.getElementById('progress-circle')const progressNum=document.getElementById('progress-num')progressCircle.style.display='flex'progressNum.innerHTML='0%'// 更新进度let source = new EventSource("/progress")source.onmessage = function(event) {progressNum.innerHTML = event.data+'%'}// 发送请求const formData=new FormData()formData.append('video', file)const response=await fetch('/shrimp',{method:'POST',body:formData})if (!response.ok) {throw new Error('Network response was not ok.'); }source.close()const data=await response.json()button.disabled=falseif(data.data==='任务被终止'){alert(data.data)}else{progressCircle.style.display='none'$('#distance').text('总路程'+data.distance)// 画图var trace=[{x: data.position_data.map(item=>item[0]),y: data.position_data.map(item=>item[1]),mode:"lines",line:{color:'rgb(32, 139, 139)'}}]var layout = {xaxis: {range: [0, 600],title: "x(像素)",titlefont: { color: 'rgb(32, 139, 139)' // 轴标签颜色 }, linecolor: 'rgb(32, 139, 139)', // 轴线颜色 tickfont: { color: 'lrgb(32, 139, 139)' // 轴刻度标签颜色 }},yaxis: {range: [0, 600],title: "y(像素)",titlefont: { color: 'rgb(32, 139, 139)' // 轴标签颜色 }, linecolor: 'rgb(32, 139, 139)', // 轴线颜色 tickfont: { color: 'lrgb(32, 139, 139)' // 轴刻度标签颜色 }}, title: "鳌虾运动轨迹",titlefont:{color:'rgb(32, 139, 139)'},plot_bgcolor: 'rgba(0,0,0,0)',paper_bgcolor:'rgba(0,0,0,0)'}Plotly.newPlot("s2", trace, layout,{scrollZoom: true,editable: true }) }}catch(error){button.disabled = falseif(error.message.startsWith('Failed to fetch')){}else{alert(error)}}}</script><fieldset><legend>输出结果</legend><P id="distance">总路程:</P></fieldset><fieldset><legend>注意事项</legend><p>本程序运行将消耗大量算力和内存,最好使用高配电脑。不支持windows10以下的操作系统。在后台有任务在跑时,切勿重复上传视频,等待后台跑完出图时再上传新的视频。如果选错视频并上传了,请点击'终止运行'再重新上传视频。有问题联系wx:m989783106</p></fieldset></section><section id="s2"><div id="progress-circle"><p id="progress-num"></p></div></section></main>
</body>
</html>plotly.js从官网下载
代码分览
总体设计是以<html>和<body>为底，<main>为主容器内使用grid2列布局，2个<section>作为内容器占据左右2个网格。
左边的<section>容纳文件上传表单，功能按钮，数据显示，使用说明
<section id="s1"><form id="form" enctype="multipart/form-data"><fieldset><legend>选择你要检测的视频</legend><input type="file" accept=".mp4" id="video" name="vedio"></fieldset><fieldset><legend>功能按键</legend><button onclick="submit_to()" id="submit">开始上传</button><button onclick="stopRun()">终止运行</button></fieldset></form><fieldset><legend>输出结果</legend><P id="distance">总路程:</P></fieldset><fieldset><legend>注意事项</legend><p>本程序运行将消耗大量算力和内存,最好使用高配电脑。不支持windows10以下的操作系统。在后台有任务在跑时,切勿重复上传视频,等待后台跑完出图时再上传新的视频。如果选错视频并上传了,请点击'终止运行'再重新上传视频。有问题联系wx:m989783106</p></fieldset></section>之间用<fieldset>做了区域划分，简单又美观。
<button>均使用onclick属性进行触发
在上传前会检测用户是否选择文件，是否选择的是MP4文件
// 获取文件
const input=document.getElementById('video')
const file=input.files[0]
if (!file){throw new Error('未选择文件')
}
if(file.type!=='video/mp4'){throw new Error('请选择MP4文件')
}一共有3个请求：
- 请求上传文件，将MP4上传给后端，然后后端运行模型发送质心坐标给前端渲染
- 请求终止程序，当用户想终止后端运行模型，重新上传文件时
- 请求获取模型处理进度，后端返回进度给前端，前端进行渲染展示
画轨迹图，前端用plotly.js将质心坐标进行渲染，同时轨迹图还有一定的交互能力。
// 画图
var trace=[{x: data.position_data.map(item=>item[0]),y: data.position_data.map(item=>item[1]),mode:"lines",line:{color:'rgb(32, 139, 139)'}
}]
var layout = {xaxis: {range: [0, 600],title: "x(像素)",titlefont: { color: 'rgb(32, 139, 139)' // 轴标签颜色 }, linecolor: 'rgb(32, 139, 139)', // 轴线颜色 tickfont: { color: 'lrgb(32, 139, 139)' // 轴刻度标签颜色 }},yaxis: {range: [0, 600],title: "y(像素)",titlefont: { color: 'rgb(32, 139, 139)' // 轴标签颜色 }, linecolor: 'rgb(32, 139, 139)', // 轴线颜色 tickfont: { color: 'lrgb(32, 139, 139)' // 轴刻度标签颜色 }}, title: "鳌虾运动轨迹",titlefont:{color:'rgb(32, 139, 139)'},plot_bgcolor: 'rgba(0,0,0,0)',paper_bgcolor:'rgba(0,0,0,0)'}
Plotly.newPlot("s2", trace, layout,{scrollZoom: true,editable: true })其余的就是代码的排布顺序，异步执行调度，错误处理能力，系统稳定性，用户交互能力的提升，细节很多，均包含在代码中
右边的<section>容纳进度圈，轨迹图
<section id="s2"><div id="progress-circle"><p id="progress-num"></p></div></section>在文件上传时，就初始化渲染进度条，然后异步请求获取进度，渲染到页面；当进度到达一定值，比如99%，就关闭获取进度的请求，同时设置进度条的display=none。当用户打断程序执行或者重新运行程序，就清理轨迹图，初始化进度条，循环往复。
后端设计
后端整体使用flask，jinjia模板，将flask与pywebview结合。把模型检测代码封装到一个类TrackShrimp，其余的就是各种请求函数。
代码总览
import webview
from flask import Flask, request, jsonify,render_template,stream_with_context,Response
import os
import time
import cv2
from onnxruntime import InferenceSession
import numpy as np
from werkzeug.utils import secure_filenameclass TrackShrimp():def __init__(self,video_path,model_path,onnx_threshold=0.7):# 获取帧数据self.cap=self.init_video(video_path)frame_width=int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))frame_height=int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 图形尺寸im_shape = np.array([[frame_height, frame_width]], dtype='float32')# y轴缩放量self.im_scale_y=640.0/frame_height# x轴缩放量self.im_scale_x=640.0/frame_widthscale_factor = np.array([[self.im_scale_y,self.im_scale_x]]).astype('float32')# 定义模型输入self.inputs_dict = {'im_shape': im_shape,'image': None,'scale_factor': scale_factor}# 初始化模型self.sess=self.init_session(model_path)# 模型输出阈值self.onnx_threshold=onnx_thresholddef init_video(self,input_path):cap=cv2.VideoCapture(input_path)if not cap.isOpened():raise ValueError(f'无法打开视频{input_path}')return capdef init_session(self,model_path):try:return InferenceSession(model_path, providers=['CUDAExecutionProvider']) except:return InferenceSession(model_path, providers=['CPUExecutionProvider'])def precess_img(self,frame):img = cv2.resize(frame, None,None,fx=self.im_scale_x,fy=self.im_scale_y,interpolation=2)img = img.astype(np.float32) / 255.0img = np.transpose(img, [2, 0, 1])img = img[np.newaxis, :, :, :]return imgdef postcess(self,results:np.ndarray,all_centers:list[np.ndarray]):results=results[(results[:, 0] == 0) & (results[:, 1] > self.onnx_threshold)]x_centers = (results[:, 2] + results[:, 4]) / 2y_centers = (results[:, 3] + results[:, 5]) / 2centers = np.column_stack((x_centers, y_centers))all_centers.extend(centers)def by_smoothfilter(self,centers:list[np.ndarray],window_size=24):""":param centers: list[np.ndarray,np.ndarray,...]:param window_size: 平滑窗口大小:return: 平滑后的质心坐标NumPy数组"""centers=np.stack(centers)# 计算滑动窗口的平均值，pad函数在序列前后补零以处理边界情况padded_centers = np.pad(centers, ((window_size//2, window_size//2), (0, 0)), mode='edge')window_sum = np.cumsum(padded_centers, axis=0)smoothed_centers = (window_sum[window_size:] - window_sum[:-window_size]) / window_sizereturn smoothed_centersdef calculate_distance(self,centers:np.ndarray):'''centers:np.ndarray n*2'''# 计算相邻点之间的差diffs = centers[1:] - centers[:-1]# 计算每个差值的欧几里得距离distances = np.linalg.norm(diffs, axis=1)# 计算总路程return int(np.sum(distances))def gain_position(self,centers:np.ndarray):position_list=centers.tolist()return position_listdef run(self):global scheduleglobal run_task# 帧数frame_count=int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))frame_number=0center_list=[]for frame_number in range(frame_count):if not run_task:returnsuccess, frame = self.cap.read()if not success:breakschedule=int(frame_number/frame_count*100)# 打印进度if frame_number%10==0:print('Process: ',schedule)# 图片预处理img=self.precess_img(frame)self.inputs_dict['image']=imgresults=self.sess.run(None,self.inputs_dict)[0]if results is not None:self.postcess(results,center_list)# 使用平滑滤波filtered_centers = self.by_smoothfilter(center_list)self.cap.release()# 返回路程，轨迹坐标return self.calculate_distance(filtered_centers),self.gain_position(filtered_centers)app = Flask(__name__)
UPLOAD_FOLDER = './temp'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
schedule=0
run_task=Falsedef run_flask():app.run(debug=False, threaded=True,host='127.0.0.1',port=5000)def video_process(video_path):return TrackShrimp(video_path,'./model.onnx').run()# 主页面
@app.route('/',methods=['POST','GET'])
def return_main_page():return render_template('index.html')# 检测视频页面
@app.route('/shrimp',methods=['POST'])
def shrimp_track():global run_taskrun_task=Truefile=request.files.get('video')filename = secure_filename(file.filename)video_path=os.path.join(app.config['UPLOAD_FOLDER'],filename)file.save(video_path)try:results=video_process(video_path)if results is not None:distance,position_data=resultsdata = {'distance': distance,'position_data': position_data}run_task=Falsereturn jsonify(data)else:return jsonify({'data':'任务被终止'})except Exception as e:print('error:',e)return jsonify({'data':'任务被终止'})finally:if os.path.exists(video_path):os.remove(video_path)@app.route('/stopRun',methods=['GET','POST'])
def stopRun():global run_taskglobal scheduleif run_task:run_task=Falseschedule=0return jsonify({'data':'正在停止任务'})else:return jsonify({'data':'当前没有任务运行'})# 进度查询路由
@app.route('/progress',methods=['GET'])
def progress():@stream_with_contextdef generate():global run_taskratio = schedulewhile ratio < 95 and run_task:yield "data:" + str(ratio) + "\n\n"ratio = scheduletime.sleep(5)return Response(generate(), mimetype='text/event-stream')if __name__=='__main__':# 启动后端 # flask_thread = threading.Thread(target=run_flask) # flask_thread.start()# time.sleep(1)# 启动前端webview.create_window('鳌虾轨迹侦测',url=app,width=900,height=600)# webview.create_window('鳌虾轨迹侦测',url=f'http://127.0.0.1:5000',width=900,height=600)webview.start()
代码分览
一个onnx部署的类TrackShrimp，详细见前面。
一些常量的定义
app = Flask(__name__)
UPLOAD_FOLDER = './temp' # 文件的上传路径，后端需要该路径保留用户上传的文件
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
schedule=0 # 实时进度，初始化进度为0
run_task=False # 一个onnx模型是否在运行的标志，用于接收用户中断信号从而终止模型运行定义一个flask的·启动函数，用于web调试，浏览器F12启动调试窗口
def run_flask():app.run(debug=False, threaded=True,host='127.0.0.1',port=5000)主页面的请求函数，该页面为主要的UI
# 主页面
@app.route('/',methods=['POST','GET'])
def return_main_page():return render_template('index.html')用户请求中断的请求函数
首先通过标志位（run_task）检测模型是否在跑，如果检测到模型正在运行，就把标志位设为False，然后把进度归0
@app.route('/stopRun',methods=['GET','POST'])
def stopRun():global run_taskglobal scheduleif run_task:run_task=Falseschedule=0return jsonify({'data':'正在停止任务'})else:return jsonify({'data':'当前没有任务运行'})进度查询
这里设置当进度为95%时，就停止查询。
@app.route('/progress',methods=['GET'])
def progress():@stream_with_contextdef generate():global run_taskratio = schedulewhile ratio < 95 and run_task:yield "data:" + str(ratio) + "\n\n"ratio = scheduletime.sleep(5)return Response(generate(), mimetype='text/event-stream')一个检测的入口函数
def video_process(video_path):return TrackShrimp(video_path,'./model.onnx').run()接收用户上传文件的函数
一旦用户上传文件，就设置运行标志位为True，然后将文件保存，再送入模型运行接口函数，当用户请求终止时，results为None，所以使用if else进行区分。模型结果出来后就把标志位设为False，同时将数据传到前端
@app.route('/shrimp',methods=['POST'])
def shrimp_track():global run_taskrun_task=Truefile=request.files.get('video')filename = secure_filename(file.filename)video_path=os.path.join(app.config['UPLOAD_FOLDER'],filename)file.save(video_path)try:results=video_process(video_path)if results is not None:distance,position_data=resultsdata = {'distance': distance,'position_data': position_data}run_task=Falsereturn jsonify(data)else:return jsonify({'data':'任务被终止'})except Exception as e:print('error:',e)return jsonify({'data':'任务被终止'})finally:if os.path.exists(video_path):os.remove(video_path)接着就是启动所有代码了，为了调试方便，我写了2份代码，一份用于调试，一份用于成品
if __name__=='__main__':# 启动前端webview.create_window('鳌虾轨迹侦测',url=app,width=900,height=600)webview.start()if __name__=='__main__':# 启动后端 flask_thread = threading.Thread(target=run_flask) flask_thread.start()time.sleep(1)# 启动前端webview.create_window('鳌虾轨迹侦测',url=f'http://127.0.0.1:5000',width=900,height=600)webview.start()pyinstaller打包
进入项目目录，命令行输入
piinstaller -D -w main.py找到生成的main.spec文件，按如下修改
# -*- mode: python ; coding: utf-8 -*-a = Analysis(['main.py'],pathex=[],binaries=[],datas=[('templates/','templates/'),('static/','static/'),('venv/Lib/site-packages/onnxruntime/capi/onnxruntime_providers_shared.dll','onnxruntime/capi/'),('venv/Lib/site-packages/onnxruntime/capi/onnxruntime_providers_cuda.dll','onnxruntime/capi/')],hiddenimports=[],hookspath=[],hooksconfig={},runtime_hooks=[],excludes=[],noarchive=False,optimize=0,
)
pyz = PYZ(a.pure)exe = EXE(pyz,a.scripts,[],exclude_binaries=True,name='main',debug=False,bootloader_ignore_signals=False,strip=False,upx=True,console=False,disable_windowed_traceback=False,argv_emulation=False,target_arch=None,codesign_identity=None,entitlements_file=None,icon='1.ico'
)
coll = COLLECT(exe,a.binaries,a.datas,strip=False,upx=True,upx_exclude=[],name='main',
)
在项目目录下放置一个图标命名为1.ico，最好是48*48像素
然后命令行运行
pyinstaller main.spec然后在venv中找到 onnxruntime_gpu-1.18.1.dist-info 文件夹，复制到 dist/main/_internal 中
同时在cuDNN中找到如下几个动态链接库，复制到 dist/main/_internal 中
cudnn_ops_infer64_8.dll
cudnn_cnn_infer64_8.dll
cudnn_adv_infer64_8.dll
cudnn64_8.dll
cudart64_110.dll
cublasLt64_11.dll
cublas64_11.dll
cufft64_10.dll然后将model.onnx放到 dist/main/ ，并在该目录创建一个目录temp
最后处理的结果如下
XXXdistmain_internalmain.exemodel.onnxtemp生成安装包
使用into setup软件，并在网站找到中文的语言包下载为 Chinese.isl 文件，放到intosetup软件安装目录的 Languages 文件夹下
接着如图所示
取消立即编译，先进入文件里修改一些东西
修改成下面这样
点击编译
然后就生成了安装包，就可以在任何win10，win11电脑里用CPU跑了，如果安装的电脑 有显卡和CUDA并把CUDA添加到了环境变量，就可以用GPU跑了
![[Unity入门01] Unity基本操作](https://i-blog.csdnimg.cn/direct/fc99f0c9da1c4aaf9efd229e1eef51b9.png)
