floraachy
/
SpireCV
forked from floratest1/SpireCV
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model manager v1.0 & add model_sync.py

This commit is contained in:
jario-jin 2024-01-17 11:00:54 +08:00
parent 62708c742c
commit 5108de8df4
7 changed files with 155 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
algorithm/common_det/cuda/common_det_cuda_impl.cpp
View File

@ -366,7 +366,7 @@ bool CommonObjectDetectorCUDAImpl::cudaSetup(CommonObjectDetectorBase* base_, bo
std::string engine_fn = get_home() + SV_MODEL_DIR + dataset + ".engine";
std::vector<std::string> files;
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", dataset + "-yolov5s-");
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "Nv-" + dataset + "-yolov5s_c");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);
@ -376,7 +376,7 @@ bool CommonObjectDetectorCUDAImpl::cudaSetup(CommonObjectDetectorBase* base_, bo
if (input_w == 1280)
{
files.clear();
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", dataset + "-yolov5s6-");
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "Nv-" + dataset + "-yolov5s6_");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);
@ -392,7 +392,7 @@ bool CommonObjectDetectorCUDAImpl::cudaSetup(CommonObjectDetectorBase* base_, bo
base_->setInputH(640);
base_->setInputW(640);
files.clear();
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", dataset + "-yolov5s-seg-");
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "Nv-" + dataset + "-yolov5s_seg_");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);

33
algorithm/common_det/intel/common_det_intel_impl.cpp
View File

@ -1,6 +1,7 @@
#include "common_det_intel_impl.h"
#include <cmath>
#include <fstream>
#include "sv_util.h"
#define SV_MODEL_DIR "/SpireCV/models/"
#define SV_ROOT_DIR "/SpireCV/"
@ -126,15 +127,43 @@ namespace sv
with_segmentation = base_->withSegmentation();
std::string openvino_fn = get_home() + SV_MODEL_DIR + dataset + ".onnx";
std::vector<std::string> files;
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.onnx", "Int-" + dataset + "-yolov5s_c");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);
openvino_fn = get_home() + SV_MODEL_DIR + files[0];
}
if (inpWidth == 1280)
{
openvino_fn = get_home() + SV_MODEL_DIR + dataset + "_HD.onnx";
files.clear();
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.onnx", "Int-" + dataset + "-yolov5s6_");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);
openvino_fn = get_home() + SV_MODEL_DIR + files[0];
}
else
{
openvino_fn = get_home() + SV_MODEL_DIR + dataset + "_HD.onnx";
}
}
if (with_segmentation)
{
base_->setInputH(640);
base_->setInputW(640);
openvino_fn = get_home() + SV_MODEL_DIR + dataset + "_SEG.onnx";
files.clear();
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.onnx", "Int-" + dataset + "-yolov5s_seg_");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);
openvino_fn = get_home() + SV_MODEL_DIR + files[0];
}
else
{
openvino_fn = get_home() + SV_MODEL_DIR + dataset + "_SEG.onnx";
}
}
std::cout << "Load: " << openvino_fn << std::endl;
if (!is_file_exist(openvino_fn))

2
algorithm/landing_det/cuda/landing_det_cuda_impl.cpp
View File

@ -53,7 +53,7 @@ bool LandingMarkerDetectorCUDAImpl::cudaSetup()
std::string trt_model_fn = get_home() + SV_MODEL_DIR + "LandingMarker.engine";
std::vector<std::string> files;
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "LandingMarker-");
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "Nv-LandingMarker-resnet34");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);

8
algorithm/landing_det/intel/landing_det_intel_impl.cpp
View File

@ -26,6 +26,14 @@ namespace sv
{
#ifdef WITH_INTEL
std::string onnx_model_fn = get_home() + SV_MODEL_DIR + "LandingMarker.onnx";
std::vector<std::string> files;
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.onnx", "Int-LandingMarker-resnet34");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);
onnx_model_fn = get_home() + SV_MODEL_DIR + files[0];
}
if (!is_file_exist(onnx_model_fn))
{
throw std::runtime_error("SpireCV (104) Error loading the LandingMarker ONNX model (File Not Exist)");

10
algorithm/sot/ocv470/sot_ocv470_impl.cpp
View File

@ -33,9 +33,9 @@ bool SingleObjectTrackerOCV470Impl::ocv470Setup(SingleObjectTrackerBase* base_)
std::string kernel_r1 = get_home() + SV_MODEL_DIR + "dasiamrpn_kernel_r1.onnx";
std::vector<std::string> files1, files2, files3;
_list_dir(get_home() + SV_MODEL_DIR, files1, "-online.engine", "DaSiamRPN-Model-");
_list_dir(get_home() + SV_MODEL_DIR, files2, "-online.engine", "DaSiamRPN-Kernel-CLS1-");
_list_dir(get_home() + SV_MODEL_DIR, files3, "-online.engine", "DaSiamRPN-Kernel-R1-");
_list_dir(get_home() + SV_MODEL_DIR, files1, ".onnx", "Ocv-DaSiamRPN-Model-");
_list_dir(get_home() + SV_MODEL_DIR, files2, ".onnx", "Ocv-DaSiamRPN-Kernel-CLS1-");
_list_dir(get_home() + SV_MODEL_DIR, files3, ".onnx", "Ocv-DaSiamRPN-Kernel-R1-");
if (files1.size() > 0 && files2.size() > 0 && files3.size() > 0)
{
std::sort(files1.rbegin(), files1.rend(), _comp_str_lesser);
@ -53,8 +53,8 @@ bool SingleObjectTrackerOCV470Impl::ocv470Setup(SingleObjectTrackerBase* base_)
std::string neckhead = get_home() + SV_MODEL_DIR + "nanotrack_head_sim.onnx";
std::vector<std::string> files4, files5;
_list_dir(get_home() + SV_MODEL_DIR, files4, "-online.engine", "NanoTrack-Backbone-SIM-");
_list_dir(get_home() + SV_MODEL_DIR, files5, "-online.engine", "NanoTrack-Head-SIM-");
_list_dir(get_home() + SV_MODEL_DIR, files4, ".onnx", "Ocv-NanoTrack-Backbone-SIM-");
_list_dir(get_home() + SV_MODEL_DIR, files5, ".onnx", "Ocv-NanoTrack-Head-SIM-");
if (files4.size() > 0 && files5.size() > 0)
{
std::sort(files4.rbegin(), files4.rend(), _comp_str_lesser);

2
algorithm/veri/cuda/veri_det_cuda_impl.cpp
View File

@ -78,7 +78,7 @@ namespace sv
std::string trt_model_fn = get_home() + SV_MODEL_DIR + "veri.engine";
std::vector<std::string> files;
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "VERI-");
_list_dir(get_home() + SV_MODEL_DIR, files, "-online.engine", "Nv-VERI-mobilenet_v3");
if (files.size() > 0)
{
std::sort(files.rbegin(), files.rend(), _comp_str_lesser);

106
scripts/model_sync.py Normal file
View File

@ -0,0 +1,106 @@
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import os
import requests
import argparse
root_url = "https://download.amovlab.com/model/SpireCV-models/"
model_list_url = root_url + "model-list.txt"
root_path = "/home/amov/SpireCV/models"
list_file = os.path.join(root_path, "model-list.txt")
def main():
parser = argparse.ArgumentParser(description="SpireCV Model SYNC")
parser.add_argument(
"-p", "--platform",
type=str,
required=True,
help="Supported Platforms: nv (Nvidia), int (Intel)",
)
args = parser.parse_args()
if args.platform in ['nv', 'nvidia', 'Nv', 'Nvidia']:
prefix = 'Nv'
elif args.platform in ['int', 'intel', 'Int', 'Intel']:
prefix = 'Int'
else:
raise Exception("Platform NOT Support!")
r = requests.get(model_list_url)
with open(list_file, "wb") as f:
f.write(r.content)
with open(list_file, 'r') as file:
lines = file.readlines()
need_switch = False
for line in lines:
line = line.strip()
if len(line) > 0:
model_file = os.path.join(root_path, line)
if not os.path.exists(model_file) and (line.startswith(prefix) or line.startswith('Ocv')):
print("[1] Downloading Model:", line, "...")
r = requests.get(root_url + line)
with open(model_file, "wb") as f:
f.write(r.content)
need_switch = True
if os.path.exists(model_file):
print("[1] Model:", line, "EXIST!")
if line.startswith('Nv'):
engine_fn = os.path.splitext(model_file)[0] + '.engine'
online_fn = os.path.splitext(model_file)[0] + '-online.engine'
if not os.path.exists(engine_fn) and not os.path.exists(online_fn):
net = line.split('-')[2]
if net.startswith("yolov5"):
if len(net.split('_')) == 3:
name, seg, ncls = net.split('_')
cmd = "SpireCVSeg -s {} {} {} {}".format(
model_file, engine_fn, ncls[1:], name[6:]
)
else:
name, ncls = net.split('_')
cmd = "SpireCVDet -s {} {} {} {}".format(
model_file, engine_fn, ncls[1:], name[6:]
)
elif line.endswith("onnx"):
cmd = ("/usr/src/tensorrt/bin/trtexec --explicitBatch --onnx={} "
"--saveEngine={} --fp16").format(
model_file, engine_fn
)
print(" [2] Converting Model:", line, "->", engine_fn, "...")
result = os.popen(cmd).read()
need_switch = True
else:
print(" [2] Model Converting FINISH!")
model_file = engine_fn
if not line.startswith('Ocv') and need_switch:
ext = os.path.splitext(model_file)[1]
fn_prefix = '-'.join(os.path.basename(model_file).split('-')[:3])
file_names = os.listdir(root_path)
selected = []
for file_name in file_names:
if file_name.startswith(fn_prefix) and file_name.endswith(ext):
selected.append(file_name)
if len(selected) > 0:
for i, sel in enumerate(selected):
if sel.endswith('-online' + ext):
os.rename(
os.path.join(root_path, sel),
os.path.join(root_path, '-'.join(sel.split('-')[:4])) + ext
)
selected[i] = '-'.join(sel.split('-')[:4]) + ext
selected.sort(reverse=True)
os.rename(
os.path.join(root_path, selected[0]),
os.path.join(root_path, os.path.splitext(selected[0])[0] + "-online" + ext)
)
online_model = os.path.splitext(selected[0])[0] + "-online" + ext
print(" [3] Model {} ONLINE *".format(online_model))
if __name__ == "__main__":
main()