duan8/yolov8/include/model.h
lindsayshuo 01191f5e61
yolov8 cls (#1453)
* yolov8 cls

* Update README.md

* Update yolov8_cls.cpp

* Update block.h

* Update block.cpp

* Update yolov8_cls.cpp

---------

Co-authored-by: Wang Xinyu <wangxinyu_es@163.com>
2024-03-12 17:53:07 +08:00

14 lines
659 B
C++

#pragma once
#include "NvInfer.h"
#include <string>
#include <assert.h>
nvinfer1::IHostMemory* buildEngineYolov8Det(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw, int& max_channels);
nvinfer1::IHostMemory* buildEngineYolov8Cls(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw);
nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder,
nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path, float& gd, float& gw, int& max_channels);