diff --git a/algorithm/common_det/cuda/common_det_cuda_impl.cpp b/algorithm/common_det/cuda/common_det_cuda_impl.cpp index 5f3d7a1..44c503b 100644 --- a/algorithm/common_det/cuda/common_det_cuda_impl.cpp +++ b/algorithm/common_det/cuda/common_det_cuda_impl.cpp @@ -366,7 +366,7 @@ bool CommonObjectDetectorCUDAImpl::cudaSetup(CommonObjectDetectorBase* base_, bo std::string engine_fn = get_home() + SV_MODEL_DIR + dataset + ".engine"; std::vector 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); diff --git a/algorithm/common_det/intel/common_det_intel_impl.cpp b/algorithm/common_det/intel/common_det_intel_impl.cpp index c61c507..2ca25cb 100644 --- a/algorithm/common_det/intel/common_det_intel_impl.cpp +++ b/algorithm/common_det/intel/common_det_intel_impl.cpp @@ -1,6 +1,7 @@ #include "common_det_intel_impl.h" #include #include +#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 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)) diff --git a/algorithm/landing_det/cuda/landing_det_cuda_impl.cpp b/algorithm/landing_det/cuda/landing_det_cuda_impl.cpp index 97fa80c..5c8c1bc 100644 --- a/algorithm/landing_det/cuda/landing_det_cuda_impl.cpp +++ b/algorithm/landing_det/cuda/landing_det_cuda_impl.cpp @@ -53,7 +53,7 @@ bool LandingMarkerDetectorCUDAImpl::cudaSetup() std::string trt_model_fn = get_home() + SV_MODEL_DIR + "LandingMarker.engine"; std::vector 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); diff --git a/algorithm/landing_det/intel/landing_det_intel_impl.cpp b/algorithm/landing_det/intel/landing_det_intel_impl.cpp index a989ebb..324ad54 100644 --- a/algorithm/landing_det/intel/landing_det_intel_impl.cpp +++ b/algorithm/landing_det/intel/landing_det_intel_impl.cpp @@ -26,6 +26,14 @@ namespace sv { #ifdef WITH_INTEL std::string onnx_model_fn = get_home() + SV_MODEL_DIR + "LandingMarker.onnx"; + std::vector 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)"); diff --git a/algorithm/sot/ocv470/sot_ocv470_impl.cpp b/algorithm/sot/ocv470/sot_ocv470_impl.cpp index 027ae2c..2d96225 100644 --- a/algorithm/sot/ocv470/sot_ocv470_impl.cpp +++ b/algorithm/sot/ocv470/sot_ocv470_impl.cpp @@ -33,9 +33,9 @@ bool SingleObjectTrackerOCV470Impl::ocv470Setup(SingleObjectTrackerBase* base_) std::string kernel_r1 = get_home() + SV_MODEL_DIR + "dasiamrpn_kernel_r1.onnx"; std::vector 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 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); diff --git a/algorithm/veri/cuda/veri_det_cuda_impl.cpp b/algorithm/veri/cuda/veri_det_cuda_impl.cpp index c42c722..1872149 100644 --- a/algorithm/veri/cuda/veri_det_cuda_impl.cpp +++ b/algorithm/veri/cuda/veri_det_cuda_impl.cpp @@ -78,7 +78,7 @@ namespace sv std::string trt_model_fn = get_home() + SV_MODEL_DIR + "veri.engine"; std::vector 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); diff --git a/scripts/model_sync.py b/scripts/model_sync.py new file mode 100644 index 0000000..1068b44 --- /dev/null +++ b/scripts/model_sync.py @@ -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()