Merge pull request #126 from triple-Mu/dev

Support YOLOv8 pose model inference with Python

config.py

@@ -36,5 +36,20 @@ MASK_COLORS = np.array([(255, 56, 56), (255, 157, 151), (255, 112, 31),
                         (255, 149, 200), (255, 55, 199)],
                        dtype=np.float32) / 255.
 
+KPS_COLORS = [[0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0],
+              [255, 128, 0], [255, 128, 0], [255, 128, 0], [255, 128, 0],
+              [255, 128, 0], [255, 128, 0], [51, 153, 255], [51, 153, 255],
+              [51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255]]
+
+SKELETON = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13],
+            [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3],
+            [2, 4], [3, 5], [4, 6], [5, 7]]
+
+LIMB_COLORS = [[51, 153, 255], [51, 153, 255], [51, 153, 255], [51, 153, 255],
+               [255, 51, 255], [255, 51, 255], [255, 51, 255], [255, 128, 0],
+               [255, 128, 0], [255, 128, 0], [255, 128, 0], [255, 128, 0],
+               [0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0], [0, 255, 0],
+               [0, 255, 0], [0, 255, 0]]
+
 # alpha for segment masks
 ALPHA = 0.5

csrc/deepstream/custom_bbox_parser/nvdsparsebbox_yoloV8.cpp

@@ -18,95 +18,85 @@
  */
 
 // This is just the function prototype. The definition is written at the end of the file.
-extern "C" bool NvDsInferParseCustomYoloV8(
+extern "C" bool NvDsInferParseCustomYoloV8(std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
-	std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
+                                           NvDsInferNetworkInfo const&            networkInfo,
-	NvDsInferNetworkInfo const& networkInfo,
+                                           NvDsInferParseDetectionParams const&   detectionParams,
-	NvDsInferParseDetectionParams const& detectionParams,
+                                           std::vector<NvDsInferParseObjectInfo>& objectList);
-	std::vector<NvDsInferParseObjectInfo>& objectList);
 
 static __inline__ float bbox_clip(const float& val, const float& minVal = 0.f, const float& maxVal = 1280.f)
 {
-	assert(minVal <= maxVal);
+    assert(minVal <= maxVal);
-	return std::max(std::min(val, (maxVal - 1)), minVal);
+    return std::max(std::min(val, (maxVal - 1)), minVal);
 }
 
-static std::vector<NvDsInferParseObjectInfo> decodeYoloV8Tensor(
+static std::vector<NvDsInferParseObjectInfo> decodeYoloV8Tensor(const int*          num_dets,
-	const int* num_dets,
+                                                                const float*        bboxes,
-	const float* bboxes,
+                                                                const float*        scores,
-	const float* scores,
+                                                                const int*          labels,
-	const int* labels,
+                                                                const unsigned int& img_w,
-	const unsigned int& img_w,
+                                                                const unsigned int& img_h)
-	const unsigned int& img_h
-)
 {
-	std::vector<NvDsInferParseObjectInfo> bboxInfo;
+    std::vector<NvDsInferParseObjectInfo> bboxInfo;
-	size_t nums = num_dets[0];
+    size_t                                nums = num_dets[0];
-	for (size_t i = 0; i < nums; i++)
+    for (size_t i = 0; i < nums; i++) {
-	{
+        float x0 = (bboxes[i * 4]);
-		float x0 = (bboxes[i * 4]);
+        float y0 = (bboxes[i * 4 + 1]);
-		float y0 = (bboxes[i * 4 + 1]);
+        float x1 = (bboxes[i * 4 + 2]);
-		float x1 = (bboxes[i * 4 + 2]);
+        float y1 = (bboxes[i * 4 + 3]);
-		float y1 = (bboxes[i * 4 + 3]);
+        x0       = bbox_clip(x0, 0.f, img_w);
-		x0 = bbox_clip(x0, 0.f, img_w);
+        y0       = bbox_clip(y0, 0.f, img_h);
-		y0 = bbox_clip(y0, 0.f, img_h);
+        x1       = bbox_clip(x1, 0.f, img_w);
-		x1 = bbox_clip(x1, 0.f, img_w);
+        y1       = bbox_clip(y1, 0.f, img_h);
-		y1 = bbox_clip(y1, 0.f, img_h);
+        NvDsInferParseObjectInfo obj;
-		NvDsInferParseObjectInfo obj;
+        obj.left                = x0;
-		obj.left = x0;
+        obj.top                 = y0;
-		obj.top = y0;
+        obj.width               = x1 - x0;
-		obj.width = x1 - x0;
+        obj.height              = y1 - y0;
-		obj.height = y1 - y0;
+        obj.detectionConfidence = scores[i];
-		obj.detectionConfidence = scores[i];
+        obj.classId             = labels[i];
-		obj.classId = labels[i];
+        bboxInfo.push_back(obj);
-		bboxInfo.push_back(obj);
+    }
-	}
 
-	return bboxInfo;
+    return bboxInfo;
 }
 
 /* C-linkage to prevent name-mangling */
-extern "C" bool NvDsInferParseCustomYoloV8(
+extern "C" bool NvDsInferParseCustomYoloV8(std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
-	std::vector<NvDsInferLayerInfo> const& outputLayersInfo,
+                                           NvDsInferNetworkInfo const&            networkInfo,
-	NvDsInferNetworkInfo const& networkInfo,
+                                           NvDsInferParseDetectionParams const&   detectionParams,
-	NvDsInferParseDetectionParams const& detectionParams,
+                                           std::vector<NvDsInferParseObjectInfo>& objectList)
-	std::vector<NvDsInferParseObjectInfo>& objectList)
 {
 
-	// Some assertions and error checking.
+    // Some assertions and error checking.
 
-	if (outputLayersInfo.empty() || outputLayersInfo.size() != 4)
+    if (outputLayersInfo.empty() || outputLayersInfo.size() != 4) {
-	{
+        std::cerr << "Could not find output layer in bbox parsing" << std::endl;
-		std::cerr << "Could not find output layer in bbox parsing" << std::endl;
+        return false;
-		return false;
+    }
-	}
 
-	// Obtaining the output layer.
+    // Obtaining the output layer.
-	const NvDsInferLayerInfo& num_dets = outputLayersInfo[0];
+    const NvDsInferLayerInfo& num_dets = outputLayersInfo[0];
-	const NvDsInferLayerInfo& bboxes = outputLayersInfo[1];
+    const NvDsInferLayerInfo& bboxes   = outputLayersInfo[1];
-	const NvDsInferLayerInfo& scores = outputLayersInfo[2];
+    const NvDsInferLayerInfo& scores   = outputLayersInfo[2];
-	const NvDsInferLayerInfo& labels = outputLayersInfo[3];
+    const NvDsInferLayerInfo& labels   = outputLayersInfo[3];
 
-	// num_dets(int) bboxes(float) scores(float) labels(int)
+    // num_dets(int) bboxes(float) scores(float) labels(int)
-	assert (num_dets.dims.numDims == 2);
+    assert(num_dets.dims.numDims == 2);
-	assert (bboxes.dims.numDims == 3);
+    assert(bboxes.dims.numDims == 3);
-	assert (scores.dims.numDims == 2);
+    assert(scores.dims.numDims == 2);
-	assert (labels.dims.numDims == 2);
+    assert(labels.dims.numDims == 2);
 
+    // Decoding the output tensor of YOLOv8 to the NvDsInferParseObjectInfo format.
+    std::vector<NvDsInferParseObjectInfo> objects = decodeYoloV8Tensor((const int*)(num_dets.buffer),
+                                                                       (const float*)(bboxes.buffer),
+                                                                       (const float*)(scores.buffer),
+                                                                       (const int*)(labels.buffer),
+                                                                       networkInfo.width,
+                                                                       networkInfo.height);
 
-	// Decoding the output tensor of YOLOv8 to the NvDsInferParseObjectInfo format.
+    objectList.clear();
-	std::vector<NvDsInferParseObjectInfo> objects =
+    objectList = objects;
-		decodeYoloV8Tensor(
+    return true;
-			(const int*)(num_dets.buffer),
-			(const float*)(bboxes.buffer),
-			(const float*)(scores.buffer),
-			(const int*)(labels.buffer),
-			networkInfo.width,
-			networkInfo.height
-		);
-
-	objectList.clear();
-	objectList = objects;
-	return true;
 }
 
 /* Check that the custom function has been defined correctly */

csrc/detect/end2end/include/common.hpp

@@ -4,153 +4,138 @@
 
 #ifndef DETECT_END2END_COMMON_HPP
 #define DETECT_END2END_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace det
+namespace det {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float> rect;
-		cv::Rect_<float> rect;
+    int              label = 0;
-		int label = 0;
+    float            prob  = 0.0;
-		float prob = 0.0;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace det
-}
+#endif  // DETECT_END2END_COMMON_HPP
-#endif //DETECT_END2END_COMMON_HPP

csrc/detect/end2end/include/yolov8.hpp

@@ -3,424 +3,281 @@
 //
 #ifndef DETECT_END2END_YOLOV8_HPP
 #define DETECT_END2END_YOLOV8_HPP
-#include "fstream"
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include "fstream"
 using namespace det;
 
-class YOLOv8
+class YOLOv8 {
-{
 public:
-	explicit YOLOv8(const std::string& engine_file_path);
+    explicit YOLOv8(const std::string& engine_file_path);
-	~YOLOv8();
+    ~YOLOv8();
-
+
-	void make_pipe(bool warmup = true);
+    void                 make_pipe(bool warmup = true);
-	void copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image);
-	void copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
-	void letterbox(
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-		const cv::Mat& image,
+    void                 infer();
-		cv::Mat& out,
+    void                 postprocess(std::vector<Object>& objs);
-		cv::Size& size
+    static void          draw_objects(const cv::Mat&                                image,
-	);
+                                      cv::Mat&                                      res,
-	void infer();
+                                      const std::vector<Object>&                    objs,
-	void postprocess(std::vector<Object>& objs);
+                                      const std::vector<std::string>&               CLASS_NAMES,
-	static void draw_objects(
+                                      const std::vector<std::vector<unsigned int>>& COLORS);
-		const cv::Mat& image,
+    int                  num_bindings;
-		cv::Mat& res,
+    int                  num_inputs  = 0;
-		const std::vector<Object>& objs,
+    int                  num_outputs = 0;
-		const std::vector<std::string>& CLASS_NAMES,
+    std::vector<Binding> input_bindings;
-		const std::vector<std::vector<unsigned int>>& COLORS
+    std::vector<Binding> output_bindings;
-	);
+    std::vector<void*>   host_ptrs;
-	int num_bindings;
+    std::vector<void*>   device_ptrs;
-	int num_inputs = 0;
+
-	int num_outputs = 0;
+    PreParam pparam;
-	std::vector<Binding> input_bindings;
-	std::vector<Binding> output_bindings;
-	std::vector<void*> host_ptrs;
-	std::vector<void*> device_ptrs;
-
-	PreParam pparam;
-private:
-	nvinfer1::ICudaEngine* engine = nullptr;
-	nvinfer1::IRuntime* runtime = nullptr;
-	nvinfer1::IExecutionContext* context = nullptr;
-	cudaStream_t stream = nullptr;
-	Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
 YOLOv8::YOLOv8(const std::string& engine_file_path)
 {
-	std::ifstream file(engine_file_path, std::ios::binary);
+    std::ifstream file(engine_file_path, std::ios::binary);
-	assert(file.good());
+    assert(file.good());
-	file.seekg(0, std::ios::end);
+    file.seekg(0, std::ios::end);
-	auto size = file.tellg();
+    auto size = file.tellg();
-	file.seekg(0, std::ios::beg);
+    file.seekg(0, std::ios::beg);
-	char* trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
-	assert(trtModelStream);
+    assert(trtModelStream);
-	file.read(trtModelStream, size);
+    file.read(trtModelStream, size);
-	file.close();
+    file.close();
-	initLibNvInferPlugins(&this->gLogger, "");
+    initLibNvInferPlugins(&this->gLogger, "");
-	this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
-	assert(this->runtime != nullptr);
+    assert(this->runtime != nullptr);
-
+
-	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
-	assert(this->engine != nullptr);
+    assert(this->engine != nullptr);
-	delete[] trtModelStream;
+    delete[] trtModelStream;
-	this->context = this->engine->createExecutionContext();
+    this->context = this->engine->createExecutionContext();
-
+
-	assert(this->context != nullptr);
+    assert(this->context != nullptr);
-	cudaStreamCreate(&this->stream);
+    cudaStreamCreate(&this->stream);
-	this->num_bindings = this->engine->getNbBindings();
+    this->num_bindings = this->engine->getNbBindings();
-
+
-	for (int i = 0; i < this->num_bindings; ++i)
+    for (int i = 0; i < this->num_bindings; ++i) {
-	{
+        Binding            binding;
-		Binding binding;
+        nvinfer1::Dims     dims;
-		nvinfer1::Dims dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-		nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
-		std::string name = this->engine->getBindingName(i);
+        binding.name             = name;
-		binding.name = name;
+        binding.dsize            = type_to_size(dtype);
-		binding.dsize = type_to_size(dtype);
+
-
+        bool IsInput = engine->bindingIsInput(i);
-		bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
-		if (IsInput)
+            this->num_inputs += 1;
-		{
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-			this->num_inputs += 1;
+            binding.size = get_size_by_dims(dims);
-			dims = this->engine->getProfileDimensions(
+            binding.dims = dims;
-				i,
+            this->input_bindings.push_back(binding);
-				0,
+            // set max opt shape
-				nvinfer1::OptProfileSelector::kMAX);
+            this->context->setBindingDimensions(i, dims);
-			binding.size = get_size_by_dims(dims);
+        }
-			binding.dims = dims;
+        else {
-			this->input_bindings.push_back(binding);
+            dims         = this->context->getBindingDimensions(i);
-			// set max opt shape
+            binding.size = get_size_by_dims(dims);
-			this->context->setBindingDimensions(i, dims);
+            binding.dims = dims;
-
+            this->output_bindings.push_back(binding);
-		}
+            this->num_outputs += 1;
-		else
+        }
-		{
+    }
-			dims = this->context->getBindingDimensions(i);
-			binding.size = get_size_by_dims(dims);
-			binding.dims = dims;
-			this->output_bindings.push_back(binding);
-			this->num_outputs += 1;
-		}
-	}
-
 }
 
 YOLOv8::~YOLOv8()
 {
-	this->context->destroy();
+    this->context->destroy();
-	this->engine->destroy();
+    this->engine->destroy();
-	this->runtime->destroy();
+    this->runtime->destroy();
-	cudaStreamDestroy(this->stream);
+    cudaStreamDestroy(this->stream);
-	for (auto& ptr : this->device_ptrs)
+    for (auto& ptr : this->device_ptrs) {
-	{
+        CHECK(cudaFree(ptr));
-		CHECK(cudaFree(ptr));
+    }
-	}
+
-
+    for (auto& ptr : this->host_ptrs) {
-	for (auto& ptr : this->host_ptrs)
+        CHECK(cudaFreeHost(ptr));
-	{
+    }
-		CHECK(cudaFreeHost(ptr));
-	}
-
 }
 void YOLOv8::make_pipe(bool warmup)
 {
 
-	for (auto& bindings : this->input_bindings)
+    for (auto& bindings : this->input_bindings) {
-	{
+        void* d_ptr;
-		void* d_ptr;
+        CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream));
-		CHECK(cudaMallocAsync(
+        this->device_ptrs.push_back(d_ptr);
-			&d_ptr,
+    }
-			bindings.size * bindings.dsize,
+
-			this->stream)
+    for (auto& bindings : this->output_bindings) {
-		);
+        void * d_ptr, *h_ptr;
-		this->device_ptrs.push_back(d_ptr);
+        size_t size = bindings.size * bindings.dsize;
-	}
+        CHECK(cudaMallocAsync(&d_ptr, size, this->stream));
-
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-	for (auto& bindings : this->output_bindings)
+        this->device_ptrs.push_back(d_ptr);
-	{
+        this->host_ptrs.push_back(h_ptr);
-		void* d_ptr, * h_ptr;
+    }
-		size_t size = bindings.size * bindings.dsize;
+
-		CHECK(cudaMallocAsync(
+    if (warmup) {
-			&d_ptr,
+        for (int i = 0; i < 10; i++) {
-			size,
+            for (auto& bindings : this->input_bindings) {
-			this->stream)
+                size_t size  = bindings.size * bindings.dsize;
-		);
+                void*  h_ptr = malloc(size);
-		CHECK(cudaHostAlloc(
+                memset(h_ptr, 0, size);
-			&h_ptr,
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-			size,
+                free(h_ptr);
-			0)
+            }
-		);
+            this->infer();
-		this->device_ptrs.push_back(d_ptr);
+        }
-		this->host_ptrs.push_back(h_ptr);
+        printf("model warmup 10 times\n");
-	}
+    }
-
-	if (warmup)
-	{
-		for (int i = 0; i < 10; i++)
-		{
-			for (auto& bindings : this->input_bindings)
-			{
-				size_t size = bindings.size * bindings.dsize;
-				void* h_ptr = malloc(size);
-				memset(h_ptr, 0, size);
-				CHECK(cudaMemcpyAsync(
-					this->device_ptrs[0],
-					h_ptr,
-					size,
-					cudaMemcpyHostToDevice,
-					this->stream)
-				);
-				free(h_ptr);
-			}
-			this->infer();
-		}
-		printf("model warmup 10 times\n");
-
-	}
 }
 
 void YOLOv8::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
 {
-	const float inp_h = size.height;
+    const float inp_h  = size.height;
-	const float inp_w = size.width;
+    const float inp_w  = size.width;
-	float height = image.rows;
+    float       height = image.rows;
-	float width = image.cols;
+    float       width  = image.cols;
-
+
-	float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-	int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-	int padh = std::round(height * r);
+    int   padh = std::round(height * r);
-
+
-	cv::Mat tmp;
+    cv::Mat tmp;
-	if ((int)width != padw || (int)height != padh)
+    if ((int)width != padw || (int)height != padh) {
-	{
+        cv::resize(image, tmp, cv::Size(padw, padh));
-		cv::resize(
+    }
-			image,
+    else {
-			tmp,
+        tmp = image.clone();
-			cv::Size(padw, padh)
+    }
-		);
+
-	}
+    float dw = inp_w - padw;
-	else
+    float dh = inp_h - padh;
-	{
+
-		tmp = image.clone();
+    dw /= 2.0f;
-	}
+    dh /= 2.0f;
-
+    int top    = int(std::round(dh - 0.1f));
-	float dw = inp_w - padw;
+    int bottom = int(std::round(dh + 0.1f));
-	float dh = inp_h - padh;
+    int left   = int(std::round(dw - 0.1f));
-
+    int right  = int(std::round(dw + 0.1f));
-	dw /= 2.0f;
+
-	dh /= 2.0f;
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-	int top = int(std::round(dh - 0.1f));
+
-	int bottom = int(std::round(dh + 0.1f));
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-	int left = int(std::round(dw - 0.1f));
+    this->pparam.ratio  = 1 / r;
-	int right = int(std::round(dw + 0.1f));
+    this->pparam.dw     = dw;
-
+    this->pparam.dh     = dh;
-	cv::copyMakeBorder(
+    this->pparam.height = height;
-		tmp,
+    this->pparam.width  = width;
-		tmp,
+    ;
-		top,
-		bottom,
-		left,
-		right,
-		cv::BORDER_CONSTANT,
-		{ 114, 114, 114 }
-	);
-
-	cv::dnn::blobFromImage(tmp,
-		out,
-		1 / 255.f,
-		cv::Size(),
-		cv::Scalar(0, 0, 0),
-		true,
-		false,
-		CV_32F
-	);
-	this->pparam.ratio = 1 / r;
-	this->pparam.dw = dw;
-	this->pparam.dh = dh;
-	this->pparam.height = height;
-	this->pparam.width = width;;
 }
 
 void YOLOv8::copy_from_Mat(const cv::Mat& image)
 {
-	cv::Mat nchw;
+    cv::Mat  nchw;
-	auto& in_binding = this->input_bindings[0];
+    auto&    in_binding = this->input_bindings[0];
-	auto width = in_binding.dims.d[3];
+    auto     width      = in_binding.dims.d[3];
-	auto height = in_binding.dims.d[2];
+    auto     height     = in_binding.dims.d[2];
-	cv::Size size{ width, height };
+    cv::Size size{width, height};
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+
-		nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-		size
+
-	);
+    CHECK(cudaMemcpyAsync(
-
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{
-				4,
-				{ 1, 3, height, width }
-			}
-	);
-
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8::copy_from_Mat(const cv::Mat& image, cv::Size& size)
 {
-	cv::Mat nchw;
+    cv::Mat nchw;
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-		nchw,
+    CHECK(cudaMemcpyAsync(
-		size
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	);
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{ 4,
-			  { 1, 3, size.height, size.width }
-			}
-	);
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8::infer()
 {
 
-	this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-		this->device_ptrs.data(),
+    for (int i = 0; i < this->num_outputs; i++) {
-		this->stream,
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		nullptr
+        CHECK(cudaMemcpyAsync(
-	);
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-	for (int i = 0; i < this->num_outputs; i++)
+    }
-	{
+    cudaStreamSynchronize(this->stream);
-		size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		CHECK(cudaMemcpyAsync(this->host_ptrs[i],
-			this->device_ptrs[i + this->num_inputs],
-			osize,
-			cudaMemcpyDeviceToHost,
-			this->stream)
-		);
-
-	}
-	cudaStreamSynchronize(this->stream);
-
 }
 
 void YOLOv8::postprocess(std::vector<Object>& objs)
 {
-	objs.clear();
+    objs.clear();
-	int* num_dets = static_cast<int*>(this->host_ptrs[0]);
+    int*  num_dets = static_cast<int*>(this->host_ptrs[0]);
-	auto* boxes = static_cast<float*>(this->host_ptrs[1]);
+    auto* boxes    = static_cast<float*>(this->host_ptrs[1]);
-	auto* scores = static_cast<float*>(this->host_ptrs[2]);
+    auto* scores   = static_cast<float*>(this->host_ptrs[2]);
-	int* labels = static_cast<int*>(this->host_ptrs[3]);
+    int*  labels   = static_cast<int*>(this->host_ptrs[3]);
-	auto& dw = this->pparam.dw;
+    auto& dw       = this->pparam.dw;
-	auto& dh = this->pparam.dh;
+    auto& dh       = this->pparam.dh;
-	auto& width = this->pparam.width;
+    auto& width    = this->pparam.width;
-	auto& height = this->pparam.height;
+    auto& height   = this->pparam.height;
-	auto& ratio = this->pparam.ratio;
+    auto& ratio    = this->pparam.ratio;
-	for (int i = 0; i < num_dets[0]; i++)
+    for (int i = 0; i < num_dets[0]; i++) {
-	{
+        float* ptr = boxes + i * 4;
-		float* ptr = boxes + i * 4;
+
-
+        float x0 = *ptr++ - dw;
-		float x0 = *ptr++ - dw;
+        float y0 = *ptr++ - dh;
-		float y0 = *ptr++ - dh;
+        float x1 = *ptr++ - dw;
-		float x1 = *ptr++ - dw;
+        float y1 = *ptr - dh;
-		float y1 = *ptr - dh;
+
-
+        x0 = clamp(x0 * ratio, 0.f, width);
-		x0 = clamp(x0 * ratio, 0.f, width);
+        y0 = clamp(y0 * ratio, 0.f, height);
-		y0 = clamp(y0 * ratio, 0.f, height);
+        x1 = clamp(x1 * ratio, 0.f, width);
-		x1 = clamp(x1 * ratio, 0.f, width);
+        y1 = clamp(y1 * ratio, 0.f, height);
-		y1 = clamp(y1 * ratio, 0.f, height);
+        Object obj;
-		Object obj;
+        obj.rect.x      = x0;
-		obj.rect.x = x0;
+        obj.rect.y      = y0;
-		obj.rect.y = y0;
+        obj.rect.width  = x1 - x0;
-		obj.rect.width = x1 - x0;
+        obj.rect.height = y1 - y0;
-		obj.rect.height = y1 - y0;
+        obj.prob        = *(scores + i);
-		obj.prob = *(scores + i);
+        obj.label       = *(labels + i);
-		obj.label = *(labels + i);
+        objs.push_back(obj);
-		objs.push_back(obj);
+    }
-	}
 }
 
-void YOLOv8::draw_objects(
+void YOLOv8::draw_objects(const cv::Mat&                                image,
-	const cv::Mat& image,
+                          cv::Mat&                                      res,
-	cv::Mat& res,
+                          const std::vector<Object>&                    objs,
-	const std::vector<Object>& objs,
+                          const std::vector<std::string>&               CLASS_NAMES,
-	const std::vector<std::string>& CLASS_NAMES,
+                          const std::vector<std::vector<unsigned int>>& COLORS)
-	const std::vector<std::vector<unsigned int>>& COLORS
-)
 {
-	res = image.clone();
+    res = image.clone();
-	for (auto& obj : objs)
+    for (auto& obj : objs) {
-	{
+        cv::Scalar color = cv::Scalar(COLORS[obj.label][0], COLORS[obj.label][1], COLORS[obj.label][2]);
-		cv::Scalar color = cv::Scalar(
+        cv::rectangle(res, obj.rect, color, 2);
-			COLORS[obj.label][0],
+
-			COLORS[obj.label][1],
+        char text[256];
-			COLORS[obj.label][2]
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[obj.label].c_str(), obj.prob * 100);
-		);
+
-		cv::rectangle(
+        int      baseLine   = 0;
-			res,
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-			obj.rect,
+
-			color,
+        int x = (int)obj.rect.x;
-			2
+        int y = (int)obj.rect.y + 1;
-		);
+
-
+        if (y > res.rows)
-		char text[256];
+            y = res.rows;
-		sprintf(
+
-			text,
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-			"%s %.1f%%",
+
-			CLASS_NAMES[obj.label].c_str(),
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-			obj.prob * 100
+    }
-		);
-
-		int baseLine = 0;
-		cv::Size label_size = cv::getTextSize(
-			text,
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			1,
-			&baseLine
-		);
-
-		int x = (int)obj.rect.x;
-		int y = (int)obj.rect.y + 1;
-
-		if (y > res.rows)
-			y = res.rows;
-
-		cv::rectangle(
-			res,
-			cv::Rect(x, y, label_size.width, label_size.height + baseLine),
-			{ 0, 0, 255 },
-			-1
-		);
-
-		cv::putText(
-			res,
-			text,
-			cv::Point(x, y + label_size.height),
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			{ 255, 255, 255 },
-			1
-		);
-	}
 }
-#endif //DETECT_END2END_YOLOV8_HPP
+#endif  // DETECT_END2END_YOLOV8_HPP

csrc/detect/end2end/main.cpp

@@ -2,160 +2,119 @@
 // Created by ubuntu on 1/20/23.
 //
 #include "chrono"
-#include "yolov8.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8.hpp"
 
 const std::vector<std::string> CLASS_NAMES = {
-	"person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",
-	"train", "truck", "boat", "traffic light", "fire hydrant",
+    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",    "parking meter", "bench",
-	"stop sign", "parking meter", "bench", "bird", "cat",
+    "bird",           "cat",        "dog",           "horse",         "sheep",        "cow",           "elephant",
-	"dog", "horse", "sheep", "cow", "elephant",
+    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",     "handbag",       "tie",
-	"bear", "zebra", "giraffe", "backpack", "umbrella",
+    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball",  "kite",          "baseball bat",
-	"handbag", "tie", "suitcase", "frisbee", "skis",
+    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",       "wine glass",    "cup",
-	"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "fork",           "knife",      "spoon",         "bowl",          "banana",       "apple",         "sandwich",
-	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
+    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",        "donut",         "cake",
-	"cup", "fork", "knife", "spoon", "bowl",
+    "chair",          "couch",      "potted plant",  "bed",           "dining table", "toilet",        "tv",
-	"banana", "apple", "sandwich", "orange", "broccoli",
+    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",   "microwave",     "oven",
-	"carrot", "hot dog", "pizza", "donut", "cake",
+    "toaster",        "sink",       "refrigerator",  "book",          "clock",        "vase",          "scissors",
-	"chair", "couch", "potted plant", "bed", "dining table",
+    "teddy bear",     "hair drier", "toothbrush"};
-	"toilet", "tv", "laptop", "mouse", "remote",
-	"keyboard", "cell phone", "microwave", "oven",
-	"toaster", "sink", "refrigerator", "book", "clock", "vase",
-	"scissors", "teddy bear", "hair drier", "toothbrush" };
 
 const std::vector<std::vector<unsigned int>> COLORS = {
-	{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
+    {0, 114, 189},   {217, 83, 25},   {237, 177, 32},  {126, 47, 142},  {119, 172, 48},  {77, 190, 238},
-	{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
+    {162, 20, 47},   {76, 76, 76},    {153, 153, 153}, {255, 0, 0},     {255, 128, 0},   {191, 191, 0},
-	{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
+    {0, 255, 0},     {0, 0, 255},     {170, 0, 255},   {85, 85, 0},     {85, 170, 0},    {85, 255, 0},
-	{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
+    {170, 85, 0},    {170, 170, 0},   {170, 255, 0},   {255, 85, 0},    {255, 170, 0},   {255, 255, 0},
-	{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
+    {0, 85, 128},    {0, 170, 128},   {0, 255, 128},   {85, 0, 128},    {85, 85, 128},   {85, 170, 128},
-	{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
+    {85, 255, 128},  {170, 0, 128},   {170, 85, 128},  {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
-	{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
+    {255, 85, 128},  {255, 170, 128}, {255, 255, 128}, {0, 85, 255},    {0, 170, 255},   {0, 255, 255},
-	{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
+    {85, 0, 255},    {85, 85, 255},   {85, 170, 255},  {85, 255, 255},  {170, 0, 255},   {170, 85, 255},
-	{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
+    {170, 170, 255}, {170, 255, 255}, {255, 0, 255},   {255, 85, 255},  {255, 170, 255}, {85, 0, 0},
-	{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
+    {128, 0, 0},     {170, 0, 0},     {212, 0, 0},     {255, 0, 0},     {0, 43, 0},      {0, 85, 0},
-	{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
+    {0, 128, 0},     {0, 170, 0},     {0, 212, 0},     {0, 255, 0},     {0, 0, 43},      {0, 0, 85},
-	{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
+    {0, 0, 128},     {0, 0, 170},     {0, 0, 212},     {0, 0, 255},     {0, 0, 0},       {36, 36, 36},
-	{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
+    {73, 73, 73},    {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
-	{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
+    {80, 183, 189},  {128, 128, 0}};
-	{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
-	{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
-	{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
-	{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
-	{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
-	{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
-	{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
-	{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
-	{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
-	{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
-	{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
-	{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
-	{ 80, 183, 189 }, { 128, 128, 0 }
-};
 
 int main(int argc, char** argv)
 {
-	// cuda:0
+    // cuda:0
-	cudaSetDevice(0);
+    cudaSetDevice(0);
 
-	const std::string engine_file_path{ argv[1] };
+    const std::string engine_file_path{argv[1]};
-	const std::string path{ argv[2] };
+    const std::string path{argv[2]};
 
-	std::vector<std::string> imagePathList;
+    std::vector<std::string> imagePathList;
-	bool isVideo{ false };
+    bool                     isVideo{false};
 
-	assert(argc == 3);
+    assert(argc == 3);
 
-	auto yolov8 = new YOLOv8(engine_file_path);
+    auto yolov8 = new YOLOv8(engine_file_path);
-	yolov8->make_pipe(true);
+    yolov8->make_pipe(true);
 
-	if (IsFile(path))
+    if (IsFile(path)) {
-	{
+        std::string suffix = path.substr(path.find_last_of('.') + 1);
-		std::string suffix = path.substr(path.find_last_of('.') + 1);
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-		if (
+            imagePathList.push_back(path);
-			suffix == "jpg" ||
+        }
-				suffix == "jpeg" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-				suffix == "png"
+                 || suffix == "mkv") {
-			)
+            isVideo = true;
-		{
+        }
-			imagePathList.push_back(path);
+        else {
-		}
+            printf("suffix %s is wrong !!!\n", suffix.c_str());
-		else if (
+            std::abort();
-			suffix == "mp4" ||
+        }
-				suffix == "avi" ||
+    }
-				suffix == "m4v" ||
+    else if (IsFolder(path)) {
-				suffix == "mpeg" ||
+        cv::glob(path + "/*.jpg", imagePathList);
-				suffix == "mov" ||
+    }
-				suffix == "mkv"
-			)
-		{
-			isVideo = true;
-		}
-		else
-		{
-			printf("suffix %s is wrong !!!\n", suffix.c_str());
-			std::abort();
-		}
-	}
-	else if (IsFolder(path))
-	{
-		cv::glob(path + "/*.jpg", imagePathList);
-	}
 
-	cv::Mat res, image;
+    cv::Mat             res, image;
-	cv::Size size = cv::Size{ 640, 640 };
+    cv::Size            size = cv::Size{640, 640};
-	std::vector<Object> objs;
+    std::vector<Object> objs;
 
-	cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
+    cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
 
-	if (isVideo)
+    if (isVideo) {
-	{
+        cv::VideoCapture cap(path);
-		cv::VideoCapture cap(path);
 
-		if (!cap.isOpened())
+        if (!cap.isOpened()) {
-		{
+            printf("can not open %s\n", path.c_str());
-			printf("can not open %s\n", path.c_str());
+            return -1;
-			return -1;
+        }
-		}
+        while (cap.read(image)) {
-		while (cap.read(image))
+            objs.clear();
-		{
+            yolov8->copy_from_Mat(image, size);
-			objs.clear();
+            auto start = std::chrono::system_clock::now();
-			yolov8->copy_from_Mat(image, size);
+            yolov8->infer();
-			auto start = std::chrono::system_clock::now();
+            auto end = std::chrono::system_clock::now();
-			yolov8->infer();
+            yolov8->postprocess(objs);
-			auto end = std::chrono::system_clock::now();
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
-			yolov8->postprocess(objs);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
+            printf("cost %2.4lf ms\n", tc);
-			auto tc = (double)
+            cv::imshow("result", res);
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+            if (cv::waitKey(10) == 'q') {
-			printf("cost %2.4lf ms\n", tc);
+                break;
-			cv::imshow("result", res);
+            }
-			if (cv::waitKey(10) == 'q')
+        }
-			{
+    }
-				break;
+    else {
-			}
+        for (auto& path : imagePathList) {
-		}
+            objs.clear();
-	}
+            image = cv::imread(path);
-	else
+            yolov8->copy_from_Mat(image, size);
-	{
+            auto start = std::chrono::system_clock::now();
-		for (auto& path : imagePathList)
+            yolov8->infer();
-		{
+            auto end = std::chrono::system_clock::now();
-			objs.clear();
+            yolov8->postprocess(objs);
-			image = cv::imread(path);
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
-			yolov8->copy_from_Mat(image, size);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			auto start = std::chrono::system_clock::now();
+            printf("cost %2.4lf ms\n", tc);
-			yolov8->infer();
+            cv::imshow("result", res);
-			auto end = std::chrono::system_clock::now();
+            cv::waitKey(0);
-			yolov8->postprocess(objs);
+        }
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
+    }
-			auto tc = (double)
+    cv::destroyAllWindows();
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+    delete yolov8;
-			printf("cost %2.4lf ms\n", tc);
+    return 0;
-			cv::imshow("result", res);
-			cv::waitKey(0);
-		}
-	}
-	cv::destroyAllWindows();
-	delete yolov8;
-	return 0;
 }

csrc/detect/normal/include/common.hpp

@@ -4,153 +4,138 @@
 
 #ifndef DETECT_NORMAL_COMMON_HPP
 #define DETECT_NORMAL_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace det
+namespace det {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float> rect;
-		cv::Rect_<float> rect;
+    int              label = 0;
-		int label = 0;
+    float            prob  = 0.0;
-		float prob = 0.0;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace det
-}
+#endif  // DETECT_NORMAL_COMMON_HPP
-#endif //DETECT_NORMAL_COMMON_HPP

csrc/detect/normal/include/yolov8.hpp

@@ -3,495 +3,319 @@
 //
 #ifndef DETECT_NORMAL_YOLOV8_HPP
 #define DETECT_NORMAL_YOLOV8_HPP
-#include "fstream"
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include "fstream"
 using namespace det;
 
-class YOLOv8
+class YOLOv8 {
-{
 public:
-	explicit YOLOv8(const std::string& engine_file_path);
+    explicit YOLOv8(const std::string& engine_file_path);
-	~YOLOv8();
+    ~YOLOv8();
-
+
-	void make_pipe(bool warmup = true);
+    void                 make_pipe(bool warmup = true);
-	void copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image);
-	void copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
-	void letterbox(
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-		const cv::Mat& image,
+    void                 infer();
-		cv::Mat& out,
+    void                 postprocess(std::vector<Object>& objs,
-		cv::Size& size
+                                     float                score_thres = 0.25f,
-	);
+                                     float                iou_thres   = 0.65f,
-	void infer();
+                                     int                  topk        = 100,
-	void postprocess(
+                                     int                  num_labels  = 80);
-		std::vector<Object>& objs,
+    static void          draw_objects(const cv::Mat&                                image,
-		float score_thres = 0.25f,
+                                      cv::Mat&                                      res,
-		float iou_thres = 0.65f,
+                                      const std::vector<Object>&                    objs,
-		int topk = 100,
+                                      const std::vector<std::string>&               CLASS_NAMES,
-		int num_labels = 80
+                                      const std::vector<std::vector<unsigned int>>& COLORS);
-	);
+    int                  num_bindings;
-	static void draw_objects(
+    int                  num_inputs  = 0;
-		const cv::Mat& image,
+    int                  num_outputs = 0;
-		cv::Mat& res,
+    std::vector<Binding> input_bindings;
-		const std::vector<Object>& objs,
+    std::vector<Binding> output_bindings;
-		const std::vector<std::string>& CLASS_NAMES,
+    std::vector<void*>   host_ptrs;
-		const std::vector<std::vector<unsigned int>>& COLORS
+    std::vector<void*>   device_ptrs;
-	);
+
-	int num_bindings;
+    PreParam pparam;
-	int num_inputs = 0;
-	int num_outputs = 0;
-	std::vector<Binding> input_bindings;
-	std::vector<Binding> output_bindings;
-	std::vector<void*> host_ptrs;
-	std::vector<void*> device_ptrs;
-
-	PreParam pparam;
-private:
-	nvinfer1::ICudaEngine* engine = nullptr;
-	nvinfer1::IRuntime* runtime = nullptr;
-	nvinfer1::IExecutionContext* context = nullptr;
-	cudaStream_t stream = nullptr;
-	Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
 YOLOv8::YOLOv8(const std::string& engine_file_path)
 {
-	std::ifstream file(engine_file_path, std::ios::binary);
+    std::ifstream file(engine_file_path, std::ios::binary);
-	assert(file.good());
+    assert(file.good());
-	file.seekg(0, std::ios::end);
+    file.seekg(0, std::ios::end);
-	auto size = file.tellg();
+    auto size = file.tellg();
-	file.seekg(0, std::ios::beg);
+    file.seekg(0, std::ios::beg);
-	char* trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
-	assert(trtModelStream);
+    assert(trtModelStream);
-	file.read(trtModelStream, size);
+    file.read(trtModelStream, size);
-	file.close();
+    file.close();
-	initLibNvInferPlugins(&this->gLogger, "");
+    initLibNvInferPlugins(&this->gLogger, "");
-	this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
-	assert(this->runtime != nullptr);
+    assert(this->runtime != nullptr);
-
+
-	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
-	assert(this->engine != nullptr);
+    assert(this->engine != nullptr);
-	delete[] trtModelStream;
+    delete[] trtModelStream;
-	this->context = this->engine->createExecutionContext();
+    this->context = this->engine->createExecutionContext();
-
+
-	assert(this->context != nullptr);
+    assert(this->context != nullptr);
-	cudaStreamCreate(&this->stream);
+    cudaStreamCreate(&this->stream);
-	this->num_bindings = this->engine->getNbBindings();
+    this->num_bindings = this->engine->getNbBindings();
-
+
-	for (int i = 0; i < this->num_bindings; ++i)
+    for (int i = 0; i < this->num_bindings; ++i) {
-	{
+        Binding            binding;
-		Binding binding;
+        nvinfer1::Dims     dims;
-		nvinfer1::Dims dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-		nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
-		std::string name = this->engine->getBindingName(i);
+        binding.name             = name;
-		binding.name = name;
+        binding.dsize            = type_to_size(dtype);
-		binding.dsize = type_to_size(dtype);
+
-
+        bool IsInput = engine->bindingIsInput(i);
-		bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
-		if (IsInput)
+            this->num_inputs += 1;
-		{
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-			this->num_inputs += 1;
+            binding.size = get_size_by_dims(dims);
-			dims = this->engine->getProfileDimensions(
+            binding.dims = dims;
-				i,
+            this->input_bindings.push_back(binding);
-				0,
+            // set max opt shape
-				nvinfer1::OptProfileSelector::kMAX);
+            this->context->setBindingDimensions(i, dims);
-			binding.size = get_size_by_dims(dims);
+        }
-			binding.dims = dims;
+        else {
-			this->input_bindings.push_back(binding);
+            dims         = this->context->getBindingDimensions(i);
-			// set max opt shape
+            binding.size = get_size_by_dims(dims);
-			this->context->setBindingDimensions(i, dims);
+            binding.dims = dims;
-
+            this->output_bindings.push_back(binding);
-		}
+            this->num_outputs += 1;
-		else
+        }
-		{
+    }
-			dims = this->context->getBindingDimensions(i);
-			binding.size = get_size_by_dims(dims);
-			binding.dims = dims;
-			this->output_bindings.push_back(binding);
-			this->num_outputs += 1;
-		}
-	}
-
 }
 
 YOLOv8::~YOLOv8()
 {
-	this->context->destroy();
+    this->context->destroy();
-	this->engine->destroy();
+    this->engine->destroy();
-	this->runtime->destroy();
+    this->runtime->destroy();
-	cudaStreamDestroy(this->stream);
+    cudaStreamDestroy(this->stream);
-	for (auto& ptr : this->device_ptrs)
+    for (auto& ptr : this->device_ptrs) {
-	{
+        CHECK(cudaFree(ptr));
-		CHECK(cudaFree(ptr));
+    }
-	}
+
-
+    for (auto& ptr : this->host_ptrs) {
-	for (auto& ptr : this->host_ptrs)
+        CHECK(cudaFreeHost(ptr));
-	{
+    }
-		CHECK(cudaFreeHost(ptr));
-	}
-
 }
 void YOLOv8::make_pipe(bool warmup)
 {
 
-	for (auto& bindings : this->input_bindings)
+    for (auto& bindings : this->input_bindings) {
-	{
+        void* d_ptr;
-		void* d_ptr;
+        CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream));
-		CHECK(cudaMallocAsync(
+        this->device_ptrs.push_back(d_ptr);
-			&d_ptr,
+    }
-			bindings.size * bindings.dsize,
+
-			this->stream)
+    for (auto& bindings : this->output_bindings) {
-		);
+        void * d_ptr, *h_ptr;
-		this->device_ptrs.push_back(d_ptr);
+        size_t size = bindings.size * bindings.dsize;
-	}
+        CHECK(cudaMallocAsync(&d_ptr, size, this->stream));
-
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-	for (auto& bindings : this->output_bindings)
+        this->device_ptrs.push_back(d_ptr);
-	{
+        this->host_ptrs.push_back(h_ptr);
-		void* d_ptr, * h_ptr;
+    }
-		size_t size = bindings.size * bindings.dsize;
+
-		CHECK(cudaMallocAsync(
+    if (warmup) {
-			&d_ptr,
+        for (int i = 0; i < 10; i++) {
-			size,
+            for (auto& bindings : this->input_bindings) {
-			this->stream)
+                size_t size  = bindings.size * bindings.dsize;
-		);
+                void*  h_ptr = malloc(size);
-		CHECK(cudaHostAlloc(
+                memset(h_ptr, 0, size);
-			&h_ptr,
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-			size,
+                free(h_ptr);
-			0)
+            }
-		);
+            this->infer();
-		this->device_ptrs.push_back(d_ptr);
+        }
-		this->host_ptrs.push_back(h_ptr);
+        printf("model warmup 10 times\n");
-	}
+    }
-
-	if (warmup)
-	{
-		for (int i = 0; i < 10; i++)
-		{
-			for (auto& bindings : this->input_bindings)
-			{
-				size_t size = bindings.size * bindings.dsize;
-				void* h_ptr = malloc(size);
-				memset(h_ptr, 0, size);
-				CHECK(cudaMemcpyAsync(
-					this->device_ptrs[0],
-					h_ptr,
-					size,
-					cudaMemcpyHostToDevice,
-					this->stream)
-				);
-				free(h_ptr);
-			}
-			this->infer();
-		}
-		printf("model warmup 10 times\n");
-
-	}
 }
 
-void YOLOv8::letterbox(
+void YOLOv8::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
-	const cv::Mat& image,
-	cv::Mat& out,
-	cv::Size& size
-)
 {
-	const float inp_h = size.height;
+    const float inp_h  = size.height;
-	const float inp_w = size.width;
+    const float inp_w  = size.width;
-	float height = image.rows;
+    float       height = image.rows;
-	float width = image.cols;
+    float       width  = image.cols;
-
+
-	float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-	int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-	int padh = std::round(height * r);
+    int   padh = std::round(height * r);
-
+
-	cv::Mat tmp;
+    cv::Mat tmp;
-	if ((int)width != padw || (int)height != padh)
+    if ((int)width != padw || (int)height != padh) {
-	{
+        cv::resize(image, tmp, cv::Size(padw, padh));
-		cv::resize(
+    }
-			image,
+    else {
-			tmp,
+        tmp = image.clone();
-			cv::Size(padw, padh)
+    }
-		);
+
-	}
+    float dw = inp_w - padw;
-	else
+    float dh = inp_h - padh;
-	{
+
-		tmp = image.clone();
+    dw /= 2.0f;
-	}
+    dh /= 2.0f;
-
+    int top    = int(std::round(dh - 0.1f));
-	float dw = inp_w - padw;
+    int bottom = int(std::round(dh + 0.1f));
-	float dh = inp_h - padh;
+    int left   = int(std::round(dw - 0.1f));
-
+    int right  = int(std::round(dw + 0.1f));
-	dw /= 2.0f;
+
-	dh /= 2.0f;
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-	int top = int(std::round(dh - 0.1f));
+
-	int bottom = int(std::round(dh + 0.1f));
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-	int left = int(std::round(dw - 0.1f));
+    this->pparam.ratio  = 1 / r;
-	int right = int(std::round(dw + 0.1f));
+    this->pparam.dw     = dw;
-
+    this->pparam.dh     = dh;
-	cv::copyMakeBorder(
+    this->pparam.height = height;
-		tmp,
+    this->pparam.width  = width;
-		tmp,
+    ;
-		top,
-		bottom,
-		left,
-		right,
-		cv::BORDER_CONSTANT,
-		{ 114, 114, 114 }
-	);
-
-	cv::dnn::blobFromImage(tmp,
-		out,
-		1 / 255.f,
-		cv::Size(),
-		cv::Scalar(0, 0, 0),
-		true,
-		false,
-		CV_32F
-	);
-	this->pparam.ratio = 1 / r;
-	this->pparam.dw = dw;
-	this->pparam.dh = dh;
-	this->pparam.height = height;
-	this->pparam.width = width;;
 }
 
 void YOLOv8::copy_from_Mat(const cv::Mat& image)
 {
-	cv::Mat nchw;
+    cv::Mat  nchw;
-	auto& in_binding = this->input_bindings[0];
+    auto&    in_binding = this->input_bindings[0];
-	auto width = in_binding.dims.d[3];
+    auto     width      = in_binding.dims.d[3];
-	auto height = in_binding.dims.d[2];
+    auto     height     = in_binding.dims.d[2];
-	cv::Size size{ width, height };
+    cv::Size size{width, height};
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+
-		nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-		size
+
-	);
+    CHECK(cudaMemcpyAsync(
-
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{
-				4,
-				{ 1, 3, height, width }
-			}
-	);
-
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8::copy_from_Mat(const cv::Mat& image, cv::Size& size)
 {
-	cv::Mat nchw;
+    cv::Mat nchw;
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-		nchw,
+    CHECK(cudaMemcpyAsync(
-		size
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	);
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{ 4,
-			  { 1, 3, size.height, size.width }
-			}
-	);
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8::infer()
 {
 
-	this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-		this->device_ptrs.data(),
+    for (int i = 0; i < this->num_outputs; i++) {
-		this->stream,
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		nullptr
+        CHECK(cudaMemcpyAsync(
-	);
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-	for (int i = 0; i < this->num_outputs; i++)
+    }
-	{
+    cudaStreamSynchronize(this->stream);
-		size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		CHECK(cudaMemcpyAsync(this->host_ptrs[i],
-			this->device_ptrs[i + this->num_inputs],
-			osize,
-			cudaMemcpyDeviceToHost,
-			this->stream)
-		);
-
-	}
-	cudaStreamSynchronize(this->stream);
-
 }
 
-void YOLOv8::postprocess(
+void YOLOv8::postprocess(std::vector<Object>& objs, float score_thres, float iou_thres, int topk, int num_labels)
-	std::vector<Object>& objs,
-	float score_thres,
-	float iou_thres,
-	int topk,
-	int num_labels
-)
 {
-	objs.clear();
+    objs.clear();
-	auto num_channels = this->output_bindings[0].dims.d[1];
+    auto num_channels = this->output_bindings[0].dims.d[1];
-	auto num_anchors = this->output_bindings[0].dims.d[2];
+    auto num_anchors  = this->output_bindings[0].dims.d[2];
-
+
-	auto& dw = this->pparam.dw;
+    auto& dw     = this->pparam.dw;
-	auto& dh = this->pparam.dh;
+    auto& dh     = this->pparam.dh;
-	auto& width = this->pparam.width;
+    auto& width  = this->pparam.width;
-	auto& height = this->pparam.height;
+    auto& height = this->pparam.height;
-	auto& ratio = this->pparam.ratio;
+    auto& ratio  = this->pparam.ratio;
-
+
-	std::vector<cv::Rect> bboxes;
+    std::vector<cv::Rect> bboxes;
-	std::vector<float> scores;
+    std::vector<float>    scores;
-	std::vector<int> labels;
+    std::vector<int>      labels;
-	std::vector<int> indices;
+    std::vector<int>      indices;
-
+
-	cv::Mat output = cv::Mat(
+    cv::Mat output = cv::Mat(num_channels, num_anchors, CV_32F, static_cast<float*>(this->host_ptrs[0]));
-		num_channels,
+    output         = output.t();
-		num_anchors,
+    for (int i = 0; i < num_anchors; i++) {
-		CV_32F,
+        auto  row_ptr    = output.row(i).ptr<float>();
-		static_cast<float*>(this->host_ptrs[0])
+        auto  bboxes_ptr = row_ptr;
-	);
+        auto  scores_ptr = row_ptr + 4;
-	output = output.t();
+        auto  max_s_ptr  = std::max_element(scores_ptr, scores_ptr + num_labels);
-	for (int i = 0; i < num_anchors; i++)
+        float score      = *max_s_ptr;
-	{
+        if (score > score_thres) {
-		auto row_ptr = output.row(i).ptr<float>();
+            float x = *bboxes_ptr++ - dw;
-		auto bboxes_ptr = row_ptr;
+            float y = *bboxes_ptr++ - dh;
-		auto scores_ptr = row_ptr + 4;
+            float w = *bboxes_ptr++;
-		auto max_s_ptr = std::max_element(scores_ptr, scores_ptr + num_labels);
+            float h = *bboxes_ptr;
-		float score = *max_s_ptr;
+
-		if (score > score_thres)
+            float x0 = clamp((x - 0.5f * w) * ratio, 0.f, width);
-		{
+            float y0 = clamp((y - 0.5f * h) * ratio, 0.f, height);
-			float x = *bboxes_ptr++ - dw;
+            float x1 = clamp((x + 0.5f * w) * ratio, 0.f, width);
-			float y = *bboxes_ptr++ - dh;
+            float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);
-			float w = *bboxes_ptr++;
+
-			float h = *bboxes_ptr;
+            int              label = max_s_ptr - scores_ptr;
-
+            cv::Rect_<float> bbox;
-			float x0 = clamp((x - 0.5f * w) * ratio, 0.f, width);
+            bbox.x      = x0;
-			float y0 = clamp((y - 0.5f * h) * ratio, 0.f, height);
+            bbox.y      = y0;
-			float x1 = clamp((x + 0.5f * w) * ratio, 0.f, width);
+            bbox.width  = x1 - x0;
-			float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);
+            bbox.height = y1 - y0;
-
+
-			int label = max_s_ptr - scores_ptr;
+            bboxes.push_back(bbox);
-			cv::Rect_<float> bbox;
+            labels.push_back(label);
-			bbox.x = x0;
+            scores.push_back(score);
-			bbox.y = y0;
+        }
-			bbox.width = x1 - x0;
+    }
-			bbox.height = y1 - y0;
-
-			bboxes.push_back(bbox);
-			labels.push_back(label);
-			scores.push_back(score);
-		}
-	}
 
 #ifdef BATCHED_NMS
-	cv::dnn::NMSBoxesBatched(
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		labels,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #else
-	cv::dnn::NMSBoxes(
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #endif
 
-	int cnt = 0;
+    int cnt = 0;
-	for (auto& i : indices)
+    for (auto& i : indices) {
-	{
+        if (cnt >= topk) {
-		if (cnt >= topk)
+            break;
-		{
+        }
-			break;
+        Object obj;
-		}
+        obj.rect  = bboxes[i];
-		Object obj;
+        obj.prob  = scores[i];
-		obj.rect = bboxes[i];
+        obj.label = labels[i];
-		obj.prob = scores[i];
+        objs.push_back(obj);
-		obj.label = labels[i];
+        cnt += 1;
-		objs.push_back(obj);
+    }
-		cnt += 1;
-	}
 }
 
-void YOLOv8::draw_objects(
+void YOLOv8::draw_objects(const cv::Mat&                                image,
-	const cv::Mat& image,
+                          cv::Mat&                                      res,
-	cv::Mat& res,
+                          const std::vector<Object>&                    objs,
-	const std::vector<Object>& objs,
+                          const std::vector<std::string>&               CLASS_NAMES,
-	const std::vector<std::string>& CLASS_NAMES,
+                          const std::vector<std::vector<unsigned int>>& COLORS)
-	const std::vector<std::vector<unsigned int>>& COLORS
-)
 {
-	res = image.clone();
+    res = image.clone();
-	for (auto& obj : objs)
+    for (auto& obj : objs) {
-	{
+        cv::Scalar color = cv::Scalar(COLORS[obj.label][0], COLORS[obj.label][1], COLORS[obj.label][2]);
-		cv::Scalar color = cv::Scalar(
+        cv::rectangle(res, obj.rect, color, 2);
-			COLORS[obj.label][0],
+
-			COLORS[obj.label][1],
+        char text[256];
-			COLORS[obj.label][2]
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[obj.label].c_str(), obj.prob * 100);
-		);
+
-		cv::rectangle(
+        int      baseLine   = 0;
-			res,
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-			obj.rect,
+
-			color,
+        int x = (int)obj.rect.x;
-			2
+        int y = (int)obj.rect.y + 1;
-		);
+
-
+        if (y > res.rows)
-		char text[256];
+            y = res.rows;
-		sprintf(
+
-			text,
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-			"%s %.1f%%",
+
-			CLASS_NAMES[obj.label].c_str(),
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-			obj.prob * 100
+    }
-		);
-
-		int baseLine = 0;
-		cv::Size label_size = cv::getTextSize(
-			text,
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			1,
-			&baseLine
-		);
-
-		int x = (int)obj.rect.x;
-		int y = (int)obj.rect.y + 1;
-
-		if (y > res.rows)
-			y = res.rows;
-
-		cv::rectangle(
-			res,
-			cv::Rect(x, y, label_size.width, label_size.height + baseLine),
-			{ 0, 0, 255 },
-			-1
-		);
-
-		cv::putText(
-			res,
-			text,
-			cv::Point(x, y + label_size.height),
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			{ 255, 255, 255 },
-			1
-		);
-	}
 }
-#endif //DETECT_NORMAL_YOLOV8_HPP
+#endif  // DETECT_NORMAL_YOLOV8_HPP

csrc/detect/normal/main.cpp

@@ -2,165 +2,124 @@
 // Created by ubuntu on 1/20/23.
 //
 #include "chrono"
-#include "yolov8.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8.hpp"
 
 const std::vector<std::string> CLASS_NAMES = {
-	"person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",
-	"train", "truck", "boat", "traffic light", "fire hydrant",
+    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",    "parking meter", "bench",
-	"stop sign", "parking meter", "bench", "bird", "cat",
+    "bird",           "cat",        "dog",           "horse",         "sheep",        "cow",           "elephant",
-	"dog", "horse", "sheep", "cow", "elephant",
+    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",     "handbag",       "tie",
-	"bear", "zebra", "giraffe", "backpack", "umbrella",
+    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball",  "kite",          "baseball bat",
-	"handbag", "tie", "suitcase", "frisbee", "skis",
+    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",       "wine glass",    "cup",
-	"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "fork",           "knife",      "spoon",         "bowl",          "banana",       "apple",         "sandwich",
-	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
+    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",        "donut",         "cake",
-	"cup", "fork", "knife", "spoon", "bowl",
+    "chair",          "couch",      "potted plant",  "bed",           "dining table", "toilet",        "tv",
-	"banana", "apple", "sandwich", "orange", "broccoli",
+    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",   "microwave",     "oven",
-	"carrot", "hot dog", "pizza", "donut", "cake",
+    "toaster",        "sink",       "refrigerator",  "book",          "clock",        "vase",          "scissors",
-	"chair", "couch", "potted plant", "bed", "dining table",
+    "teddy bear",     "hair drier", "toothbrush"};
-	"toilet", "tv", "laptop", "mouse", "remote",
-	"keyboard", "cell phone", "microwave", "oven",
-	"toaster", "sink", "refrigerator", "book", "clock", "vase",
-	"scissors", "teddy bear", "hair drier", "toothbrush" };
 
 const std::vector<std::vector<unsigned int>> COLORS = {
-	{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
+    {0, 114, 189},   {217, 83, 25},   {237, 177, 32},  {126, 47, 142},  {119, 172, 48},  {77, 190, 238},
-	{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
+    {162, 20, 47},   {76, 76, 76},    {153, 153, 153}, {255, 0, 0},     {255, 128, 0},   {191, 191, 0},
-	{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
+    {0, 255, 0},     {0, 0, 255},     {170, 0, 255},   {85, 85, 0},     {85, 170, 0},    {85, 255, 0},
-	{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
+    {170, 85, 0},    {170, 170, 0},   {170, 255, 0},   {255, 85, 0},    {255, 170, 0},   {255, 255, 0},
-	{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
+    {0, 85, 128},    {0, 170, 128},   {0, 255, 128},   {85, 0, 128},    {85, 85, 128},   {85, 170, 128},
-	{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
+    {85, 255, 128},  {170, 0, 128},   {170, 85, 128},  {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
-	{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
+    {255, 85, 128},  {255, 170, 128}, {255, 255, 128}, {0, 85, 255},    {0, 170, 255},   {0, 255, 255},
-	{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
+    {85, 0, 255},    {85, 85, 255},   {85, 170, 255},  {85, 255, 255},  {170, 0, 255},   {170, 85, 255},
-	{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
+    {170, 170, 255}, {170, 255, 255}, {255, 0, 255},   {255, 85, 255},  {255, 170, 255}, {85, 0, 0},
-	{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
+    {128, 0, 0},     {170, 0, 0},     {212, 0, 0},     {255, 0, 0},     {0, 43, 0},      {0, 85, 0},
-	{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
+    {0, 128, 0},     {0, 170, 0},     {0, 212, 0},     {0, 255, 0},     {0, 0, 43},      {0, 0, 85},
-	{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
+    {0, 0, 128},     {0, 0, 170},     {0, 0, 212},     {0, 0, 255},     {0, 0, 0},       {36, 36, 36},
-	{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
+    {73, 73, 73},    {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
-	{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
+    {80, 183, 189},  {128, 128, 0}};
-	{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
-	{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
-	{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
-	{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
-	{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
-	{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
-	{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
-	{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
-	{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
-	{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
-	{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
-	{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
-	{ 80, 183, 189 }, { 128, 128, 0 }
-};
 
 int main(int argc, char** argv)
 {
-	// cuda:0
+    // cuda:0
-	cudaSetDevice(0);
+    cudaSetDevice(0);
 
-	const std::string engine_file_path{ argv[1] };
+    const std::string engine_file_path{argv[1]};
-	const std::string path{ argv[2] };
+    const std::string path{argv[2]};
 
-	std::vector<std::string> imagePathList;
+    std::vector<std::string> imagePathList;
-	bool isVideo{ false };
+    bool                     isVideo{false};
 
-	assert(argc == 3);
+    assert(argc == 3);
 
-	auto yolov8 = new YOLOv8(engine_file_path);
+    auto yolov8 = new YOLOv8(engine_file_path);
-	yolov8->make_pipe(true);
+    yolov8->make_pipe(true);
 
-	if (IsFile(path))
+    if (IsFile(path)) {
-	{
+        std::string suffix = path.substr(path.find_last_of('.') + 1);
-		std::string suffix = path.substr(path.find_last_of('.') + 1);
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-		if (
+            imagePathList.push_back(path);
-			suffix == "jpg" ||
+        }
-				suffix == "jpeg" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-				suffix == "png"
+                 || suffix == "mkv") {
-			)
+            isVideo = true;
-		{
+        }
-			imagePathList.push_back(path);
+        else {
-		}
+            printf("suffix %s is wrong !!!\n", suffix.c_str());
-		else if (
+            std::abort();
-			suffix == "mp4" ||
+        }
-				suffix == "avi" ||
+    }
-				suffix == "m4v" ||
+    else if (IsFolder(path)) {
-				suffix == "mpeg" ||
+        cv::glob(path + "/*.jpg", imagePathList);
-				suffix == "mov" ||
+    }
-				suffix == "mkv"
-			)
-		{
-			isVideo = true;
-		}
-		else
-		{
-			printf("suffix %s is wrong !!!\n", suffix.c_str());
-			std::abort();
-		}
-	}
-	else if (IsFolder(path))
-	{
-		cv::glob(path + "/*.jpg", imagePathList);
-	}
 
-	cv::Mat res, image;
+    cv::Mat  res, image;
-	cv::Size size = cv::Size{ 640, 640 };
+    cv::Size size        = cv::Size{640, 640};
-	int num_labels = 80;
+    int      num_labels  = 80;
-	int topk = 100;
+    int      topk        = 100;
-	float score_thres = 0.25f;
+    float    score_thres = 0.25f;
-	float iou_thres = 0.65f;
+    float    iou_thres   = 0.65f;
 
-	std::vector<Object> objs;
+    std::vector<Object> objs;
 
-	cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
+    cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
 
-	if (isVideo)
+    if (isVideo) {
-	{
+        cv::VideoCapture cap(path);
-		cv::VideoCapture cap(path);
 
-		if (!cap.isOpened())
+        if (!cap.isOpened()) {
-		{
+            printf("can not open %s\n", path.c_str());
-			printf("can not open %s\n", path.c_str());
+            return -1;
-			return -1;
+        }
-		}
+        while (cap.read(image)) {
-		while (cap.read(image))
+            objs.clear();
-		{
+            yolov8->copy_from_Mat(image, size);
-			objs.clear();
+            auto start = std::chrono::system_clock::now();
-			yolov8->copy_from_Mat(image, size);
+            yolov8->infer();
-			auto start = std::chrono::system_clock::now();
+            auto end = std::chrono::system_clock::now();
-			yolov8->infer();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, num_labels);
-			auto end = std::chrono::system_clock::now();
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, num_labels);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
+            printf("cost %2.4lf ms\n", tc);
-			auto tc = (double)
+            cv::imshow("result", res);
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+            if (cv::waitKey(10) == 'q') {
-			printf("cost %2.4lf ms\n", tc);
+                break;
-			cv::imshow("result", res);
+            }
-			if (cv::waitKey(10) == 'q')
+        }
-			{
+    }
-				break;
+    else {
-			}
+        for (auto& path : imagePathList) {
-		}
+            objs.clear();
-	}
+            image = cv::imread(path);
-	else
+            yolov8->copy_from_Mat(image, size);
-	{
+            auto start = std::chrono::system_clock::now();
-		for (auto& path : imagePathList)
+            yolov8->infer();
-		{
+            auto end = std::chrono::system_clock::now();
-			objs.clear();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, num_labels);
-			image = cv::imread(path);
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
-			yolov8->copy_from_Mat(image, size);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			auto start = std::chrono::system_clock::now();
+            printf("cost %2.4lf ms\n", tc);
-			yolov8->infer();
+            cv::imshow("result", res);
-			auto end = std::chrono::system_clock::now();
+            cv::waitKey(0);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, num_labels);
+        }
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
+    }
-			auto tc = (double)
+    cv::destroyAllWindows();
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+    delete yolov8;
-			printf("cost %2.4lf ms\n", tc);
+    return 0;
-			cv::imshow("result", res);
-			cv::waitKey(0);
-		}
-	}
-	cv::destroyAllWindows();
-	delete yolov8;
-	return 0;
 }

csrc/jetson/detect/include/common.hpp

@@ -4,153 +4,138 @@
 
 #ifndef JETSON_DETECT_COMMON_HPP
 #define JETSON_DETECT_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace det
+namespace det {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float> rect;
-		cv::Rect_<float> rect;
+    int              label = 0;
-		int label = 0;
+    float            prob  = 0.0;
-		float prob = 0.0;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace det
-}
+#endif  // JETSON_DETECT_COMMON_HPP
-#endif //JETSON_DETECT_COMMON_HPP

csrc/jetson/detect/include/yolov8.hpp

@@ -3,422 +3,281 @@
 //
 #ifndef JETSON_DETECT_YOLOV8_HPP
 #define JETSON_DETECT_YOLOV8_HPP
-#include "fstream"
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include "fstream"
 using namespace det;
 
-class YOLOv8
+class YOLOv8 {
-{
 public:
-	explicit YOLOv8(const std::string& engine_file_path);
+    explicit YOLOv8(const std::string& engine_file_path);
-	~YOLOv8();
+    ~YOLOv8();
-
+
-	void make_pipe(bool warmup = true);
+    void                 make_pipe(bool warmup = true);
-	void copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image);
-	void copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
-	void letterbox(
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-		const cv::Mat& image,
+    void                 infer();
-		cv::Mat& out,
+    void                 postprocess(std::vector<Object>& objs);
-		cv::Size& size
+    static void          draw_objects(const cv::Mat&                                image,
-	);
+                                      cv::Mat&                                      res,
-	void infer();
+                                      const std::vector<Object>&                    objs,
-	void postprocess(std::vector<Object>& objs);
+                                      const std::vector<std::string>&               CLASS_NAMES,
-	static void draw_objects(
+                                      const std::vector<std::vector<unsigned int>>& COLORS);
-		const cv::Mat& image,
+    int                  num_bindings;
-		cv::Mat& res,
+    int                  num_inputs  = 0;
-		const std::vector<Object>& objs,
+    int                  num_outputs = 0;
-		const std::vector<std::string>& CLASS_NAMES,
+    std::vector<Binding> input_bindings;
-		const std::vector<std::vector<unsigned int>>& COLORS
+    std::vector<Binding> output_bindings;
-	);
+    std::vector<void*>   host_ptrs;
-	int num_bindings;
+    std::vector<void*>   device_ptrs;
-	int num_inputs = 0;
+
-	int num_outputs = 0;
+    PreParam pparam;
-	std::vector<Binding> input_bindings;
-	std::vector<Binding> output_bindings;
-	std::vector<void*> host_ptrs;
-	std::vector<void*> device_ptrs;
-
-	PreParam pparam;
-private:
-	nvinfer1::ICudaEngine* engine = nullptr;
-	nvinfer1::IRuntime* runtime = nullptr;
-	nvinfer1::IExecutionContext* context = nullptr;
-	cudaStream_t stream = nullptr;
-	Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
 YOLOv8::YOLOv8(const std::string& engine_file_path)
 {
-	std::ifstream file(engine_file_path, std::ios::binary);
+    std::ifstream file(engine_file_path, std::ios::binary);
-	assert(file.good());
+    assert(file.good());
-	file.seekg(0, std::ios::end);
+    file.seekg(0, std::ios::end);
-	auto size = file.tellg();
+    auto size = file.tellg();
-	file.seekg(0, std::ios::beg);
+    file.seekg(0, std::ios::beg);
-	char* trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
-	assert(trtModelStream);
+    assert(trtModelStream);
-	file.read(trtModelStream, size);
+    file.read(trtModelStream, size);
-	file.close();
+    file.close();
-	initLibNvInferPlugins(&this->gLogger, "");
+    initLibNvInferPlugins(&this->gLogger, "");
-	this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
-	assert(this->runtime != nullptr);
+    assert(this->runtime != nullptr);
-
+
-	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
-	assert(this->engine != nullptr);
+    assert(this->engine != nullptr);
-	delete[] trtModelStream;
+    delete[] trtModelStream;
-	this->context = this->engine->createExecutionContext();
+    this->context = this->engine->createExecutionContext();
-
+
-	assert(this->context != nullptr);
+    assert(this->context != nullptr);
-	cudaStreamCreate(&this->stream);
+    cudaStreamCreate(&this->stream);
-	this->num_bindings = this->engine->getNbBindings();
+    this->num_bindings = this->engine->getNbBindings();
-
+
-	for (int i = 0; i < this->num_bindings; ++i)
+    for (int i = 0; i < this->num_bindings; ++i) {
-	{
+        Binding            binding;
-		Binding binding;
+        nvinfer1::Dims     dims;
-		nvinfer1::Dims dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-		nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
-		std::string name = this->engine->getBindingName(i);
+        binding.name             = name;
-		binding.name = name;
+        binding.dsize            = type_to_size(dtype);
-		binding.dsize = type_to_size(dtype);
+
-
+        bool IsInput = engine->bindingIsInput(i);
-		bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
-		if (IsInput)
+            this->num_inputs += 1;
-		{
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-			this->num_inputs += 1;
+            binding.size = get_size_by_dims(dims);
-			dims = this->engine->getProfileDimensions(
+            binding.dims = dims;
-				i,
+            this->input_bindings.push_back(binding);
-				0,
+            // set max opt shape
-				nvinfer1::OptProfileSelector::kMAX);
+            this->context->setBindingDimensions(i, dims);
-			binding.size = get_size_by_dims(dims);
+        }
-			binding.dims = dims;
+        else {
-			this->input_bindings.push_back(binding);
+            dims         = this->context->getBindingDimensions(i);
-			// set max opt shape
+            binding.size = get_size_by_dims(dims);
-			this->context->setBindingDimensions(i, dims);
+            binding.dims = dims;
-
+            this->output_bindings.push_back(binding);
-		}
+            this->num_outputs += 1;
-		else
+        }
-		{
+    }
-			dims = this->context->getBindingDimensions(i);
-			binding.size = get_size_by_dims(dims);
-			binding.dims = dims;
-			this->output_bindings.push_back(binding);
-			this->num_outputs += 1;
-		}
-	}
-
 }
 
 YOLOv8::~YOLOv8()
 {
-	this->context->destroy();
+    this->context->destroy();
-	this->engine->destroy();
+    this->engine->destroy();
-	this->runtime->destroy();
+    this->runtime->destroy();
-	cudaStreamDestroy(this->stream);
+    cudaStreamDestroy(this->stream);
-	for (auto& ptr : this->device_ptrs)
+    for (auto& ptr : this->device_ptrs) {
-	{
+        CHECK(cudaFree(ptr));
-		CHECK(cudaFree(ptr));
+    }
-	}
+
-
+    for (auto& ptr : this->host_ptrs) {
-	for (auto& ptr : this->host_ptrs)
+        CHECK(cudaFreeHost(ptr));
-	{
+    }
-		CHECK(cudaFreeHost(ptr));
-	}
-
 }
 void YOLOv8::make_pipe(bool warmup)
 {
 
-	for (auto& bindings : this->input_bindings)
+    for (auto& bindings : this->input_bindings) {
-	{
+        void* d_ptr;
-		void* d_ptr;
+        CHECK(cudaMalloc(&d_ptr, bindings.size * bindings.dsize));
-		CHECK(cudaMalloc(
+        this->device_ptrs.push_back(d_ptr);
-			&d_ptr,
+    }
-			bindings.size * bindings.dsize)
+
-		);
+    for (auto& bindings : this->output_bindings) {
-		this->device_ptrs.push_back(d_ptr);
+        void * d_ptr, *h_ptr;
-	}
+        size_t size = bindings.size * bindings.dsize;
-
+        CHECK(cudaMalloc(&d_ptr, size));
-	for (auto& bindings : this->output_bindings)
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-	{
+        this->device_ptrs.push_back(d_ptr);
-		void* d_ptr, * h_ptr;
+        this->host_ptrs.push_back(h_ptr);
-		size_t size = bindings.size * bindings.dsize;
+    }
-		CHECK(cudaMalloc(
+
-			&d_ptr,
+    if (warmup) {
-			size)
+        for (int i = 0; i < 10; i++) {
-		);
+            for (auto& bindings : this->input_bindings) {
-		CHECK(cudaHostAlloc(
+                size_t size  = bindings.size * bindings.dsize;
-			&h_ptr,
+                void*  h_ptr = malloc(size);
-			size,
+                memset(h_ptr, 0, size);
-			0)
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-		);
+                free(h_ptr);
-		this->device_ptrs.push_back(d_ptr);
+            }
-		this->host_ptrs.push_back(h_ptr);
+            this->infer();
-	}
+        }
-
+        printf("model warmup 10 times\n");
-	if (warmup)
+    }
-	{
-		for (int i = 0; i < 10; i++)
-		{
-			for (auto& bindings : this->input_bindings)
-			{
-				size_t size = bindings.size * bindings.dsize;
-				void* h_ptr = malloc(size);
-				memset(h_ptr, 0, size);
-				CHECK(cudaMemcpyAsync(
-					this->device_ptrs[0],
-					h_ptr,
-					size,
-					cudaMemcpyHostToDevice,
-					this->stream)
-				);
-				free(h_ptr);
-			}
-			this->infer();
-		}
-		printf("model warmup 10 times\n");
-
-	}
 }
 
 void YOLOv8::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
 {
-	const float inp_h = size.height;
+    const float inp_h  = size.height;
-	const float inp_w = size.width;
+    const float inp_w  = size.width;
-	float height = image.rows;
+    float       height = image.rows;
-	float width = image.cols;
+    float       width  = image.cols;
-
+
-	float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-	int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-	int padh = std::round(height * r);
+    int   padh = std::round(height * r);
-
+
-	cv::Mat tmp;
+    cv::Mat tmp;
-	if ((int)width != padw || (int)height != padh)
+    if ((int)width != padw || (int)height != padh) {
-	{
+        cv::resize(image, tmp, cv::Size(padw, padh));
-		cv::resize(
+    }
-			image,
+    else {
-			tmp,
+        tmp = image.clone();
-			cv::Size(padw, padh)
+    }
-		);
+
-	}
+    float dw = inp_w - padw;
-	else
+    float dh = inp_h - padh;
-	{
+
-		tmp = image.clone();
+    dw /= 2.0f;
-	}
+    dh /= 2.0f;
-
+    int top    = int(std::round(dh - 0.1f));
-	float dw = inp_w - padw;
+    int bottom = int(std::round(dh + 0.1f));
-	float dh = inp_h - padh;
+    int left   = int(std::round(dw - 0.1f));
-
+    int right  = int(std::round(dw + 0.1f));
-	dw /= 2.0f;
+
-	dh /= 2.0f;
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-	int top = int(std::round(dh - 0.1f));
+
-	int bottom = int(std::round(dh + 0.1f));
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-	int left = int(std::round(dw - 0.1f));
+    this->pparam.ratio  = 1 / r;
-	int right = int(std::round(dw + 0.1f));
+    this->pparam.dw     = dw;
-
+    this->pparam.dh     = dh;
-	cv::copyMakeBorder(
+    this->pparam.height = height;
-		tmp,
+    this->pparam.width  = width;
-		tmp,
+    ;
-		top,
-		bottom,
-		left,
-		right,
-		cv::BORDER_CONSTANT,
-		{ 114, 114, 114 }
-	);
-
-	cv::dnn::blobFromImage(tmp,
-		out,
-		1 / 255.f,
-		cv::Size(),
-		cv::Scalar(0, 0, 0),
-		true,
-		false,
-		CV_32F
-	);
-	this->pparam.ratio = 1 / r;
-	this->pparam.dw = dw;
-	this->pparam.dh = dh;
-	this->pparam.height = height;
-	this->pparam.width = width;;
 }
 
 void YOLOv8::copy_from_Mat(const cv::Mat& image)
 {
-	cv::Mat nchw;
+    cv::Mat  nchw;
-	auto& in_binding = this->input_bindings[0];
+    auto&    in_binding = this->input_bindings[0];
-	auto width = in_binding.dims.d[3];
+    auto     width      = in_binding.dims.d[3];
-	auto height = in_binding.dims.d[2];
+    auto     height     = in_binding.dims.d[2];
-	cv::Size size{ width, height };
+    cv::Size size{width, height};
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+
-		nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-		size
+
-	);
+    CHECK(cudaMemcpyAsync(
-
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{
-				4,
-				{ 1, 3, height, width }
-			}
-	);
-
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8::copy_from_Mat(const cv::Mat& image, cv::Size& size)
 {
-	cv::Mat nchw;
+    cv::Mat nchw;
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-		nchw,
+    CHECK(cudaMemcpyAsync(
-		size
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	);
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{ 4,
-			  { 1, 3, size.height, size.width }
-			}
-	);
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8::infer()
 {
 
-	this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-		this->device_ptrs.data(),
+    for (int i = 0; i < this->num_outputs; i++) {
-		this->stream,
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		nullptr
+        CHECK(cudaMemcpyAsync(
-	);
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-	for (int i = 0; i < this->num_outputs; i++)
+    }
-	{
+    cudaStreamSynchronize(this->stream);
-		size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		CHECK(cudaMemcpyAsync(this->host_ptrs[i],
-			this->device_ptrs[i + this->num_inputs],
-			osize,
-			cudaMemcpyDeviceToHost,
-			this->stream)
-		);
-
-	}
-	cudaStreamSynchronize(this->stream);
-
 }
 
 void YOLOv8::postprocess(std::vector<Object>& objs)
 {
-	objs.clear();
+    objs.clear();
-	int* num_dets = static_cast<int*>(this->host_ptrs[0]);
+    int*  num_dets = static_cast<int*>(this->host_ptrs[0]);
-	auto* boxes = static_cast<float*>(this->host_ptrs[1]);
+    auto* boxes    = static_cast<float*>(this->host_ptrs[1]);
-	auto* scores = static_cast<float*>(this->host_ptrs[2]);
+    auto* scores   = static_cast<float*>(this->host_ptrs[2]);
-	int* labels = static_cast<int*>(this->host_ptrs[3]);
+    int*  labels   = static_cast<int*>(this->host_ptrs[3]);
-	auto& dw = this->pparam.dw;
+    auto& dw       = this->pparam.dw;
-	auto& dh = this->pparam.dh;
+    auto& dh       = this->pparam.dh;
-	auto& width = this->pparam.width;
+    auto& width    = this->pparam.width;
-	auto& height = this->pparam.height;
+    auto& height   = this->pparam.height;
-	auto& ratio = this->pparam.ratio;
+    auto& ratio    = this->pparam.ratio;
-	for (int i = 0; i < num_dets[0]; i++)
+    for (int i = 0; i < num_dets[0]; i++) {
-	{
+        float* ptr = boxes + i * 4;
-		float* ptr = boxes + i * 4;
+
-
+        float x0 = *ptr++ - dw;
-		float x0 = *ptr++ - dw;
+        float y0 = *ptr++ - dh;
-		float y0 = *ptr++ - dh;
+        float x1 = *ptr++ - dw;
-		float x1 = *ptr++ - dw;
+        float y1 = *ptr - dh;
-		float y1 = *ptr - dh;
+
-
+        x0 = clamp(x0 * ratio, 0.f, width);
-		x0 = clamp(x0 * ratio, 0.f, width);
+        y0 = clamp(y0 * ratio, 0.f, height);
-		y0 = clamp(y0 * ratio, 0.f, height);
+        x1 = clamp(x1 * ratio, 0.f, width);
-		x1 = clamp(x1 * ratio, 0.f, width);
+        y1 = clamp(y1 * ratio, 0.f, height);
-		y1 = clamp(y1 * ratio, 0.f, height);
+        Object obj;
-		Object obj;
+        obj.rect.x      = x0;
-		obj.rect.x = x0;
+        obj.rect.y      = y0;
-		obj.rect.y = y0;
+        obj.rect.width  = x1 - x0;
-		obj.rect.width = x1 - x0;
+        obj.rect.height = y1 - y0;
-		obj.rect.height = y1 - y0;
+        obj.prob        = *(scores + i);
-		obj.prob = *(scores + i);
+        obj.label       = *(labels + i);
-		obj.label = *(labels + i);
+        objs.push_back(obj);
-		objs.push_back(obj);
+    }
-	}
 }
 
-void YOLOv8::draw_objects(
+void YOLOv8::draw_objects(const cv::Mat&                                image,
-	const cv::Mat& image,
+                          cv::Mat&                                      res,
-	cv::Mat& res,
+                          const std::vector<Object>&                    objs,
-	const std::vector<Object>& objs,
+                          const std::vector<std::string>&               CLASS_NAMES,
-	const std::vector<std::string>& CLASS_NAMES,
+                          const std::vector<std::vector<unsigned int>>& COLORS)
-	const std::vector<std::vector<unsigned int>>& COLORS
-)
 {
-	res = image.clone();
+    res = image.clone();
-	for (auto& obj : objs)
+    for (auto& obj : objs) {
-	{
+        cv::Scalar color = cv::Scalar(COLORS[obj.label][0], COLORS[obj.label][1], COLORS[obj.label][2]);
-		cv::Scalar color = cv::Scalar(
+        cv::rectangle(res, obj.rect, color, 2);
-			COLORS[obj.label][0],
+
-			COLORS[obj.label][1],
+        char text[256];
-			COLORS[obj.label][2]
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[obj.label].c_str(), obj.prob * 100);
-		);
+
-		cv::rectangle(
+        int      baseLine   = 0;
-			res,
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-			obj.rect,
+
-			color,
+        int x = (int)obj.rect.x;
-			2
+        int y = (int)obj.rect.y + 1;
-		);
+
-
+        if (y > res.rows)
-		char text[256];
+            y = res.rows;
-		sprintf(
+
-			text,
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-			"%s %.1f%%",
+
-			CLASS_NAMES[obj.label].c_str(),
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-			obj.prob * 100
+    }
-		);
-
-		int baseLine = 0;
-		cv::Size label_size = cv::getTextSize(
-			text,
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			1,
-			&baseLine
-		);
-
-		int x = (int)obj.rect.x;
-		int y = (int)obj.rect.y + 1;
-
-		if (y > res.rows)
-			y = res.rows;
-
-		cv::rectangle(
-			res,
-			cv::Rect(x, y, label_size.width, label_size.height + baseLine),
-			{ 0, 0, 255 },
-			-1
-		);
-
-		cv::putText(
-			res,
-			text,
-			cv::Point(x, y + label_size.height),
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			{ 255, 255, 255 },
-			1
-		);
-	}
 }
-#endif //JETSON_DETECT_YOLOV8_HPP
+#endif  // JETSON_DETECT_YOLOV8_HPP

csrc/jetson/detect/main.cpp

@@ -2,157 +2,116 @@
 // Created by ubuntu on 3/16/23.
 //
 #include "chrono"
-#include "yolov8.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8.hpp"
 
 const std::vector<std::string> CLASS_NAMES = {
-	"person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",
-	"train", "truck", "boat", "traffic light", "fire hydrant",
+    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",    "parking meter", "bench",
-	"stop sign", "parking meter", "bench", "bird", "cat",
+    "bird",           "cat",        "dog",           "horse",         "sheep",        "cow",           "elephant",
-	"dog", "horse", "sheep", "cow", "elephant",
+    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",     "handbag",       "tie",
-	"bear", "zebra", "giraffe", "backpack", "umbrella",
+    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball",  "kite",          "baseball bat",
-	"handbag", "tie", "suitcase", "frisbee", "skis",
+    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",       "wine glass",    "cup",
-	"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "fork",           "knife",      "spoon",         "bowl",          "banana",       "apple",         "sandwich",
-	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
+    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",        "donut",         "cake",
-	"cup", "fork", "knife", "spoon", "bowl",
+    "chair",          "couch",      "potted plant",  "bed",           "dining table", "toilet",        "tv",
-	"banana", "apple", "sandwich", "orange", "broccoli",
+    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",   "microwave",     "oven",
-	"carrot", "hot dog", "pizza", "donut", "cake",
+    "toaster",        "sink",       "refrigerator",  "book",          "clock",        "vase",          "scissors",
-	"chair", "couch", "potted plant", "bed", "dining table",
+    "teddy bear",     "hair drier", "toothbrush"};
-	"toilet", "tv", "laptop", "mouse", "remote",
-	"keyboard", "cell phone", "microwave", "oven",
-	"toaster", "sink", "refrigerator", "book", "clock", "vase",
-	"scissors", "teddy bear", "hair drier", "toothbrush" };
 
 const std::vector<std::vector<unsigned int>> COLORS = {
-	{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
+    {0, 114, 189},   {217, 83, 25},   {237, 177, 32},  {126, 47, 142},  {119, 172, 48},  {77, 190, 238},
-	{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
+    {162, 20, 47},   {76, 76, 76},    {153, 153, 153}, {255, 0, 0},     {255, 128, 0},   {191, 191, 0},
-	{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
+    {0, 255, 0},     {0, 0, 255},     {170, 0, 255},   {85, 85, 0},     {85, 170, 0},    {85, 255, 0},
-	{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
+    {170, 85, 0},    {170, 170, 0},   {170, 255, 0},   {255, 85, 0},    {255, 170, 0},   {255, 255, 0},
-	{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
+    {0, 85, 128},    {0, 170, 128},   {0, 255, 128},   {85, 0, 128},    {85, 85, 128},   {85, 170, 128},
-	{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
+    {85, 255, 128},  {170, 0, 128},   {170, 85, 128},  {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
-	{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
+    {255, 85, 128},  {255, 170, 128}, {255, 255, 128}, {0, 85, 255},    {0, 170, 255},   {0, 255, 255},
-	{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
+    {85, 0, 255},    {85, 85, 255},   {85, 170, 255},  {85, 255, 255},  {170, 0, 255},   {170, 85, 255},
-	{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
+    {170, 170, 255}, {170, 255, 255}, {255, 0, 255},   {255, 85, 255},  {255, 170, 255}, {85, 0, 0},
-	{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
+    {128, 0, 0},     {170, 0, 0},     {212, 0, 0},     {255, 0, 0},     {0, 43, 0},      {0, 85, 0},
-	{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
+    {0, 128, 0},     {0, 170, 0},     {0, 212, 0},     {0, 255, 0},     {0, 0, 43},      {0, 0, 85},
-	{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
+    {0, 0, 128},     {0, 0, 170},     {0, 0, 212},     {0, 0, 255},     {0, 0, 0},       {36, 36, 36},
-	{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
+    {73, 73, 73},    {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
-	{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
+    {80, 183, 189},  {128, 128, 0}};
-	{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
-	{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
-	{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
-	{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
-	{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
-	{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
-	{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
-	{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
-	{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
-	{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
-	{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
-	{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
-	{ 80, 183, 189 }, { 128, 128, 0 }
-};
 
 int main(int argc, char** argv)
 {
-	const std::string engine_file_path{ argv[1] };
+    const std::string engine_file_path{argv[1]};
-	const std::string path{ argv[2] };
+    const std::string path{argv[2]};
 
-	std::vector<std::string> imagePathList;
+    std::vector<std::string> imagePathList;
-	bool isVideo{ false };
+    bool                     isVideo{false};
 
-	assert(argc == 3);
+    assert(argc == 3);
 
-	auto yolov8 = new YOLOv8(engine_file_path);
+    auto yolov8 = new YOLOv8(engine_file_path);
-	yolov8->make_pipe(true);
+    yolov8->make_pipe(true);
 
-	if (IsFile(path))
+    if (IsFile(path)) {
-	{
+        std::string suffix = path.substr(path.find_last_of('.') + 1);
-		std::string suffix = path.substr(path.find_last_of('.') + 1);
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-		if (
+            imagePathList.push_back(path);
-			suffix == "jpg" ||
+        }
-				suffix == "jpeg" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-				suffix == "png"
+                 || suffix == "mkv") {
-			)
+            isVideo = true;
-		{
+        }
-			imagePathList.push_back(path);
+        else {
-		}
+            printf("suffix %s is wrong !!!\n", suffix.c_str());
-		else if (
+            std::abort();
-			suffix == "mp4" ||
+        }
-				suffix == "avi" ||
+    }
-				suffix == "m4v" ||
+    else if (IsFolder(path)) {
-				suffix == "mpeg" ||
+        cv::glob(path + "/*.jpg", imagePathList);
-				suffix == "mov" ||
+    }
-				suffix == "mkv"
-			)
-		{
-			isVideo = true;
-		}
-		else
-		{
-			printf("suffix %s is wrong !!!\n", suffix.c_str());
-			std::abort();
-		}
-	}
-	else if (IsFolder(path))
-	{
-		cv::glob(path + "/*.jpg", imagePathList);
-	}
 
-	cv::Mat res, image;
+    cv::Mat             res, image;
-	cv::Size size = cv::Size{ 640, 640 };
+    cv::Size            size = cv::Size{640, 640};
-	std::vector<Object> objs;
+    std::vector<Object> objs;
 
-	cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
+    cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
 
-	if (isVideo)
+    if (isVideo) {
-	{
+        cv::VideoCapture cap(path);
-		cv::VideoCapture cap(path);
 
-		if (!cap.isOpened())
+        if (!cap.isOpened()) {
-		{
+            printf("can not open %s\n", path.c_str());
-			printf("can not open %s\n", path.c_str());
+            return -1;
-			return -1;
+        }
-		}
+        while (cap.read(image)) {
-		while (cap.read(image))
+            objs.clear();
-		{
+            yolov8->copy_from_Mat(image, size);
-			objs.clear();
+            auto start = std::chrono::system_clock::now();
-			yolov8->copy_from_Mat(image, size);
+            yolov8->infer();
-			auto start = std::chrono::system_clock::now();
+            auto end = std::chrono::system_clock::now();
-			yolov8->infer();
+            yolov8->postprocess(objs);
-			auto end = std::chrono::system_clock::now();
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
-			yolov8->postprocess(objs);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
+            printf("cost %2.4lf ms\n", tc);
-			auto tc = (double)
+            cv::imshow("result", res);
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+            if (cv::waitKey(10) == 'q') {
-			printf("cost %2.4lf ms\n", tc);
+                break;
-			cv::imshow("result", res);
+            }
-			if (cv::waitKey(10) == 'q')
+        }
-			{
+    }
-				break;
+    else {
-			}
+        for (auto& path : imagePathList) {
-		}
+            objs.clear();
-	}
+            image = cv::imread(path);
-	else
+            yolov8->copy_from_Mat(image, size);
-	{
+            auto start = std::chrono::system_clock::now();
-		for (auto& path : imagePathList)
+            yolov8->infer();
-		{
+            auto end = std::chrono::system_clock::now();
-			objs.clear();
+            yolov8->postprocess(objs);
-			image = cv::imread(path);
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
-			yolov8->copy_from_Mat(image, size);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			auto start = std::chrono::system_clock::now();
+            printf("cost %2.4lf ms\n", tc);
-			yolov8->infer();
+            cv::imshow("result", res);
-			auto end = std::chrono::system_clock::now();
+            cv::waitKey(0);
-			yolov8->postprocess(objs);
+        }
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);
+    }
-			auto tc = (double)
+    cv::destroyAllWindows();
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+    delete yolov8;
-			printf("cost %2.4lf ms\n", tc);
+    return 0;
-			cv::imshow("result", res);
-			cv::waitKey(0);
-		}
-	}
-	cv::destroyAllWindows();
-	delete yolov8;
-	return 0;
 }

csrc/jetson/pose/include/common.hpp

@@ -4,154 +4,139 @@
 
 #ifndef JETSON_POSE_COMMON_HPP
 #define JETSON_POSE_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace pose
+namespace pose {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float>   rect;
-		cv::Rect_<float> rect;
+    int                label = 0;
-		int label = 0;
+    float              prob  = 0.0;
-		float prob = 0.0;
+    std::vector<float> kps;
-        std::vector<float> kps;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace pose
-}
+#endif  // JETSON_POSE_COMMON_HPP
-#endif //JETSON_POSE_COMMON_HPP

csrc/jetson/pose/include/yolov8-pose.hpp

@@ -4,73 +4,63 @@
 #ifndef JETSON_POSE_YOLOV8_POSE_HPP
 #define JETSON_POSE_YOLOV8_POSE_HPP
 
-#include "fstream"
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include "fstream"
 
 using namespace pose;
 
 class YOLOv8_pose {
 public:
-    explicit YOLOv8_pose(const std::string &engine_file_path);
+    explicit YOLOv8_pose(const std::string& engine_file_path);
 
     ~YOLOv8_pose();
 
     void make_pipe(bool warmup = true);
 
-    void copy_from_Mat(const cv::Mat &image);
+    void copy_from_Mat(const cv::Mat& image);
 
-    void copy_from_Mat(const cv::Mat &image, cv::Size &size);
+    void copy_from_Mat(const cv::Mat& image, cv::Size& size);
 
-    void letterbox(
+    void letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-            const cv::Mat &image,
-            cv::Mat &out,
-            cv::Size &size
-    );
 
     void infer();
 
-    void postprocess(
+    void postprocess(std::vector<Object>& objs, float score_thres = 0.25f, float iou_thres = 0.65f, int topk = 100);
-            std::vector<Object> &objs,
+
-            float score_thres = 0.25f,
+    static void draw_objects(const cv::Mat&                                image,
-            float iou_thres = 0.65f,
+                             cv::Mat&                                      res,
-            int topk = 100
+                             const std::vector<Object>&                    objs,
-    );
+                             const std::vector<std::vector<unsigned int>>& SKELETON,
-
+                             const std::vector<std::vector<unsigned int>>& KPS_COLORS,
-    static void draw_objects(
+                             const std::vector<std::vector<unsigned int>>& LIMB_COLORS);
-            const cv::Mat &image,
+
-            cv::Mat &res,
+    int                  num_bindings;
-            const std::vector<Object> &objs,
+    int                  num_inputs  = 0;
-            const std::vector<std::vector<unsigned int>> &SKELETON,
+    int                  num_outputs = 0;
-            const std::vector<std::vector<unsigned int>> &KPS_COLORS,
-            const std::vector<std::vector<unsigned int>> &LIMB_COLORS
-    );
-
-    int num_bindings;
-    int num_inputs = 0;
-    int num_outputs = 0;
     std::vector<Binding> input_bindings;
     std::vector<Binding> output_bindings;
-    std::vector<void *> host_ptrs;
+    std::vector<void*>   host_ptrs;
-    std::vector<void *> device_ptrs;
+    std::vector<void*>   device_ptrs;
 
     PreParam pparam;
-private:
-    nvinfer1::ICudaEngine *engine = nullptr;
-    nvinfer1::IRuntime *runtime = nullptr;
-    nvinfer1::IExecutionContext *context = nullptr;
-    cudaStream_t stream = nullptr;
-    Logger gLogger{nvinfer1::ILogger::Severity::kERROR};
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
-YOLOv8_pose::YOLOv8_pose(const std::string &engine_file_path) {
+YOLOv8_pose::YOLOv8_pose(const std::string& engine_file_path)
+{
     std::ifstream file(engine_file_path, std::ios::binary);
     assert(file.good());
     file.seekg(0, std::ios::end);
     auto size = file.tellg();
     file.seekg(0, std::ios::beg);
-    char *trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
     assert(trtModelStream);
     file.read(trtModelStream, size);
     file.close();
@@ -88,124 +78,97 @@ YOLOv8_pose::YOLOv8_pose(const std::string &engine_file_path) {
     this->num_bindings = this->engine->getNbBindings();
 
     for (int i = 0; i < this->num_bindings; ++i) {
-        Binding binding;
+        Binding            binding;
-        nvinfer1::Dims dims;
+        nvinfer1::Dims     dims;
         nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-        std::string name = this->engine->getBindingName(i);
+        std::string        name  = this->engine->getBindingName(i);
-        binding.name = name;
+        binding.name             = name;
-        binding.dsize = type_to_size(dtype);
+        binding.dsize            = type_to_size(dtype);
 
         bool IsInput = engine->bindingIsInput(i);
         if (IsInput) {
             this->num_inputs += 1;
-            dims = this->engine->getProfileDimensions(
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-                    i,
-                    0,
-                    nvinfer1::OptProfileSelector::kMAX);
             binding.size = get_size_by_dims(dims);
             binding.dims = dims;
             this->input_bindings.push_back(binding);
             // set max opt shape
             this->context->setBindingDimensions(i, dims);
-
+        }
-        } else {
+        else {
-            dims = this->context->getBindingDimensions(i);
+            dims         = this->context->getBindingDimensions(i);
             binding.size = get_size_by_dims(dims);
             binding.dims = dims;
             this->output_bindings.push_back(binding);
             this->num_outputs += 1;
         }
     }
-
 }
 
-YOLOv8_pose::~YOLOv8_pose() {
+YOLOv8_pose::~YOLOv8_pose()
+{
     this->context->destroy();
     this->engine->destroy();
     this->runtime->destroy();
     cudaStreamDestroy(this->stream);
-    for (auto &ptr: this->device_ptrs) {
+    for (auto& ptr : this->device_ptrs) {
         CHECK(cudaFree(ptr));
     }
 
-    for (auto &ptr: this->host_ptrs) {
+    for (auto& ptr : this->host_ptrs) {
         CHECK(cudaFreeHost(ptr));
     }
-
 }
 
-void YOLOv8_pose::make_pipe(bool warmup) {
+void YOLOv8_pose::make_pipe(bool warmup)
+{
 
-    for (auto &bindings: this->input_bindings) {
+    for (auto& bindings : this->input_bindings) {
-        void *d_ptr;
+        void* d_ptr;
-        CHECK(cudaMalloc(
+        CHECK(cudaMalloc(&d_ptr, bindings.size * bindings.dsize));
-                &d_ptr,
-                bindings.size * bindings.dsize
-                )
-        );
         this->device_ptrs.push_back(d_ptr);
     }
 
-    for (auto &bindings: this->output_bindings) {
+    for (auto& bindings : this->output_bindings) {
-        void *d_ptr, *h_ptr;
+        void * d_ptr, *h_ptr;
         size_t size = bindings.size * bindings.dsize;
-        CHECK(cudaMalloc(
+        CHECK(cudaMalloc(&d_ptr, size));
-                &d_ptr,
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-                size)
-        );
-        CHECK(cudaHostAlloc(
-                &h_ptr,
-                size,
-                0)
-        );
         this->device_ptrs.push_back(d_ptr);
         this->host_ptrs.push_back(h_ptr);
     }
 
     if (warmup) {
         for (int i = 0; i < 10; i++) {
-            for (auto &bindings: this->input_bindings) {
+            for (auto& bindings : this->input_bindings) {
-                size_t size = bindings.size * bindings.dsize;
+                size_t size  = bindings.size * bindings.dsize;
-                void *h_ptr = malloc(size);
+                void*  h_ptr = malloc(size);
                 memset(h_ptr, 0, size);
-                CHECK(cudaMemcpyAsync(
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-                        this->device_ptrs[0],
-                        h_ptr,
-                        size,
-                        cudaMemcpyHostToDevice,
-                        this->stream)
-                );
                 free(h_ptr);
             }
             this->infer();
         }
         printf("model warmup 10 times\n");
-
     }
 }
 
-void YOLOv8_pose::letterbox(
+void YOLOv8_pose::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
-        const cv::Mat &image,
+{
-        cv::Mat &out,
+    const float inp_h  = size.height;
-        cv::Size &size
+    const float inp_w  = size.width;
-) {
+    float       height = image.rows;
-    const float inp_h = size.height;
+    float       width  = image.cols;
-    const float inp_w = size.width;
-    float height = image.rows;
-    float width = image.cols;
 
-    float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-    int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-    int padh = std::round(height * r);
+    int   padh = std::round(height * r);
 
     cv::Mat tmp;
-    if ((int) width != padw || (int) height != padh) {
+    if ((int)width != padw || (int)height != padh) {
-        cv::resize(
+        cv::resize(image, tmp, cv::Size(padw, padh));
-                image,
+    }
-                tmp,
+    else {
-                cv::Size(padw, padh)
-        );
-    } else {
         tmp = image.clone();
     }
 
@@ -214,146 +177,83 @@ void YOLOv8_pose::letterbox(
 
     dw /= 2.0f;
     dh /= 2.0f;
-    int top = int(std::round(dh - 0.1f));
+    int top    = int(std::round(dh - 0.1f));
     int bottom = int(std::round(dh + 0.1f));
-    int left = int(std::round(dw - 0.1f));
+    int left   = int(std::round(dw - 0.1f));
-    int right = int(std::round(dw + 0.1f));
+    int right  = int(std::round(dw + 0.1f));
-
+
-    cv::copyMakeBorder(
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-            tmp,
+
-            tmp,
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-            top,
+    this->pparam.ratio  = 1 / r;
-            bottom,
+    this->pparam.dw     = dw;
-            left,
+    this->pparam.dh     = dh;
-            right,
-            cv::BORDER_CONSTANT,
-            {114, 114, 114}
-    );
-
-    cv::dnn::blobFromImage(tmp,
-                           out,
-                           1 / 255.f,
-                           cv::Size(),
-                           cv::Scalar(0, 0, 0),
-                           true,
-                           false,
-                           CV_32F
-    );
-    this->pparam.ratio = 1 / r;
-    this->pparam.dw = dw;
-    this->pparam.dh = dh;
     this->pparam.height = height;
-    this->pparam.width = width;;
+    this->pparam.width  = width;
+    ;
 }
 
-void YOLOv8_pose::copy_from_Mat(const cv::Mat &image) {
+void YOLOv8_pose::copy_from_Mat(const cv::Mat& image)
-    cv::Mat nchw;
+{
-    auto &in_binding = this->input_bindings[0];
+    cv::Mat  nchw;
-    auto width = in_binding.dims.d[3];
+    auto&    in_binding = this->input_bindings[0];
-    auto height = in_binding.dims.d[2];
+    auto     width      = in_binding.dims.d[3];
+    auto     height     = in_binding.dims.d[2];
     cv::Size size{width, height};
-    this->letterbox(
+    this->letterbox(image, nchw, size);
-            image,
+
-            nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-            size
-    );
-
-    this->context->setBindingDimensions(
-            0,
-            nvinfer1::Dims
-                    {
-                            4,
-                            {1, 3, height, width}
-                    }
-    );
 
     CHECK(cudaMemcpyAsync(
-            this->device_ptrs[0],
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-            nchw.ptr<float>(),
-            nchw.total() * nchw.elemSize(),
-            cudaMemcpyHostToDevice,
-            this->stream)
-    );
 }
 
-void YOLOv8_pose::copy_from_Mat(const cv::Mat &image, cv::Size &size) {
+void YOLOv8_pose::copy_from_Mat(const cv::Mat& image, cv::Size& size)
+{
     cv::Mat nchw;
-    this->letterbox(
+    this->letterbox(image, nchw, size);
-            image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-            nchw,
-            size
-    );
-    this->context->setBindingDimensions(
-            0,
-            nvinfer1::Dims
-                    {4,
-                     {1, 3, size.height, size.width}
-                    }
-    );
     CHECK(cudaMemcpyAsync(
-            this->device_ptrs[0],
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-            nchw.ptr<float>(),
-            nchw.total() * nchw.elemSize(),
-            cudaMemcpyHostToDevice,
-            this->stream)
-    );
 }
 
-void YOLOv8_pose::infer() {
+void YOLOv8_pose::infer()
+{
 
-    this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-            this->device_ptrs.data(),
-            this->stream,
-            nullptr
-    );
     for (int i = 0; i < this->num_outputs; i++) {
         size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-        CHECK(cudaMemcpyAsync(this->host_ptrs[i],
+        CHECK(cudaMemcpyAsync(
-                              this->device_ptrs[i + this->num_inputs],
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-                              osize,
-                              cudaMemcpyDeviceToHost,
-                              this->stream)
-        );
-
     }
     cudaStreamSynchronize(this->stream);
-
 }
 
-void YOLOv8_pose::postprocess(
+void YOLOv8_pose::postprocess(std::vector<Object>& objs, float score_thres, float iou_thres, int topk)
-        std::vector<Object> &objs,
+{
-        float score_thres,
-        float iou_thres,
-        int topk
-) {
     objs.clear();
     auto num_channels = this->output_bindings[0].dims.d[1];
-    auto num_anchors = this->output_bindings[0].dims.d[2];
+    auto num_anchors  = this->output_bindings[0].dims.d[2];
-
+
-    auto &dw = this->pparam.dw;
+    auto& dw     = this->pparam.dw;
-    auto &dh = this->pparam.dh;
+    auto& dh     = this->pparam.dh;
-    auto &width = this->pparam.width;
+    auto& width  = this->pparam.width;
-    auto &height = this->pparam.height;
+    auto& height = this->pparam.height;
-    auto &ratio = this->pparam.ratio;
+    auto& ratio  = this->pparam.ratio;
-
+
-    std::vector<cv::Rect> bboxes;
+    std::vector<cv::Rect>           bboxes;
-    std::vector<float> scores;
+    std::vector<float>              scores;
-    std::vector<int> labels;
+    std::vector<int>                labels;
-    std::vector<int> indices;
+    std::vector<int>                indices;
     std::vector<std::vector<float>> kpss;
 
-    cv::Mat output = cv::Mat(
+    cv::Mat output = cv::Mat(num_channels, num_anchors, CV_32F, static_cast<float*>(this->host_ptrs[0]));
-            num_channels,
+    output         = output.t();
-            num_anchors,
-            CV_32F,
-            static_cast<float *>(this->host_ptrs[0])
-    );
-    output = output.t();
     for (int i = 0; i < num_anchors; i++) {
-        auto row_ptr = output.row(i).ptr<float>();
+        auto row_ptr    = output.row(i).ptr<float>();
         auto bboxes_ptr = row_ptr;
         auto scores_ptr = row_ptr + 4;
-        auto kps_ptr = row_ptr + 5;
+        auto kps_ptr    = row_ptr + 5;
 
         float score = *scores_ptr;
         if (score > score_thres) {
@@ -368,17 +268,17 @@ void YOLOv8_pose::postprocess(
             float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);
 
             cv::Rect_<float> bbox;
-            bbox.x = x0;
+            bbox.x      = x0;
-            bbox.y = y0;
+            bbox.y      = y0;
-            bbox.width = x1 - x0;
+            bbox.width  = x1 - x0;
             bbox.height = y1 - y0;
             std::vector<float> kps;
             for (int k = 0; k < 17; k++) {
                 float kps_x = (*(kps_ptr + 3 * k) - dw) * ratio;
                 float kps_y = (*(kps_ptr + 3 * k + 1) - dh) * ratio;
                 float kps_s = *(kps_ptr + 3 * k + 2);
-                kps_x = clamp(kps_x, 0.f, width);
+                kps_x       = clamp(kps_x, 0.f, width);
-                kps_y = clamp(kps_y, 0.f, height);
+                kps_y       = clamp(kps_y, 0.f, height);
                 kps.push_back(kps_x);
                 kps.push_back(kps_y);
                 kps.push_back(kps_s);
@@ -392,110 +292,68 @@ void YOLOv8_pose::postprocess(
     }
 
 #ifdef BATCHED_NMS
-    cv::dnn::NMSBoxesBatched(
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
-            bboxes,
-            scores,
-            labels,
-            score_thres,
-            iou_thres,
-            indices
-    );
 #else
-    cv::dnn::NMSBoxes(
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
-            bboxes,
-            scores,
-            score_thres,
-            iou_thres,
-            indices
-    );
 #endif
 
     int cnt = 0;
-    for (auto &i: indices) {
+    for (auto& i : indices) {
         if (cnt >= topk) {
             break;
         }
         Object obj;
-        obj.rect = bboxes[i];
+        obj.rect  = bboxes[i];
-        obj.prob = scores[i];
+        obj.prob  = scores[i];
         obj.label = labels[i];
-        obj.kps = kpss[i];
+        obj.kps   = kpss[i];
         objs.push_back(obj);
         cnt += 1;
     }
 }
 
-void YOLOv8_pose::draw_objects(
+void YOLOv8_pose::draw_objects(const cv::Mat&                                image,
-        const cv::Mat &image,
+                               cv::Mat&                                      res,
-        cv::Mat &res,
+                               const std::vector<Object>&                    objs,
-        const std::vector<Object> &objs,
+                               const std::vector<std::vector<unsigned int>>& SKELETON,
-        const std::vector<std::vector<unsigned int>> &SKELETON,
+                               const std::vector<std::vector<unsigned int>>& KPS_COLORS,
-        const std::vector<std::vector<unsigned int>> &KPS_COLORS,
+                               const std::vector<std::vector<unsigned int>>& LIMB_COLORS)
-        const std::vector<std::vector<unsigned int>> &LIMB_COLORS
+{
-) {
+    res                 = image.clone();
-    res = image.clone();
     const int num_point = 17;
-    for (auto &obj: objs) {
+    for (auto& obj : objs) {
-        cv::rectangle(
+        cv::rectangle(res, obj.rect, {0, 0, 255}, 2);
-                res,
-                obj.rect,
-                {0, 0, 255},
-                2
-        );
 
         char text[256];
-        sprintf(
+        sprintf(text, "person %.1f%%", obj.prob * 100);
-                text,
+
-                "person %.1f%%",
+        int      baseLine   = 0;
-                obj.prob * 100
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-        );
+
-
+        int x = (int)obj.rect.x;
-        int baseLine = 0;
+        int y = (int)obj.rect.y + 1;
-        cv::Size label_size = cv::getTextSize(
-                text,
-                cv::FONT_HERSHEY_SIMPLEX,
-                0.4,
-                1,
-                &baseLine
-        );
-
-        int x = (int) obj.rect.x;
-        int y = (int) obj.rect.y + 1;
 
         if (y > res.rows)
             y = res.rows;
 
-        cv::rectangle(
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-                res,
+
-                cv::Rect(x, y, label_size.width, label_size.height + baseLine),
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-                {0, 0, 255},
+
-                -1
+        auto& kps = obj.kps;
-        );
-
-        cv::putText(
-                res,
-                text,
-                cv::Point(x, y + label_size.height),
-                cv::FONT_HERSHEY_SIMPLEX,
-                0.4,
-                {255, 255, 255},
-                1
-        );
-
-        auto &kps = obj.kps;
         for (int k = 0; k < num_point + 2; k++) {
             if (k < num_point) {
-                int kps_x = std::round(kps[k * 3]);
+                int   kps_x = std::round(kps[k * 3]);
-                int kps_y = std::round(kps[k * 3 + 1]);
+                int   kps_y = std::round(kps[k * 3 + 1]);
                 float kps_s = kps[k * 3 + 2];
                 if (kps_s > 0.5f) {
                     cv::Scalar kps_color = cv::Scalar(KPS_COLORS[k][0], KPS_COLORS[k][1], KPS_COLORS[k][2]);
                     cv::circle(res, {kps_x, kps_y}, 5, kps_color, -1);
                 }
             }
-            auto &ske = SKELETON[k];
+            auto& ske    = SKELETON[k];
-            int pos1_x = std::round(kps[(ske[0] - 1) * 3]);
+            int   pos1_x = std::round(kps[(ske[0] - 1) * 3]);
-            int pos1_y = std::round(kps[(ske[0] - 1) * 3 + 1]);
+            int   pos1_y = std::round(kps[(ske[0] - 1) * 3 + 1]);
 
             int pos2_x = std::round(kps[(ske[1] - 1) * 3]);
             int pos2_y = std::round(kps[(ske[1] - 1) * 3 + 1]);
@@ -503,7 +361,6 @@ void YOLOv8_pose::draw_objects(
             float pos1_s = kps[(ske[0] - 1) * 3 + 2];
             float pos2_s = kps[(ske[1] - 1) * 3 + 2];
 
-
             if (pos1_s > 0.5f && pos2_s > 0.5f) {
                 cv::Scalar limb_color = cv::Scalar(LIMB_COLORS[k][0], LIMB_COLORS[k][1], LIMB_COLORS[k][2]);
                 cv::line(res, {pos1_x, pos1_y}, {pos2_x, pos2_y}, limb_color, 2);
@@ -512,4 +369,4 @@ void YOLOv8_pose::draw_objects(
     }
 }
 
-#endif //JETSON_POSE_YOLOV8_POSE_HPP
+#endif  // JETSON_POSE_YOLOV8_POSE_HPP

csrc/jetson/pose/main.cpp

@@ -2,70 +2,69 @@
 // Created by ubuntu on 4/7/23.
 //
 #include "chrono"
-#include "yolov8-pose.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8-pose.hpp"
 
-
+const std::vector<std::vector<unsigned int>> KPS_COLORS = {{0, 255, 0},
-const std::vector<std::vector<unsigned int>> KPS_COLORS =
+                                                           {0, 255, 0},
-        {{0,   255, 0},
+                                                           {0, 255, 0},
-         {0,   255, 0},
+                                                           {0, 255, 0},
-         {0,   255, 0},
+                                                           {0, 255, 0},
-         {0,   255, 0},
+                                                           {255, 128, 0},
-         {0,   255, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {51, 153, 255},
-         {255, 128, 0},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255}};
-         {51,  153, 255},
-         {51,  153, 255}};
 
 const std::vector<std::vector<unsigned int>> SKELETON = {{16, 14},
                                                          {14, 12},
                                                          {17, 15},
                                                          {15, 13},
                                                          {12, 13},
-                                                         {6,  12},
+                                                         {6, 12},
-                                                         {7,  13},
+                                                         {7, 13},
-                                                         {6,  7},
+                                                         {6, 7},
-                                                         {6,  8},
+                                                         {6, 8},
-                                                         {7,  9},
+                                                         {7, 9},
-                                                         {8,  10},
+                                                         {8, 10},
-                                                         {9,  11},
+                                                         {9, 11},
-                                                         {2,  3},
+                                                         {2, 3},
-                                                         {1,  2},
+                                                         {1, 2},
-                                                         {1,  3},
+                                                         {1, 3},
-                                                         {2,  4},
+                                                         {2, 4},
-                                                         {3,  5},
+                                                         {3, 5},
-                                                         {4,  6},
+                                                         {4, 6},
-                                                         {5,  7}};
+                                                         {5, 7}};
 
-const std::vector<std::vector<unsigned int>> LIMB_COLORS = {{51,  153, 255},
+const std::vector<std::vector<unsigned int>> LIMB_COLORS = {{51, 153, 255},
-                                                            {51,  153, 255},
+                                                            {51, 153, 255},
-                                                            {51,  153, 255},
+                                                            {51, 153, 255},
-                                                            {51,  153, 255},
+                                                            {51, 153, 255},
-                                                            {255, 51,  255},
+                                                            {255, 51, 255},
-                                                            {255, 51,  255},
+                                                            {255, 51, 255},
-                                                            {255, 51,  255},
+                                                            {255, 51, 255},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0}};
+                                                            {0, 255, 0}};
 
-int main(int argc, char **argv) {
+int main(int argc, char** argv)
+{
     // cuda:0
     cudaSetDevice(0);
 
@@ -73,7 +72,7 @@ int main(int argc, char **argv) {
     const std::string path{argv[2]};
 
     std::vector<std::string> imagePathList;
-    bool isVideo{false};
+    bool                     isVideo{false};
 
     assert(argc == 3);
 
@@ -82,34 +81,27 @@ int main(int argc, char **argv) {
 
     if (IsFile(path)) {
         std::string suffix = path.substr(path.find_last_of('.') + 1);
-        if (
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-                suffix == "jpg" ||
-                suffix == "jpeg" ||
-                suffix == "png"
-                ) {
             imagePathList.push_back(path);
-        } else if (
+        }
-                suffix == "mp4" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-                suffix == "avi" ||
+                 || suffix == "mkv") {
-                suffix == "m4v" ||
-                suffix == "mpeg" ||
-                suffix == "mov" ||
-                suffix == "mkv"
-                ) {
             isVideo = true;
-        } else {
+        }
+        else {
             printf("suffix %s is wrong !!!\n", suffix.c_str());
             std::abort();
         }
-    } else if (IsFolder(path)) {
+    }
+    else if (IsFolder(path)) {
         cv::glob(path + "/*.jpg", imagePathList);
     }
 
-    cv::Mat res, image;
+    cv::Mat  res, image;
-    cv::Size size = cv::Size{640, 640};
+    cv::Size size        = cv::Size{640, 640};
-    int topk = 100;
+    int      topk        = 100;
-    float score_thres = 0.25f;
+    float    score_thres = 0.25f;
-    float iou_thres = 0.65f;
+    float    iou_thres   = 0.65f;
 
     std::vector<Object> objs;
 
@@ -130,16 +122,16 @@ int main(int argc, char **argv) {
             auto end = std::chrono::system_clock::now();
             yolov8_pose->postprocess(objs, score_thres, iou_thres, topk);
             yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS);
-            auto tc = (double)
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-                              std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
             printf("cost %2.4lf ms\n", tc);
             cv::imshow("result", res);
             if (cv::waitKey(10) == 'q') {
                 break;
             }
         }
-    } else {
+    }
-        for (auto &path: imagePathList) {
+    else {
+        for (auto& path : imagePathList) {
             objs.clear();
             image = cv::imread(path);
             yolov8_pose->copy_from_Mat(image, size);
@@ -148,8 +140,7 @@ int main(int argc, char **argv) {
             auto end = std::chrono::system_clock::now();
             yolov8_pose->postprocess(objs, score_thres, iou_thres, topk);
             yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS);
-            auto tc = (double)
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-                              std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
             printf("cost %2.4lf ms\n", tc);
             cv::imshow("result", res);
             cv::waitKey(0);

csrc/jetson/segment/include/common.hpp

@@ -4,154 +4,139 @@
 
 #ifndef JETSON_SEGMENT_COMMON_HPP
 #define JETSON_SEGMENT_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace seg
+namespace seg {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float> rect;
-		cv::Rect_<float> rect;
+    int              label = 0;
-		int label = 0;
+    float            prob  = 0.0;
-		float prob = 0.0;
+    cv::Mat          boxMask;
-		cv::Mat boxMask;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace seg
-}
+#endif  // JETSON_SEGMENT_COMMON_HPP
-#endif //JETSON_SEGMENT_COMMON_HPP

csrc/jetson/segment/include/yolov8-seg.hpp

@@ -3,546 +3,355 @@
 //
 #ifndef JETSON_SEGMENT_YOLOV8_SEG_HPP
 #define JETSON_SEGMENT_YOLOV8_SEG_HPP
-#include <fstream>
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include <fstream>
 
 using namespace seg;
 
-class YOLOv8_seg
+class YOLOv8_seg {
-{
 public:
-	explicit YOLOv8_seg(const std::string& engine_file_path);
+    explicit YOLOv8_seg(const std::string& engine_file_path);
-	~YOLOv8_seg();
+    ~YOLOv8_seg();
-
+
-	void make_pipe(bool warmup = true);
+    void                 make_pipe(bool warmup = true);
-	void copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image);
-	void copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
-	void letterbox(
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-		const cv::Mat& image,
+    void                 infer();
-		cv::Mat& out,
+    void                 postprocess(std::vector<Object>& objs,
-		cv::Size& size
+                                     float                score_thres  = 0.25f,
-	);
+                                     float                iou_thres    = 0.65f,
-	void infer();
+                                     int                  topk         = 100,
-	void postprocess(
+                                     int                  seg_channels = 32,
-		std::vector<Object>& objs,
+                                     int                  seg_h        = 160,
-		float score_thres = 0.25f,
+                                     int                  seg_w        = 160);
-		float iou_thres = 0.65f,
+    static void          draw_objects(const cv::Mat&                                image,
-		int topk = 100,
+                                      cv::Mat&                                      res,
-		int seg_channels = 32,
+                                      const std::vector<Object>&                    objs,
-		int seg_h = 160,
+                                      const std::vector<std::string>&               CLASS_NAMES,
-		int seg_w = 160
+                                      const std::vector<std::vector<unsigned int>>& COLORS,
-	);
+                                      const std::vector<std::vector<unsigned int>>& MASK_COLORS);
-	static void draw_objects(
+    int                  num_bindings;
-		const cv::Mat& image,
+    int                  num_inputs  = 0;
-		cv::Mat& res,
+    int                  num_outputs = 0;
-		const std::vector<Object>& objs,
+    std::vector<Binding> input_bindings;
-		const std::vector<std::string>& CLASS_NAMES,
+    std::vector<Binding> output_bindings;
-		const std::vector<std::vector<unsigned int>>& COLORS,
+    std::vector<void*>   host_ptrs;
-		const std::vector<std::vector<unsigned int>>& MASK_COLORS
+    std::vector<void*>   device_ptrs;
-	);
+
-	int num_bindings;
+    PreParam pparam;
-	int num_inputs = 0;
-	int num_outputs = 0;
-	std::vector<Binding> input_bindings;
-	std::vector<Binding> output_bindings;
-	std::vector<void*> host_ptrs;
-	std::vector<void*> device_ptrs;
-
-	PreParam pparam;
-private:
-	nvinfer1::ICudaEngine* engine = nullptr;
-	nvinfer1::IRuntime* runtime = nullptr;
-	nvinfer1::IExecutionContext* context = nullptr;
-	cudaStream_t stream = nullptr;
-	Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
 YOLOv8_seg::YOLOv8_seg(const std::string& engine_file_path)
 {
-	std::ifstream file(engine_file_path, std::ios::binary);
+    std::ifstream file(engine_file_path, std::ios::binary);
-	assert(file.good());
+    assert(file.good());
-	file.seekg(0, std::ios::end);
+    file.seekg(0, std::ios::end);
-	auto size = file.tellg();
+    auto size = file.tellg();
-	file.seekg(0, std::ios::beg);
+    file.seekg(0, std::ios::beg);
-	char* trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
-	assert(trtModelStream);
+    assert(trtModelStream);
-	file.read(trtModelStream, size);
+    file.read(trtModelStream, size);
-	file.close();
+    file.close();
-	initLibNvInferPlugins(&this->gLogger, "");
+    initLibNvInferPlugins(&this->gLogger, "");
-	this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
-	assert(this->runtime != nullptr);
+    assert(this->runtime != nullptr);
-
+
-	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
-	assert(this->engine != nullptr);
+    assert(this->engine != nullptr);
-	delete[] trtModelStream;
+    delete[] trtModelStream;
-	this->context = this->engine->createExecutionContext();
+    this->context = this->engine->createExecutionContext();
-
+
-	assert(this->context != nullptr);
+    assert(this->context != nullptr);
-	cudaStreamCreate(&this->stream);
+    cudaStreamCreate(&this->stream);
-	this->num_bindings = this->engine->getNbBindings();
+    this->num_bindings = this->engine->getNbBindings();
-
+
-	for (int i = 0; i < this->num_bindings; ++i)
+    for (int i = 0; i < this->num_bindings; ++i) {
-	{
+        Binding            binding;
-		Binding binding;
+        nvinfer1::Dims     dims;
-		nvinfer1::Dims dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-		nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
-		std::string name = this->engine->getBindingName(i);
+        binding.name             = name;
-		binding.name = name;
+        binding.dsize            = type_to_size(dtype);
-		binding.dsize = type_to_size(dtype);
+
-
+        bool IsInput = engine->bindingIsInput(i);
-		bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
-		if (IsInput)
+            this->num_inputs += 1;
-		{
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-			this->num_inputs += 1;
+            binding.size = get_size_by_dims(dims);
-			dims = this->engine->getProfileDimensions(
+            binding.dims = dims;
-				i,
+            this->input_bindings.push_back(binding);
-				0,
+            // set max opt shape
-				nvinfer1::OptProfileSelector::kMAX);
+            this->context->setBindingDimensions(i, dims);
-			binding.size = get_size_by_dims(dims);
+        }
-			binding.dims = dims;
+        else {
-			this->input_bindings.push_back(binding);
+            dims         = this->context->getBindingDimensions(i);
-			// set max opt shape
+            binding.size = get_size_by_dims(dims);
-			this->context->setBindingDimensions(i, dims);
+            binding.dims = dims;
-
+            this->output_bindings.push_back(binding);
-		}
+            this->num_outputs += 1;
-		else
+        }
-		{
+    }
-			dims = this->context->getBindingDimensions(i);
-			binding.size = get_size_by_dims(dims);
-			binding.dims = dims;
-			this->output_bindings.push_back(binding);
-			this->num_outputs += 1;
-		}
-	}
-
 }
 
 YOLOv8_seg::~YOLOv8_seg()
 {
-	this->context->destroy();
+    this->context->destroy();
-	this->engine->destroy();
+    this->engine->destroy();
-	this->runtime->destroy();
+    this->runtime->destroy();
-	cudaStreamDestroy(this->stream);
+    cudaStreamDestroy(this->stream);
-	for (auto& ptr : this->device_ptrs)
+    for (auto& ptr : this->device_ptrs) {
-	{
+        CHECK(cudaFree(ptr));
-		CHECK(cudaFree(ptr));
+    }
-	}
+
-
+    for (auto& ptr : this->host_ptrs) {
-	for (auto& ptr : this->host_ptrs)
+        CHECK(cudaFreeHost(ptr));
-	{
+    }
-		CHECK(cudaFreeHost(ptr));
-	}
 }
 
 void YOLOv8_seg::make_pipe(bool warmup)
 {
 
-	for (auto& bindings : this->input_bindings)
+    for (auto& bindings : this->input_bindings) {
-	{
+        void* d_ptr;
-		void* d_ptr;
+        CHECK(cudaMalloc(&d_ptr, bindings.size * bindings.dsize));
-		CHECK(cudaMalloc(
+        this->device_ptrs.push_back(d_ptr);
-			&d_ptr,
+    }
-			bindings.size * bindings.dsize)
+
-		);
+    for (auto& bindings : this->output_bindings) {
-		this->device_ptrs.push_back(d_ptr);
+        void * d_ptr, *h_ptr;
-	}
+        size_t size = bindings.size * bindings.dsize;
-
+        CHECK(cudaMalloc(&d_ptr, size));
-	for (auto& bindings : this->output_bindings)
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-	{
+        this->device_ptrs.push_back(d_ptr);
-		void* d_ptr, * h_ptr;
+        this->host_ptrs.push_back(h_ptr);
-		size_t size = bindings.size * bindings.dsize;
+    }
-		CHECK(cudaMalloc(
+
-			&d_ptr,
+    if (warmup) {
-			size)
+        for (int i = 0; i < 10; i++) {
-		);
+            for (auto& bindings : this->input_bindings) {
-		CHECK(cudaHostAlloc(
+                size_t size  = bindings.size * bindings.dsize;
-			&h_ptr,
+                void*  h_ptr = malloc(size);
-			size,
+                memset(h_ptr, 0, size);
-			0)
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-		);
+                free(h_ptr);
-		this->device_ptrs.push_back(d_ptr);
+            }
-		this->host_ptrs.push_back(h_ptr);
+            this->infer();
-	}
+        }
-
+        printf("model warmup 10 times\n");
-	if (warmup)
+    }
-	{
-		for (int i = 0; i < 10; i++)
-		{
-			for (auto& bindings : this->input_bindings)
-			{
-				size_t size = bindings.size * bindings.dsize;
-				void* h_ptr = malloc(size);
-				memset(h_ptr, 0, size);
-				CHECK(cudaMemcpyAsync(
-					this->device_ptrs[0],
-					h_ptr,
-					size,
-					cudaMemcpyHostToDevice,
-					this->stream)
-				);
-				free(h_ptr);
-			}
-			this->infer();
-		}
-		printf("model warmup 10 times\n");
-
-	}
 }
 
-void YOLOv8_seg::letterbox(
+void YOLOv8_seg::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
-	const cv::Mat& image,
-	cv::Mat& out,
-	cv::Size& size
-)
 {
-	const float inp_h = size.height;
+    const float inp_h  = size.height;
-	const float inp_w = size.width;
+    const float inp_w  = size.width;
-	float height = image.rows;
+    float       height = image.rows;
-	float width = image.cols;
+    float       width  = image.cols;
-
+
-	float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-	int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-	int padh = std::round(height * r);
+    int   padh = std::round(height * r);
-
+
-	cv::Mat tmp;
+    cv::Mat tmp;
-	if ((int)width != padw || (int)height != padh)
+    if ((int)width != padw || (int)height != padh) {
-	{
+        cv::resize(image, tmp, cv::Size(padw, padh));
-		cv::resize(
+    }
-			image,
+    else {
-			tmp,
+        tmp = image.clone();
-			cv::Size(padw, padh)
+    }
-		);
+
-	}
+    float dw = inp_w - padw;
-	else
+    float dh = inp_h - padh;
-	{
+
-		tmp = image.clone();
+    dw /= 2.0f;
-	}
+    dh /= 2.0f;
-
+    int top    = int(std::round(dh - 0.1f));
-	float dw = inp_w - padw;
+    int bottom = int(std::round(dh + 0.1f));
-	float dh = inp_h - padh;
+    int left   = int(std::round(dw - 0.1f));
-
+    int right  = int(std::round(dw + 0.1f));
-	dw /= 2.0f;
+
-	dh /= 2.0f;
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-	int top = int(std::round(dh - 0.1f));
+
-	int bottom = int(std::round(dh + 0.1f));
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-	int left = int(std::round(dw - 0.1f));
+    this->pparam.ratio  = 1 / r;
-	int right = int(std::round(dw + 0.1f));
+    this->pparam.dw     = dw;
-
+    this->pparam.dh     = dh;
-	cv::copyMakeBorder(
+    this->pparam.height = height;
-		tmp,
+    this->pparam.width  = width;
-		tmp,
+    ;
-		top,
-		bottom,
-		left,
-		right,
-		cv::BORDER_CONSTANT,
-		{ 114, 114, 114 }
-	);
-
-	cv::dnn::blobFromImage(tmp,
-		out,
-		1 / 255.f,
-		cv::Size(),
-		cv::Scalar(0, 0, 0),
-		true,
-		false,
-		CV_32F
-	);
-	this->pparam.ratio = 1 / r;
-	this->pparam.dw = dw;
-	this->pparam.dh = dh;
-	this->pparam.height = height;
-	this->pparam.width = width;;
 }
 
 void YOLOv8_seg::copy_from_Mat(const cv::Mat& image)
 {
-	cv::Mat nchw;
+    cv::Mat  nchw;
-	auto& in_binding = this->input_bindings[0];
+    auto&    in_binding = this->input_bindings[0];
-	auto width = in_binding.dims.d[3];
+    auto     width      = in_binding.dims.d[3];
-	auto height = in_binding.dims.d[2];
+    auto     height     = in_binding.dims.d[2];
-	cv::Size size{ width, height };
+    cv::Size size{width, height};
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+
-		nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-		size
+
-	);
+    CHECK(cudaMemcpyAsync(
-
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{
-				4,
-				{ 1, 3, height, width }
-			}
-	);
-
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8_seg::copy_from_Mat(const cv::Mat& image, cv::Size& size)
 {
-	cv::Mat nchw;
+    cv::Mat nchw;
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-		nchw,
+    CHECK(cudaMemcpyAsync(
-		size
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	);
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{ 4,
-			  { 1, 3, size.height, size.width }
-			}
-	);
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8_seg::infer()
 {
 
-	this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-		this->device_ptrs.data(),
+    for (int i = 0; i < this->num_outputs; i++) {
-		this->stream,
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		nullptr
+        CHECK(cudaMemcpyAsync(
-	);
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-	for (int i = 0; i < this->num_outputs; i++)
+    }
-	{
+    cudaStreamSynchronize(this->stream);
-		size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		CHECK(cudaMemcpyAsync(this->host_ptrs[i],
-			this->device_ptrs[i + this->num_inputs],
-			osize,
-			cudaMemcpyDeviceToHost,
-			this->stream)
-		);
-
-	}
-	cudaStreamSynchronize(this->stream);
-
 }
 
-void YOLOv8_seg::postprocess(std::vector<Object>& objs,
+void YOLOv8_seg::postprocess(
-	float score_thres,
+    std::vector<Object>& objs, float score_thres, float iou_thres, int topk, int seg_channels, int seg_h, int seg_w)
-	float iou_thres,
-	int topk,
-	int seg_channels,
-	int seg_h,
-	int seg_w
-)
 {
-	objs.clear();
+    objs.clear();
-	auto input_h = this->input_bindings[0].dims.d[2];
+    auto input_h      = this->input_bindings[0].dims.d[2];
-	auto input_w = this->input_bindings[0].dims.d[3];
+    auto input_w      = this->input_bindings[0].dims.d[3];
-	auto num_anchors = this->output_bindings[0].dims.d[1];
+    auto num_anchors  = this->output_bindings[0].dims.d[1];
-	auto num_channels = this->output_bindings[0].dims.d[2];
+    auto num_channels = this->output_bindings[0].dims.d[2];
-
+
-	auto& dw = this->pparam.dw;
+    auto& dw     = this->pparam.dw;
-	auto& dh = this->pparam.dh;
+    auto& dh     = this->pparam.dh;
-	auto& width = this->pparam.width;
+    auto& width  = this->pparam.width;
-	auto& height = this->pparam.height;
+    auto& height = this->pparam.height;
-	auto& ratio = this->pparam.ratio;
+    auto& ratio  = this->pparam.ratio;
-
+
-	auto* output = static_cast<float*>(this->host_ptrs[0]);
+    auto*   output = static_cast<float*>(this->host_ptrs[0]);
-	cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F,
+    cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F, static_cast<float*>(this->host_ptrs[1]));
-		static_cast<float*>(this->host_ptrs[1]));
+
-
+    std::vector<int>      labels;
-	std::vector<int> labels;
+    std::vector<float>    scores;
-	std::vector<float> scores;
+    std::vector<cv::Rect> bboxes;
-	std::vector<cv::Rect> bboxes;
+    std::vector<cv::Mat>  mask_confs;
-	std::vector<cv::Mat> mask_confs;
+    std::vector<int>      indices;
-	std::vector<int> indices;
+
-
+    for (int i = 0; i < num_anchors; i++) {
-	for (int i = 0; i < num_anchors; i++)
+        float* ptr   = output + i * num_channels;
-	{
+        float  score = *(ptr + 4);
-		float* ptr = output + i * num_channels;
+        if (score > score_thres) {
-		float score = *(ptr + 4);
+            float x0 = *ptr++ - dw;
-		if (score > score_thres)
+            float y0 = *ptr++ - dh;
-		{
+            float x1 = *ptr++ - dw;
-			float x0 = *ptr++ - dw;
+            float y1 = *ptr++ - dh;
-			float y0 = *ptr++ - dh;
+
-			float x1 = *ptr++ - dw;
+            x0 = clamp(x0 * ratio, 0.f, width);
-			float y1 = *ptr++ - dh;
+            y0 = clamp(y0 * ratio, 0.f, height);
-
+            x1 = clamp(x1 * ratio, 0.f, width);
-			x0 = clamp(x0 * ratio, 0.f, width);
+            y1 = clamp(y1 * ratio, 0.f, height);
-			y0 = clamp(y0 * ratio, 0.f, height);
+
-			x1 = clamp(x1 * ratio, 0.f, width);
+            int     label     = *(++ptr);
-			y1 = clamp(y1 * ratio, 0.f, height);
+            cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, ++ptr);
-
+            mask_confs.push_back(mask_conf);
-			int label = *(++ptr);
+            labels.push_back(label);
-			cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, ++ptr);
+            scores.push_back(score);
-			mask_confs.push_back(mask_conf);
+            bboxes.push_back(cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0));
-			labels.push_back(label);
+        }
-			scores.push_back(score);
+    }
-			bboxes.push_back(cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0));
-
-		}
-	}
 
 #if defined(BATCHED_NMS)
-	cv::dnn::NMSBoxesBatched(
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		labels,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #else
-	cv::dnn::NMSBoxes(
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #endif
 
-	cv::Mat masks;
+    cv::Mat masks;
-	int cnt = 0;
+    int     cnt = 0;
-	for (auto& i : indices)
+    for (auto& i : indices) {
-	{
+        if (cnt >= topk) {
-		if (cnt >= topk)
+            break;
-		{
+        }
-			break;
+        cv::Rect tmp = bboxes[i];
-		}
+        Object   obj;
-		cv::Rect tmp = bboxes[i];
+        obj.label = labels[i];
-		Object obj;
+        obj.rect  = tmp;
-		obj.label = labels[i];
+        obj.prob  = scores[i];
-		obj.rect = tmp;
+        masks.push_back(mask_confs[i]);
-		obj.prob = scores[i];
+        objs.push_back(obj);
-		masks.push_back(mask_confs[i]);
+        cnt += 1;
-		objs.push_back(obj);
+    }
-		cnt += 1;
+    if (masks.empty()) {
-	}
+        // masks is empty
-	if(masks.empty())
+    }
-	{
+    else {
-		//masks is empty
+        cv::Mat matmulRes = (masks * protos).t();
-	}
+        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_w, seg_h});
-	else
+
-	{
+        std::vector<cv::Mat> maskChannels;
-		cv::Mat matmulRes = (masks * protos).t();
+        cv::split(maskMat, maskChannels);
-		cv::Mat maskMat = matmulRes.reshape(indices.size(), { seg_w, seg_h });
+        int scale_dw = dw / input_w * seg_w;
-
+        int scale_dh = dh / input_h * seg_h;
-		std::vector<cv::Mat> maskChannels;
+
-		cv::split(maskMat, maskChannels);
+        cv::Rect roi(scale_dw, scale_dh, seg_w - 2 * scale_dw, seg_h - 2 * scale_dh);
-		int scale_dw = dw / input_w * seg_w;
+
-		int scale_dh = dh / input_h * seg_h;
+        for (int i = 0; i < indices.size(); i++) {
-
+            cv::Mat dest, mask;
-		cv::Rect roi(
+            cv::exp(-maskChannels[i], dest);
-			scale_dw,
+            dest = 1.0 / (1.0 + dest);
-			scale_dh,
+            dest = dest(roi);
-			seg_w - 2 * scale_dw,
+            cv::resize(dest, mask, cv::Size((int)width, (int)height), cv::INTER_LINEAR);
-			seg_h - 2 * scale_dh);
+            objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-
+        }
-		for (int i = 0; i < indices.size(); i++)
+    }
-		{
-			cv::Mat dest, mask;
-			cv::exp(-maskChannels[i], dest);
-			dest = 1.0 / (1.0 + dest);
-			dest = dest(roi);
-			cv::resize(
-				dest,
-				mask,
-				cv::Size((int)width, (int)height),
-				cv::INTER_LINEAR
-			);
-			objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-		}
-	}
 }
 
-void YOLOv8_seg::draw_objects(const cv::Mat& image,
+void YOLOv8_seg::draw_objects(const cv::Mat&                                image,
-	cv::Mat& res,
+                              cv::Mat&                                      res,
-	const std::vector<Object>& objs,
+                              const std::vector<Object>&                    objs,
-	const std::vector<std::string>& CLASS_NAMES,
+                              const std::vector<std::string>&               CLASS_NAMES,
-	const std::vector<std::vector<unsigned int>>& COLORS,
+                              const std::vector<std::vector<unsigned int>>& COLORS,
-	const std::vector<std::vector<unsigned int>>& MASK_COLORS
+                              const std::vector<std::vector<unsigned int>>& MASK_COLORS)
-)
 {
-	res = image.clone();
+    res          = image.clone();
-	cv::Mat mask = image.clone();
+    cv::Mat mask = image.clone();
-	for (auto& obj : objs)
+    for (auto& obj : objs) {
-	{
+        int        idx   = obj.label;
-		int idx = obj.label;
+        cv::Scalar color = cv::Scalar(COLORS[idx][0], COLORS[idx][1], COLORS[idx][2]);
-		cv::Scalar color = cv::Scalar(
+        cv::Scalar mask_color =
-			COLORS[idx][0],
+            cv::Scalar(MASK_COLORS[idx % 20][0], MASK_COLORS[idx % 20][1], MASK_COLORS[idx % 20][2]);
-			COLORS[idx][1],
+        cv::rectangle(res, obj.rect, color, 2);
-			COLORS[idx][2]
+
-		);
+        char text[256];
-		cv::Scalar mask_color = cv::Scalar(
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[idx].c_str(), obj.prob * 100);
-			MASK_COLORS[idx % 20][0],
+        mask(obj.rect).setTo(mask_color, obj.boxMask);
-			MASK_COLORS[idx % 20][1],
+
-			MASK_COLORS[idx % 20][2]
+        int      baseLine   = 0;
-		);
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-		cv::rectangle(
+
-			res,
+        int x = (int)obj.rect.x;
-			obj.rect,
+        int y = (int)obj.rect.y + 1;
-			color,
+
-			2
+        if (y > res.rows)
-		);
+            y = res.rows;
-
+
-		char text[256];
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-		sprintf(
+
-			text,
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-			"%s %.1f%%",
+    }
-			CLASS_NAMES[idx].c_str(),
+    cv::addWeighted(res, 0.5, mask, 0.8, 1, res);
-			obj.prob * 100
-		);
-		mask(obj.rect).setTo(mask_color, obj.boxMask);
-
-		int baseLine = 0;
-		cv::Size label_size = cv::getTextSize(
-			text,
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			1,
-			&baseLine
-		);
-
-		int x = (int)obj.rect.x;
-		int y = (int)obj.rect.y + 1;
-
-		if (y > res.rows)
-			y = res.rows;
-
-		cv::rectangle(
-			res,
-			cv::Rect(x, y, label_size.width, label_size.height + baseLine),
-			{ 0, 0, 255 },
-			-1
-		);
-
-		cv::putText(
-			res,
-			text,
-			cv::Point(x, y + label_size.height),
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			{ 255, 255, 255 },
-			1
-		);
-	}
-	cv::addWeighted(
-		res,
-		0.5,
-		mask,
-		0.8,
-		1,
-		res
-	);
 }
-#endif //JETSON_SEGMENT_YOLOV8_SEG_HPP
+#endif  // JETSON_SEGMENT_YOLOV8_SEG_HPP

csrc/jetson/segment/main.cpp

@@ -2,177 +2,131 @@
 // Created by ubuntu on 3/16/23.
 //
 #include "chrono"
-#include "yolov8-seg.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8-seg.hpp"
 
 const std::vector<std::string> CLASS_NAMES = {
-	"person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",
-	"train", "truck", "boat", "traffic light", "fire hydrant",
+    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",    "parking meter", "bench",
-	"stop sign", "parking meter", "bench", "bird", "cat",
+    "bird",           "cat",        "dog",           "horse",         "sheep",        "cow",           "elephant",
-	"dog", "horse", "sheep", "cow", "elephant",
+    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",     "handbag",       "tie",
-	"bear", "zebra", "giraffe", "backpack", "umbrella",
+    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball",  "kite",          "baseball bat",
-	"handbag", "tie", "suitcase", "frisbee", "skis",
+    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",       "wine glass",    "cup",
-	"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "fork",           "knife",      "spoon",         "bowl",          "banana",       "apple",         "sandwich",
-	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
+    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",        "donut",         "cake",
-	"cup", "fork", "knife", "spoon", "bowl",
+    "chair",          "couch",      "potted plant",  "bed",           "dining table", "toilet",        "tv",
-	"banana", "apple", "sandwich", "orange", "broccoli",
+    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",   "microwave",     "oven",
-	"carrot", "hot dog", "pizza", "donut", "cake",
+    "toaster",        "sink",       "refrigerator",  "book",          "clock",        "vase",          "scissors",
-	"chair", "couch", "potted plant", "bed", "dining table",
+    "teddy bear",     "hair drier", "toothbrush"};
-	"toilet", "tv", "laptop", "mouse", "remote",
-	"keyboard", "cell phone", "microwave", "oven",
-	"toaster", "sink", "refrigerator", "book", "clock", "vase",
-	"scissors", "teddy bear", "hair drier", "toothbrush" };
 
 const std::vector<std::vector<unsigned int>> COLORS = {
-	{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
+    {0, 114, 189},   {217, 83, 25},   {237, 177, 32},  {126, 47, 142},  {119, 172, 48},  {77, 190, 238},
-	{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
+    {162, 20, 47},   {76, 76, 76},    {153, 153, 153}, {255, 0, 0},     {255, 128, 0},   {191, 191, 0},
-	{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
+    {0, 255, 0},     {0, 0, 255},     {170, 0, 255},   {85, 85, 0},     {85, 170, 0},    {85, 255, 0},
-	{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
+    {170, 85, 0},    {170, 170, 0},   {170, 255, 0},   {255, 85, 0},    {255, 170, 0},   {255, 255, 0},
-	{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
+    {0, 85, 128},    {0, 170, 128},   {0, 255, 128},   {85, 0, 128},    {85, 85, 128},   {85, 170, 128},
-	{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
+    {85, 255, 128},  {170, 0, 128},   {170, 85, 128},  {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
-	{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
+    {255, 85, 128},  {255, 170, 128}, {255, 255, 128}, {0, 85, 255},    {0, 170, 255},   {0, 255, 255},
-	{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
+    {85, 0, 255},    {85, 85, 255},   {85, 170, 255},  {85, 255, 255},  {170, 0, 255},   {170, 85, 255},
-	{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
+    {170, 170, 255}, {170, 255, 255}, {255, 0, 255},   {255, 85, 255},  {255, 170, 255}, {85, 0, 0},
-	{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
+    {128, 0, 0},     {170, 0, 0},     {212, 0, 0},     {255, 0, 0},     {0, 43, 0},      {0, 85, 0},
-	{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
+    {0, 128, 0},     {0, 170, 0},     {0, 212, 0},     {0, 255, 0},     {0, 0, 43},      {0, 0, 85},
-	{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
+    {0, 0, 128},     {0, 0, 170},     {0, 0, 212},     {0, 0, 255},     {0, 0, 0},       {36, 36, 36},
-	{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
+    {73, 73, 73},    {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
-	{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
+    {80, 183, 189},  {128, 128, 0}};
-	{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
-	{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
-	{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
-	{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
-	{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
-	{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
-	{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
-	{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
-	{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
-	{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
-	{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
-	{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
-	{ 80, 183, 189 }, { 128, 128, 0 }
-};
 
 const std::vector<std::vector<unsigned int>> MASK_COLORS = {
-	{ 255, 56, 56 }, { 255, 157, 151 }, { 255, 112, 31 },
+    {255, 56, 56},  {255, 157, 151}, {255, 112, 31}, {255, 178, 29}, {207, 210, 49},  {72, 249, 10}, {146, 204, 23},
-	{ 255, 178, 29 }, { 207, 210, 49 }, { 72, 249, 10 },
+    {61, 219, 134}, {26, 147, 52},   {0, 212, 187},  {44, 153, 168}, {0, 194, 255},   {52, 69, 147}, {100, 115, 255},
-	{ 146, 204, 23 }, { 61, 219, 134 }, { 26, 147, 52 },
+    {0, 24, 236},   {132, 56, 255},  {82, 0, 133},   {203, 56, 255}, {255, 149, 200}, {255, 55, 199}};
-	{ 0, 212, 187 }, { 44, 153, 168 }, { 0, 194, 255 },
-	{ 52, 69, 147 }, { 100, 115, 255 }, { 0, 24, 236 },
-	{ 132, 56, 255 }, { 82, 0, 133 }, { 203, 56, 255 },
-	{ 255, 149, 200 }, { 255, 55, 199 }
-};
 
 int main(int argc, char** argv)
 {
-	// cuda:0
+    // cuda:0
-	cudaSetDevice(0);
+    cudaSetDevice(0);
 
-	const std::string engine_file_path{ argv[1] };
+    const std::string engine_file_path{argv[1]};
-	const std::string path{ argv[2] };
+    const std::string path{argv[2]};
 
-	std::vector<std::string> imagePathList;
+    std::vector<std::string> imagePathList;
-	bool isVideo{ false };
+    bool                     isVideo{false};
 
-	assert(argc == 3);
+    assert(argc == 3);
 
-	auto yolov8 = new YOLOv8_seg(engine_file_path);
+    auto yolov8 = new YOLOv8_seg(engine_file_path);
-	yolov8->make_pipe(true);
+    yolov8->make_pipe(true);
 
-	if (IsFile(path))
+    if (IsFile(path)) {
-	{
+        std::string suffix = path.substr(path.find_last_of('.') + 1);
-		std::string suffix = path.substr(path.find_last_of('.') + 1);
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-		if (
+            imagePathList.push_back(path);
-			suffix == "jpg" ||
+        }
-				suffix == "jpeg" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-				suffix == "png"
+                 || suffix == "mkv") {
-			)
+            isVideo = true;
-		{
+        }
-			imagePathList.push_back(path);
+        else {
-		}
+            printf("suffix %s is wrong !!!\n", suffix.c_str());
-		else if (
+            std::abort();
-			suffix == "mp4" ||
+        }
-				suffix == "avi" ||
+    }
-				suffix == "m4v" ||
+    else if (IsFolder(path)) {
-				suffix == "mpeg" ||
+        cv::glob(path + "/*.jpg", imagePathList);
-				suffix == "mov" ||
+    }
-				suffix == "mkv"
-			)
-		{
-			isVideo = true;
-		}
-		else
-		{
-			printf("suffix %s is wrong !!!\n", suffix.c_str());
-			std::abort();
-		}
-	}
-	else if (IsFolder(path))
-	{
-		cv::glob(path + "/*.jpg", imagePathList);
-	}
 
-	cv::Mat res, image;
+    cv::Mat  res, image;
-	cv::Size size = cv::Size{ 640, 640 };
+    cv::Size size         = cv::Size{640, 640};
-	int topk = 100;
+    int      topk         = 100;
-	int seg_h = 160;
+    int      seg_h        = 160;
-	int seg_w = 160;
+    int      seg_w        = 160;
-	int seg_channels = 32;
+    int      seg_channels = 32;
-	float score_thres = 0.25f;
+    float    score_thres  = 0.25f;
-	float iou_thres = 0.65f;
+    float    iou_thres    = 0.65f;
 
-	std::vector<Object> objs;
+    std::vector<Object> objs;
 
-	cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
+    cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
 
-	if (isVideo)
+    if (isVideo) {
-	{
+        cv::VideoCapture cap(path);
-		cv::VideoCapture cap(path);
 
-		if (!cap.isOpened())
+        if (!cap.isOpened()) {
-		{
+            printf("can not open %s\n", path.c_str());
-			printf("can not open %s\n", path.c_str());
+            return -1;
-			return -1;
+        }
-		}
+        while (cap.read(image)) {
-		while (cap.read(image))
+            objs.clear();
-		{
+            yolov8->copy_from_Mat(image, size);
-			objs.clear();
+            auto start = std::chrono::system_clock::now();
-			yolov8->copy_from_Mat(image, size);
+            yolov8->infer();
-			auto start = std::chrono::system_clock::now();
+            auto end = std::chrono::system_clock::now();
-			yolov8->infer();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
-			auto end = std::chrono::system_clock::now();
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
+            printf("cost %2.4lf ms\n", tc);
-			auto tc = (double)
+            cv::imshow("result", res);
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+            if (cv::waitKey(10) == 'q') {
-			printf("cost %2.4lf ms\n", tc);
+                break;
-			cv::imshow("result", res);
+            }
-			if (cv::waitKey(10) == 'q')
+        }
-			{
+    }
-				break;
+    else {
-			}
+        for (auto& path : imagePathList) {
-		}
+            objs.clear();
-	}
+            image = cv::imread(path);
-	else
+            yolov8->copy_from_Mat(image, size);
-	{
+            auto start = std::chrono::system_clock::now();
-		for (auto& path : imagePathList)
+            yolov8->infer();
-		{
+            auto end = std::chrono::system_clock::now();
-			objs.clear();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
-			image = cv::imread(path);
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
-			yolov8->copy_from_Mat(image, size);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			auto start = std::chrono::system_clock::now();
+            printf("cost %2.4lf ms\n", tc);
-			yolov8->infer();
+            cv::imshow("result", res);
-			auto end = std::chrono::system_clock::now();
+            cv::waitKey(0);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
+        }
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
+    }
-			auto tc = (double)
+    cv::destroyAllWindows();
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+    delete yolov8;
-			printf("cost %2.4lf ms\n", tc);
+    return 0;
-			cv::imshow("result", res);
-			cv::waitKey(0);
-		}
-	}
-	cv::destroyAllWindows();
-	delete yolov8;
-	return 0;
 }

csrc/pose/normal/include/common.hpp

@@ -4,154 +4,139 @@
 
 #ifndef POSE_NORMAL_COMMON_HPP
 #define POSE_NORMAL_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace pose
+namespace pose {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float>   rect;
-		cv::Rect_<float> rect;
+    int                label = 0;
-		int label = 0;
+    float              prob  = 0.0;
-		float prob = 0.0;
+    std::vector<float> kps;
-        std::vector<float> kps;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace pose
-}
+#endif  // POSE_NORMAL_COMMON_HPP
-#endif //POSE_NORMAL_COMMON_HPP

csrc/pose/normal/include/yolov8-pose.hpp

@@ -4,73 +4,63 @@
 #ifndef POSE_NORMAL_YOLOv8_pose_HPP
 #define POSE_NORMAL_YOLOv8_pose_HPP
 
-#include "fstream"
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include "fstream"
 
 using namespace pose;
 
 class YOLOv8_pose {
 public:
-    explicit YOLOv8_pose(const std::string &engine_file_path);
+    explicit YOLOv8_pose(const std::string& engine_file_path);
 
     ~YOLOv8_pose();
 
     void make_pipe(bool warmup = true);
 
-    void copy_from_Mat(const cv::Mat &image);
+    void copy_from_Mat(const cv::Mat& image);
 
-    void copy_from_Mat(const cv::Mat &image, cv::Size &size);
+    void copy_from_Mat(const cv::Mat& image, cv::Size& size);
 
-    void letterbox(
+    void letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-            const cv::Mat &image,
-            cv::Mat &out,
-            cv::Size &size
-    );
 
     void infer();
 
-    void postprocess(
+    void postprocess(std::vector<Object>& objs, float score_thres = 0.25f, float iou_thres = 0.65f, int topk = 100);
-            std::vector<Object> &objs,
+
-            float score_thres = 0.25f,
+    static void draw_objects(const cv::Mat&                                image,
-            float iou_thres = 0.65f,
+                             cv::Mat&                                      res,
-            int topk = 100
+                             const std::vector<Object>&                    objs,
-    );
+                             const std::vector<std::vector<unsigned int>>& SKELETON,
-
+                             const std::vector<std::vector<unsigned int>>& KPS_COLORS,
-    static void draw_objects(
+                             const std::vector<std::vector<unsigned int>>& LIMB_COLORS);
-            const cv::Mat &image,
+
-            cv::Mat &res,
+    int                  num_bindings;
-            const std::vector<Object> &objs,
+    int                  num_inputs  = 0;
-            const std::vector<std::vector<unsigned int>> &SKELETON,
+    int                  num_outputs = 0;
-            const std::vector<std::vector<unsigned int>> &KPS_COLORS,
-            const std::vector<std::vector<unsigned int>> &LIMB_COLORS
-    );
-
-    int num_bindings;
-    int num_inputs = 0;
-    int num_outputs = 0;
     std::vector<Binding> input_bindings;
     std::vector<Binding> output_bindings;
-    std::vector<void *> host_ptrs;
+    std::vector<void*>   host_ptrs;
-    std::vector<void *> device_ptrs;
+    std::vector<void*>   device_ptrs;
 
     PreParam pparam;
-private:
-    nvinfer1::ICudaEngine *engine = nullptr;
-    nvinfer1::IRuntime *runtime = nullptr;
-    nvinfer1::IExecutionContext *context = nullptr;
-    cudaStream_t stream = nullptr;
-    Logger gLogger{nvinfer1::ILogger::Severity::kERROR};
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
-YOLOv8_pose::YOLOv8_pose(const std::string &engine_file_path) {
+YOLOv8_pose::YOLOv8_pose(const std::string& engine_file_path)
+{
     std::ifstream file(engine_file_path, std::ios::binary);
     assert(file.good());
     file.seekg(0, std::ios::end);
     auto size = file.tellg();
     file.seekg(0, std::ios::beg);
-    char *trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
     assert(trtModelStream);
     file.read(trtModelStream, size);
     file.close();
@@ -88,125 +78,97 @@ YOLOv8_pose::YOLOv8_pose(const std::string &engine_file_path) {
     this->num_bindings = this->engine->getNbBindings();
 
     for (int i = 0; i < this->num_bindings; ++i) {
-        Binding binding;
+        Binding            binding;
-        nvinfer1::Dims dims;
+        nvinfer1::Dims     dims;
         nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-        std::string name = this->engine->getBindingName(i);
+        std::string        name  = this->engine->getBindingName(i);
-        binding.name = name;
+        binding.name             = name;
-        binding.dsize = type_to_size(dtype);
+        binding.dsize            = type_to_size(dtype);
 
         bool IsInput = engine->bindingIsInput(i);
         if (IsInput) {
             this->num_inputs += 1;
-            dims = this->engine->getProfileDimensions(
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-                    i,
-                    0,
-                    nvinfer1::OptProfileSelector::kMAX);
             binding.size = get_size_by_dims(dims);
             binding.dims = dims;
             this->input_bindings.push_back(binding);
             // set max opt shape
             this->context->setBindingDimensions(i, dims);
-
+        }
-        } else {
+        else {
-            dims = this->context->getBindingDimensions(i);
+            dims         = this->context->getBindingDimensions(i);
             binding.size = get_size_by_dims(dims);
             binding.dims = dims;
             this->output_bindings.push_back(binding);
             this->num_outputs += 1;
         }
     }
-
 }
 
-YOLOv8_pose::~YOLOv8_pose() {
+YOLOv8_pose::~YOLOv8_pose()
+{
     this->context->destroy();
     this->engine->destroy();
     this->runtime->destroy();
     cudaStreamDestroy(this->stream);
-    for (auto &ptr: this->device_ptrs) {
+    for (auto& ptr : this->device_ptrs) {
         CHECK(cudaFree(ptr));
     }
 
-    for (auto &ptr: this->host_ptrs) {
+    for (auto& ptr : this->host_ptrs) {
         CHECK(cudaFreeHost(ptr));
     }
-
 }
 
-void YOLOv8_pose::make_pipe(bool warmup) {
+void YOLOv8_pose::make_pipe(bool warmup)
+{
 
-    for (auto &bindings: this->input_bindings) {
+    for (auto& bindings : this->input_bindings) {
-        void *d_ptr;
+        void* d_ptr;
-        CHECK(cudaMallocAsync(
+        CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream));
-                &d_ptr,
-                bindings.size * bindings.dsize,
-                this->stream)
-        );
         this->device_ptrs.push_back(d_ptr);
     }
 
-    for (auto &bindings: this->output_bindings) {
+    for (auto& bindings : this->output_bindings) {
-        void *d_ptr, *h_ptr;
+        void * d_ptr, *h_ptr;
         size_t size = bindings.size * bindings.dsize;
-        CHECK(cudaMallocAsync(
+        CHECK(cudaMallocAsync(&d_ptr, size, this->stream));
-                &d_ptr,
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-                size,
-                this->stream)
-        );
-        CHECK(cudaHostAlloc(
-                &h_ptr,
-                size,
-                0)
-        );
         this->device_ptrs.push_back(d_ptr);
         this->host_ptrs.push_back(h_ptr);
     }
 
     if (warmup) {
         for (int i = 0; i < 10; i++) {
-            for (auto &bindings: this->input_bindings) {
+            for (auto& bindings : this->input_bindings) {
-                size_t size = bindings.size * bindings.dsize;
+                size_t size  = bindings.size * bindings.dsize;
-                void *h_ptr = malloc(size);
+                void*  h_ptr = malloc(size);
                 memset(h_ptr, 0, size);
-                CHECK(cudaMemcpyAsync(
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-                        this->device_ptrs[0],
-                        h_ptr,
-                        size,
-                        cudaMemcpyHostToDevice,
-                        this->stream)
-                );
                 free(h_ptr);
             }
             this->infer();
         }
         printf("model warmup 10 times\n");
-
     }
 }
 
-void YOLOv8_pose::letterbox(
+void YOLOv8_pose::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
-        const cv::Mat &image,
+{
-        cv::Mat &out,
+    const float inp_h  = size.height;
-        cv::Size &size
+    const float inp_w  = size.width;
-) {
+    float       height = image.rows;
-    const float inp_h = size.height;
+    float       width  = image.cols;
-    const float inp_w = size.width;
-    float height = image.rows;
-    float width = image.cols;
 
-    float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-    int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-    int padh = std::round(height * r);
+    int   padh = std::round(height * r);
 
     cv::Mat tmp;
-    if ((int) width != padw || (int) height != padh) {
+    if ((int)width != padw || (int)height != padh) {
-        cv::resize(
+        cv::resize(image, tmp, cv::Size(padw, padh));
-                image,
+    }
-                tmp,
+    else {
-                cv::Size(padw, padh)
-        );
-    } else {
         tmp = image.clone();
     }
 
@@ -215,146 +177,83 @@ void YOLOv8_pose::letterbox(
 
     dw /= 2.0f;
     dh /= 2.0f;
-    int top = int(std::round(dh - 0.1f));
+    int top    = int(std::round(dh - 0.1f));
     int bottom = int(std::round(dh + 0.1f));
-    int left = int(std::round(dw - 0.1f));
+    int left   = int(std::round(dw - 0.1f));
-    int right = int(std::round(dw + 0.1f));
+    int right  = int(std::round(dw + 0.1f));
-
+
-    cv::copyMakeBorder(
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-            tmp,
+
-            tmp,
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-            top,
+    this->pparam.ratio  = 1 / r;
-            bottom,
+    this->pparam.dw     = dw;
-            left,
+    this->pparam.dh     = dh;
-            right,
-            cv::BORDER_CONSTANT,
-            {114, 114, 114}
-    );
-
-    cv::dnn::blobFromImage(tmp,
-                           out,
-                           1 / 255.f,
-                           cv::Size(),
-                           cv::Scalar(0, 0, 0),
-                           true,
-                           false,
-                           CV_32F
-    );
-    this->pparam.ratio = 1 / r;
-    this->pparam.dw = dw;
-    this->pparam.dh = dh;
     this->pparam.height = height;
-    this->pparam.width = width;;
+    this->pparam.width  = width;
+    ;
 }
 
-void YOLOv8_pose::copy_from_Mat(const cv::Mat &image) {
+void YOLOv8_pose::copy_from_Mat(const cv::Mat& image)
-    cv::Mat nchw;
+{
-    auto &in_binding = this->input_bindings[0];
+    cv::Mat  nchw;
-    auto width = in_binding.dims.d[3];
+    auto&    in_binding = this->input_bindings[0];
-    auto height = in_binding.dims.d[2];
+    auto     width      = in_binding.dims.d[3];
+    auto     height     = in_binding.dims.d[2];
     cv::Size size{width, height};
-    this->letterbox(
+    this->letterbox(image, nchw, size);
-            image,
+
-            nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-            size
-    );
-
-    this->context->setBindingDimensions(
-            0,
-            nvinfer1::Dims
-                    {
-                            4,
-                            {1, 3, height, width}
-                    }
-    );
 
     CHECK(cudaMemcpyAsync(
-            this->device_ptrs[0],
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-            nchw.ptr<float>(),
-            nchw.total() * nchw.elemSize(),
-            cudaMemcpyHostToDevice,
-            this->stream)
-    );
 }
 
-void YOLOv8_pose::copy_from_Mat(const cv::Mat &image, cv::Size &size) {
+void YOLOv8_pose::copy_from_Mat(const cv::Mat& image, cv::Size& size)
+{
     cv::Mat nchw;
-    this->letterbox(
+    this->letterbox(image, nchw, size);
-            image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-            nchw,
-            size
-    );
-    this->context->setBindingDimensions(
-            0,
-            nvinfer1::Dims
-                    {4,
-                     {1, 3, size.height, size.width}
-                    }
-    );
     CHECK(cudaMemcpyAsync(
-            this->device_ptrs[0],
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-            nchw.ptr<float>(),
-            nchw.total() * nchw.elemSize(),
-            cudaMemcpyHostToDevice,
-            this->stream)
-    );
 }
 
-void YOLOv8_pose::infer() {
+void YOLOv8_pose::infer()
+{
 
-    this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-            this->device_ptrs.data(),
-            this->stream,
-            nullptr
-    );
     for (int i = 0; i < this->num_outputs; i++) {
         size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-        CHECK(cudaMemcpyAsync(this->host_ptrs[i],
+        CHECK(cudaMemcpyAsync(
-                              this->device_ptrs[i + this->num_inputs],
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-                              osize,
-                              cudaMemcpyDeviceToHost,
-                              this->stream)
-        );
-
     }
     cudaStreamSynchronize(this->stream);
-
 }
 
-void YOLOv8_pose::postprocess(
+void YOLOv8_pose::postprocess(std::vector<Object>& objs, float score_thres, float iou_thres, int topk)
-        std::vector<Object> &objs,
+{
-        float score_thres,
-        float iou_thres,
-        int topk
-) {
     objs.clear();
     auto num_channels = this->output_bindings[0].dims.d[1];
-    auto num_anchors = this->output_bindings[0].dims.d[2];
+    auto num_anchors  = this->output_bindings[0].dims.d[2];
-
+
-    auto &dw = this->pparam.dw;
+    auto& dw     = this->pparam.dw;
-    auto &dh = this->pparam.dh;
+    auto& dh     = this->pparam.dh;
-    auto &width = this->pparam.width;
+    auto& width  = this->pparam.width;
-    auto &height = this->pparam.height;
+    auto& height = this->pparam.height;
-    auto &ratio = this->pparam.ratio;
+    auto& ratio  = this->pparam.ratio;
-
+
-    std::vector<cv::Rect> bboxes;
+    std::vector<cv::Rect>           bboxes;
-    std::vector<float> scores;
+    std::vector<float>              scores;
-    std::vector<int> labels;
+    std::vector<int>                labels;
-    std::vector<int> indices;
+    std::vector<int>                indices;
     std::vector<std::vector<float>> kpss;
 
-    cv::Mat output = cv::Mat(
+    cv::Mat output = cv::Mat(num_channels, num_anchors, CV_32F, static_cast<float*>(this->host_ptrs[0]));
-            num_channels,
+    output         = output.t();
-            num_anchors,
-            CV_32F,
-            static_cast<float *>(this->host_ptrs[0])
-    );
-    output = output.t();
     for (int i = 0; i < num_anchors; i++) {
-        auto row_ptr = output.row(i).ptr<float>();
+        auto row_ptr    = output.row(i).ptr<float>();
         auto bboxes_ptr = row_ptr;
         auto scores_ptr = row_ptr + 4;
-        auto kps_ptr = row_ptr + 5;
+        auto kps_ptr    = row_ptr + 5;
 
         float score = *scores_ptr;
         if (score > score_thres) {
@@ -369,17 +268,17 @@ void YOLOv8_pose::postprocess(
             float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);
 
             cv::Rect_<float> bbox;
-            bbox.x = x0;
+            bbox.x      = x0;
-            bbox.y = y0;
+            bbox.y      = y0;
-            bbox.width = x1 - x0;
+            bbox.width  = x1 - x0;
             bbox.height = y1 - y0;
             std::vector<float> kps;
             for (int k = 0; k < 17; k++) {
                 float kps_x = (*(kps_ptr + 3 * k) - dw) * ratio;
                 float kps_y = (*(kps_ptr + 3 * k + 1) - dh) * ratio;
                 float kps_s = *(kps_ptr + 3 * k + 2);
-                kps_x = clamp(kps_x, 0.f, width);
+                kps_x       = clamp(kps_x, 0.f, width);
-                kps_y = clamp(kps_y, 0.f, height);
+                kps_y       = clamp(kps_y, 0.f, height);
                 kps.push_back(kps_x);
                 kps.push_back(kps_y);
                 kps.push_back(kps_s);
@@ -393,110 +292,68 @@ void YOLOv8_pose::postprocess(
     }
 
 #ifdef BATCHED_NMS
-    cv::dnn::NMSBoxesBatched(
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
-            bboxes,
-            scores,
-            labels,
-            score_thres,
-            iou_thres,
-            indices
-    );
 #else
-    cv::dnn::NMSBoxes(
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
-        bboxes,
-        scores,
-        score_thres,
-        iou_thres,
-        indices
-    );
 #endif
 
     int cnt = 0;
-    for (auto &i: indices) {
+    for (auto& i : indices) {
         if (cnt >= topk) {
             break;
         }
         Object obj;
-        obj.rect = bboxes[i];
+        obj.rect  = bboxes[i];
-        obj.prob = scores[i];
+        obj.prob  = scores[i];
         obj.label = labels[i];
-        obj.kps = kpss[i];
+        obj.kps   = kpss[i];
         objs.push_back(obj);
         cnt += 1;
     }
 }
 
-void YOLOv8_pose::draw_objects(
+void YOLOv8_pose::draw_objects(const cv::Mat&                                image,
-        const cv::Mat &image,
+                               cv::Mat&                                      res,
-        cv::Mat &res,
+                               const std::vector<Object>&                    objs,
-        const std::vector<Object> &objs,
+                               const std::vector<std::vector<unsigned int>>& SKELETON,
-        const std::vector<std::vector<unsigned int>> &SKELETON,
+                               const std::vector<std::vector<unsigned int>>& KPS_COLORS,
-        const std::vector<std::vector<unsigned int>> &KPS_COLORS,
+                               const std::vector<std::vector<unsigned int>>& LIMB_COLORS)
-        const std::vector<std::vector<unsigned int>> &LIMB_COLORS
+{
-) {
+    res                 = image.clone();
-    res = image.clone();
     const int num_point = 17;
-    for (auto &obj: objs) {
+    for (auto& obj : objs) {
-        cv::rectangle(
+        cv::rectangle(res, obj.rect, {0, 0, 255}, 2);
-                res,
-                obj.rect,
-                {0, 0, 255},
-                2
-        );
 
         char text[256];
-        sprintf(
+        sprintf(text, "person %.1f%%", obj.prob * 100);
-                text,
+
-                "person %.1f%%",
+        int      baseLine   = 0;
-                obj.prob * 100
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-        );
+
-
+        int x = (int)obj.rect.x;
-        int baseLine = 0;
+        int y = (int)obj.rect.y + 1;
-        cv::Size label_size = cv::getTextSize(
-                text,
-                cv::FONT_HERSHEY_SIMPLEX,
-                0.4,
-                1,
-                &baseLine
-        );
-
-        int x = (int) obj.rect.x;
-        int y = (int) obj.rect.y + 1;
 
         if (y > res.rows)
             y = res.rows;
 
-        cv::rectangle(
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-                res,
+
-                cv::Rect(x, y, label_size.width, label_size.height + baseLine),
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-                {0, 0, 255},
+
-                -1
+        auto& kps = obj.kps;
-        );
-
-        cv::putText(
-                res,
-                text,
-                cv::Point(x, y + label_size.height),
-                cv::FONT_HERSHEY_SIMPLEX,
-                0.4,
-                {255, 255, 255},
-                1
-        );
-
-        auto &kps = obj.kps;
         for (int k = 0; k < num_point + 2; k++) {
             if (k < num_point) {
-                int kps_x = std::round(kps[k * 3]);
+                int   kps_x = std::round(kps[k * 3]);
-                int kps_y = std::round(kps[k * 3 + 1]);
+                int   kps_y = std::round(kps[k * 3 + 1]);
                 float kps_s = kps[k * 3 + 2];
                 if (kps_s > 0.5f) {
                     cv::Scalar kps_color = cv::Scalar(KPS_COLORS[k][0], KPS_COLORS[k][1], KPS_COLORS[k][2]);
                     cv::circle(res, {kps_x, kps_y}, 5, kps_color, -1);
                 }
             }
-            auto &ske = SKELETON[k];
+            auto& ske    = SKELETON[k];
-            int pos1_x = std::round(kps[(ske[0] - 1) * 3]);
+            int   pos1_x = std::round(kps[(ske[0] - 1) * 3]);
-            int pos1_y = std::round(kps[(ske[0] - 1) * 3 + 1]);
+            int   pos1_y = std::round(kps[(ske[0] - 1) * 3 + 1]);
 
             int pos2_x = std::round(kps[(ske[1] - 1) * 3]);
             int pos2_y = std::round(kps[(ske[1] - 1) * 3 + 1]);
@@ -504,7 +361,6 @@ void YOLOv8_pose::draw_objects(
             float pos1_s = kps[(ske[0] - 1) * 3 + 2];
             float pos2_s = kps[(ske[1] - 1) * 3 + 2];
 
-
             if (pos1_s > 0.5f && pos2_s > 0.5f) {
                 cv::Scalar limb_color = cv::Scalar(LIMB_COLORS[k][0], LIMB_COLORS[k][1], LIMB_COLORS[k][2]);
                 cv::line(res, {pos1_x, pos1_y}, {pos2_x, pos2_y}, limb_color, 2);
@@ -513,4 +369,4 @@ void YOLOv8_pose::draw_objects(
     }
 }
 
-#endif //POSE_NORMAL_YOLOv8_pose_HPP
+#endif  // POSE_NORMAL_YOLOv8_pose_HPP

csrc/pose/normal/main.cpp

@@ -2,70 +2,69 @@
 // Created by ubuntu on 4/7/23.
 //
 #include "chrono"
-#include "yolov8-pose.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8-pose.hpp"
 
-
+const std::vector<std::vector<unsigned int>> KPS_COLORS = {{0, 255, 0},
-const std::vector<std::vector<unsigned int>> KPS_COLORS =
+                                                           {0, 255, 0},
-        {{0,   255, 0},
+                                                           {0, 255, 0},
-         {0,   255, 0},
+                                                           {0, 255, 0},
-         {0,   255, 0},
+                                                           {0, 255, 0},
-         {0,   255, 0},
+                                                           {255, 128, 0},
-         {0,   255, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {255, 128, 0},
-         {255, 128, 0},
+                                                           {51, 153, 255},
-         {255, 128, 0},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255},
-         {51,  153, 255},
+                                                           {51, 153, 255}};
-         {51,  153, 255},
-         {51,  153, 255}};
 
 const std::vector<std::vector<unsigned int>> SKELETON = {{16, 14},
                                                          {14, 12},
                                                          {17, 15},
                                                          {15, 13},
                                                          {12, 13},
-                                                         {6,  12},
+                                                         {6, 12},
-                                                         {7,  13},
+                                                         {7, 13},
-                                                         {6,  7},
+                                                         {6, 7},
-                                                         {6,  8},
+                                                         {6, 8},
-                                                         {7,  9},
+                                                         {7, 9},
-                                                         {8,  10},
+                                                         {8, 10},
-                                                         {9,  11},
+                                                         {9, 11},
-                                                         {2,  3},
+                                                         {2, 3},
-                                                         {1,  2},
+                                                         {1, 2},
-                                                         {1,  3},
+                                                         {1, 3},
-                                                         {2,  4},
+                                                         {2, 4},
-                                                         {3,  5},
+                                                         {3, 5},
-                                                         {4,  6},
+                                                         {4, 6},
-                                                         {5,  7}};
+                                                         {5, 7}};
 
-const std::vector<std::vector<unsigned int>> LIMB_COLORS = {{51,  153, 255},
+const std::vector<std::vector<unsigned int>> LIMB_COLORS = {{51, 153, 255},
-                                                            {51,  153, 255},
+                                                            {51, 153, 255},
-                                                            {51,  153, 255},
+                                                            {51, 153, 255},
-                                                            {51,  153, 255},
+                                                            {51, 153, 255},
-                                                            {255, 51,  255},
+                                                            {255, 51, 255},
-                                                            {255, 51,  255},
+                                                            {255, 51, 255},
-                                                            {255, 51,  255},
+                                                            {255, 51, 255},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
                                                             {255, 128, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0},
+                                                            {0, 255, 0},
-                                                            {0,   255, 0}};
+                                                            {0, 255, 0}};
 
-int main(int argc, char **argv) {
+int main(int argc, char** argv)
+{
     // cuda:0
     cudaSetDevice(0);
 
@@ -73,7 +72,7 @@ int main(int argc, char **argv) {
     const std::string path{argv[2]};
 
     std::vector<std::string> imagePathList;
-    bool isVideo{false};
+    bool                     isVideo{false};
 
     assert(argc == 3);
 
@@ -82,34 +81,27 @@ int main(int argc, char **argv) {
 
     if (IsFile(path)) {
         std::string suffix = path.substr(path.find_last_of('.') + 1);
-        if (
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-                suffix == "jpg" ||
-                suffix == "jpeg" ||
-                suffix == "png"
-                ) {
             imagePathList.push_back(path);
-        } else if (
+        }
-                suffix == "mp4" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-                suffix == "avi" ||
+                 || suffix == "mkv") {
-                suffix == "m4v" ||
-                suffix == "mpeg" ||
-                suffix == "mov" ||
-                suffix == "mkv"
-                ) {
             isVideo = true;
-        } else {
+        }
+        else {
             printf("suffix %s is wrong !!!\n", suffix.c_str());
             std::abort();
         }
-    } else if (IsFolder(path)) {
+    }
+    else if (IsFolder(path)) {
         cv::glob(path + "/*.jpg", imagePathList);
     }
 
-    cv::Mat res, image;
+    cv::Mat  res, image;
-    cv::Size size = cv::Size{640, 640};
+    cv::Size size        = cv::Size{640, 640};
-    int topk = 100;
+    int      topk        = 100;
-    float score_thres = 0.25f;
+    float    score_thres = 0.25f;
-    float iou_thres = 0.65f;
+    float    iou_thres   = 0.65f;
 
     std::vector<Object> objs;
 
@@ -130,16 +122,16 @@ int main(int argc, char **argv) {
             auto end = std::chrono::system_clock::now();
             yolov8_pose->postprocess(objs, score_thres, iou_thres, topk);
             yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS);
-            auto tc = (double)
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-                              std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
             printf("cost %2.4lf ms\n", tc);
             cv::imshow("result", res);
             if (cv::waitKey(10) == 'q') {
                 break;
             }
         }
-    } else {
+    }
-        for (auto &path: imagePathList) {
+    else {
+        for (auto& path : imagePathList) {
             objs.clear();
             image = cv::imread(path);
             yolov8_pose->copy_from_Mat(image, size);
@@ -148,8 +140,7 @@ int main(int argc, char **argv) {
             auto end = std::chrono::system_clock::now();
             yolov8_pose->postprocess(objs, score_thres, iou_thres, topk);
             yolov8_pose->draw_objects(image, res, objs, SKELETON, KPS_COLORS, LIMB_COLORS);
-            auto tc = (double)
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-                              std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
             printf("cost %2.4lf ms\n", tc);
             cv::imshow("result", res);
             cv::waitKey(0);

csrc/segment/normal/include/common.hpp

@@ -4,154 +4,139 @@
 
 #ifndef SEGMENT_NORMAL_COMMON_HPP
 #define SEGMENT_NORMAL_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace seg
+namespace seg {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float> rect;
-		cv::Rect_<float> rect;
+    int              label = 0;
-		int label = 0;
+    float            prob  = 0.0;
-		float prob = 0.0;
+    cv::Mat          boxMask;
-		cv::Mat boxMask;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace seg
-}
+#endif  // SEGMENT_NORMAL_COMMON_HPP
-#endif //SEGMENT_NORMAL_COMMON_HPP

csrc/segment/normal/include/yolov8-seg.hpp

@@ -3,578 +3,382 @@
 //
 #ifndef SEGMENT_NORMAL_YOLOV8_SEG_HPP
 #define SEGMENT_NORMAL_YOLOV8_SEG_HPP
-#include <fstream>
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include <fstream>
 
 using namespace seg;
 
-class YOLOv8_seg
+class YOLOv8_seg {
-{
 public:
-	explicit YOLOv8_seg(const std::string& engine_file_path);
+    explicit YOLOv8_seg(const std::string& engine_file_path);
-	~YOLOv8_seg();
+    ~YOLOv8_seg();
-
+
-	void make_pipe(bool warmup = true);
+    void                 make_pipe(bool warmup = true);
-	void copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image);
-	void copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
-	void letterbox(
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-		const cv::Mat& image,
+    void                 infer();
-		cv::Mat& out,
+    void                 postprocess(std::vector<Object>& objs,
-		cv::Size& size
+                                     float                score_thres  = 0.25f,
-	);
+                                     float                iou_thres    = 0.65f,
-	void infer();
+                                     int                  topk         = 100,
-	void postprocess(
+                                     int                  seg_channels = 32,
-		std::vector<Object>& objs,
+                                     int                  seg_h        = 160,
-		float score_thres = 0.25f,
+                                     int                  seg_w        = 160);
-		float iou_thres = 0.65f,
+    static void          draw_objects(const cv::Mat&                                image,
-		int topk = 100,
+                                      cv::Mat&                                      res,
-		int seg_channels = 32,
+                                      const std::vector<Object>&                    objs,
-		int seg_h = 160,
+                                      const std::vector<std::string>&               CLASS_NAMES,
-		int seg_w = 160
+                                      const std::vector<std::vector<unsigned int>>& COLORS,
-	);
+                                      const std::vector<std::vector<unsigned int>>& MASK_COLORS);
-	static void draw_objects(
+    int                  num_bindings;
-		const cv::Mat& image,
+    int                  num_inputs  = 0;
-		cv::Mat& res,
+    int                  num_outputs = 0;
-		const std::vector<Object>& objs,
+    std::vector<Binding> input_bindings;
-		const std::vector<std::string>& CLASS_NAMES,
+    std::vector<Binding> output_bindings;
-		const std::vector<std::vector<unsigned int>>& COLORS,
+    std::vector<void*>   host_ptrs;
-		const std::vector<std::vector<unsigned int>>& MASK_COLORS
+    std::vector<void*>   device_ptrs;
-	);
+
-	int num_bindings;
+    PreParam pparam;
-	int num_inputs = 0;
-	int num_outputs = 0;
-	std::vector<Binding> input_bindings;
-	std::vector<Binding> output_bindings;
-	std::vector<void*> host_ptrs;
-	std::vector<void*> device_ptrs;
-
-	PreParam pparam;
-private:
-	nvinfer1::ICudaEngine* engine = nullptr;
-	nvinfer1::IRuntime* runtime = nullptr;
-	nvinfer1::IExecutionContext* context = nullptr;
-	cudaStream_t stream = nullptr;
-	Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
 YOLOv8_seg::YOLOv8_seg(const std::string& engine_file_path)
 {
-	std::ifstream file(engine_file_path, std::ios::binary);
+    std::ifstream file(engine_file_path, std::ios::binary);
-	assert(file.good());
+    assert(file.good());
-	file.seekg(0, std::ios::end);
+    file.seekg(0, std::ios::end);
-	auto size = file.tellg();
+    auto size = file.tellg();
-	file.seekg(0, std::ios::beg);
+    file.seekg(0, std::ios::beg);
-	char* trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
-	assert(trtModelStream);
+    assert(trtModelStream);
-	file.read(trtModelStream, size);
+    file.read(trtModelStream, size);
-	file.close();
+    file.close();
-	initLibNvInferPlugins(&this->gLogger, "");
+    initLibNvInferPlugins(&this->gLogger, "");
-	this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
-	assert(this->runtime != nullptr);
+    assert(this->runtime != nullptr);
-
+
-	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
-	assert(this->engine != nullptr);
+    assert(this->engine != nullptr);
-	delete[] trtModelStream;
+    delete[] trtModelStream;
-	this->context = this->engine->createExecutionContext();
+    this->context = this->engine->createExecutionContext();
-
+
-	assert(this->context != nullptr);
+    assert(this->context != nullptr);
-	cudaStreamCreate(&this->stream);
+    cudaStreamCreate(&this->stream);
-	this->num_bindings = this->engine->getNbBindings();
+    this->num_bindings = this->engine->getNbBindings();
-
+
-	for (int i = 0; i < this->num_bindings; ++i)
+    for (int i = 0; i < this->num_bindings; ++i) {
-	{
+        Binding            binding;
-		Binding binding;
+        nvinfer1::Dims     dims;
-		nvinfer1::Dims dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-		nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
-		std::string name = this->engine->getBindingName(i);
+        binding.name             = name;
-		binding.name = name;
+        binding.dsize            = type_to_size(dtype);
-		binding.dsize = type_to_size(dtype);
+
-
+        bool IsInput = engine->bindingIsInput(i);
-		bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
-		if (IsInput)
+            this->num_inputs += 1;
-		{
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-			this->num_inputs += 1;
+            binding.size = get_size_by_dims(dims);
-			dims = this->engine->getProfileDimensions(
+            binding.dims = dims;
-				i,
+            this->input_bindings.push_back(binding);
-				0,
+            // set max opt shape
-				nvinfer1::OptProfileSelector::kMAX);
+            this->context->setBindingDimensions(i, dims);
-			binding.size = get_size_by_dims(dims);
+        }
-			binding.dims = dims;
+        else {
-			this->input_bindings.push_back(binding);
+            dims         = this->context->getBindingDimensions(i);
-			// set max opt shape
+            binding.size = get_size_by_dims(dims);
-			this->context->setBindingDimensions(i, dims);
+            binding.dims = dims;
-
+            this->output_bindings.push_back(binding);
-		}
+            this->num_outputs += 1;
-		else
+        }
-		{
+    }
-			dims = this->context->getBindingDimensions(i);
-			binding.size = get_size_by_dims(dims);
-			binding.dims = dims;
-			this->output_bindings.push_back(binding);
-			this->num_outputs += 1;
-		}
-	}
-
 }
 
 YOLOv8_seg::~YOLOv8_seg()
 {
-	this->context->destroy();
+    this->context->destroy();
-	this->engine->destroy();
+    this->engine->destroy();
-	this->runtime->destroy();
+    this->runtime->destroy();
-	cudaStreamDestroy(this->stream);
+    cudaStreamDestroy(this->stream);
-	for (auto& ptr : this->device_ptrs)
+    for (auto& ptr : this->device_ptrs) {
-	{
+        CHECK(cudaFree(ptr));
-		CHECK(cudaFree(ptr));
+    }
-	}
+
-
+    for (auto& ptr : this->host_ptrs) {
-	for (auto& ptr : this->host_ptrs)
+        CHECK(cudaFreeHost(ptr));
-	{
+    }
-		CHECK(cudaFreeHost(ptr));
-	}
 }
 
 void YOLOv8_seg::make_pipe(bool warmup)
 {
 
-	for (auto& bindings : this->input_bindings)
+    for (auto& bindings : this->input_bindings) {
-	{
+        void* d_ptr;
-		void* d_ptr;
+        CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream));
-		CHECK(cudaMallocAsync(
+        this->device_ptrs.push_back(d_ptr);
-			&d_ptr,
+    }
-			bindings.size * bindings.dsize,
+
-			this->stream)
+    for (auto& bindings : this->output_bindings) {
-		);
+        void * d_ptr, *h_ptr;
-		this->device_ptrs.push_back(d_ptr);
+        size_t size = bindings.size * bindings.dsize;
-	}
+        CHECK(cudaMallocAsync(&d_ptr, size, this->stream));
-
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-	for (auto& bindings : this->output_bindings)
+        this->device_ptrs.push_back(d_ptr);
-	{
+        this->host_ptrs.push_back(h_ptr);
-		void* d_ptr, * h_ptr;
+    }
-		size_t size = bindings.size * bindings.dsize;
+
-		CHECK(cudaMallocAsync(
+    if (warmup) {
-			&d_ptr,
+        for (int i = 0; i < 10; i++) {
-			size,
+            for (auto& bindings : this->input_bindings) {
-			this->stream)
+                size_t size  = bindings.size * bindings.dsize;
-		);
+                void*  h_ptr = malloc(size);
-		CHECK(cudaHostAlloc(
+                memset(h_ptr, 0, size);
-			&h_ptr,
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-			size,
+                free(h_ptr);
-			0)
+            }
-		);
+            this->infer();
-		this->device_ptrs.push_back(d_ptr);
+        }
-		this->host_ptrs.push_back(h_ptr);
+        printf("model warmup 10 times\n");
-	}
+    }
-
-	if (warmup)
-	{
-		for (int i = 0; i < 10; i++)
-		{
-			for (auto& bindings : this->input_bindings)
-			{
-				size_t size = bindings.size * bindings.dsize;
-				void* h_ptr = malloc(size);
-				memset(h_ptr, 0, size);
-				CHECK(cudaMemcpyAsync(
-					this->device_ptrs[0],
-					h_ptr,
-					size,
-					cudaMemcpyHostToDevice,
-					this->stream)
-				);
-				free(h_ptr);
-			}
-			this->infer();
-		}
-		printf("model warmup 10 times\n");
-
-	}
 }
 
-void YOLOv8_seg::letterbox(
+void YOLOv8_seg::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
-	const cv::Mat& image,
-	cv::Mat& out,
-	cv::Size& size
-)
 {
-	const float inp_h = size.height;
+    const float inp_h  = size.height;
-	const float inp_w = size.width;
+    const float inp_w  = size.width;
-	float height = image.rows;
+    float       height = image.rows;
-	float width = image.cols;
+    float       width  = image.cols;
-
+
-	float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-	int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-	int padh = std::round(height * r);
+    int   padh = std::round(height * r);
-
+
-	cv::Mat tmp;
+    cv::Mat tmp;
-	if ((int)width != padw || (int)height != padh)
+    if ((int)width != padw || (int)height != padh) {
-	{
+        cv::resize(image, tmp, cv::Size(padw, padh));
-		cv::resize(
+    }
-			image,
+    else {
-			tmp,
+        tmp = image.clone();
-			cv::Size(padw, padh)
+    }
-		);
+
-	}
+    float dw = inp_w - padw;
-	else
+    float dh = inp_h - padh;
-	{
+
-		tmp = image.clone();
+    dw /= 2.0f;
-	}
+    dh /= 2.0f;
-
+    int top    = int(std::round(dh - 0.1f));
-	float dw = inp_w - padw;
+    int bottom = int(std::round(dh + 0.1f));
-	float dh = inp_h - padh;
+    int left   = int(std::round(dw - 0.1f));
-
+    int right  = int(std::round(dw + 0.1f));
-	dw /= 2.0f;
+
-	dh /= 2.0f;
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-	int top = int(std::round(dh - 0.1f));
+
-	int bottom = int(std::round(dh + 0.1f));
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-	int left = int(std::round(dw - 0.1f));
+    this->pparam.ratio  = 1 / r;
-	int right = int(std::round(dw + 0.1f));
+    this->pparam.dw     = dw;
-
+    this->pparam.dh     = dh;
-	cv::copyMakeBorder(
+    this->pparam.height = height;
-		tmp,
+    this->pparam.width  = width;
-		tmp,
+    ;
-		top,
-		bottom,
-		left,
-		right,
-		cv::BORDER_CONSTANT,
-		{ 114, 114, 114 }
-	);
-
-	cv::dnn::blobFromImage(tmp,
-		out,
-		1 / 255.f,
-		cv::Size(),
-		cv::Scalar(0, 0, 0),
-		true,
-		false,
-		CV_32F
-	);
-	this->pparam.ratio = 1 / r;
-	this->pparam.dw = dw;
-	this->pparam.dh = dh;
-	this->pparam.height = height;
-	this->pparam.width = width;;
 }
 
 void YOLOv8_seg::copy_from_Mat(const cv::Mat& image)
 {
-	cv::Mat nchw;
+    cv::Mat  nchw;
-	auto& in_binding = this->input_bindings[0];
+    auto&    in_binding = this->input_bindings[0];
-	auto width = in_binding.dims.d[3];
+    auto     width      = in_binding.dims.d[3];
-	auto height = in_binding.dims.d[2];
+    auto     height     = in_binding.dims.d[2];
-	cv::Size size{ width, height };
+    cv::Size size{width, height};
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+
-		nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-		size
+
-	);
+    CHECK(cudaMemcpyAsync(
-
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{
-				4,
-				{ 1, 3, height, width }
-			}
-	);
-
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8_seg::copy_from_Mat(const cv::Mat& image, cv::Size& size)
 {
-	cv::Mat nchw;
+    cv::Mat nchw;
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-		nchw,
+    CHECK(cudaMemcpyAsync(
-		size
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	);
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{ 4,
-			  { 1, 3, size.height, size.width }
-			}
-	);
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8_seg::infer()
 {
 
-	this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-		this->device_ptrs.data(),
+    for (int i = 0; i < this->num_outputs; i++) {
-		this->stream,
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		nullptr
+        CHECK(cudaMemcpyAsync(
-	);
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-	for (int i = 0; i < this->num_outputs; i++)
+    }
-	{
+    cudaStreamSynchronize(this->stream);
-		size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		CHECK(cudaMemcpyAsync(this->host_ptrs[i],
-			this->device_ptrs[i + this->num_inputs],
-			osize,
-			cudaMemcpyDeviceToHost,
-			this->stream)
-		);
-
-	}
-	cudaStreamSynchronize(this->stream);
-
 }
 
-void YOLOv8_seg::postprocess(std::vector<Object>& objs,
+void YOLOv8_seg::postprocess(
-	float score_thres,
+    std::vector<Object>& objs, float score_thres, float iou_thres, int topk, int seg_channels, int seg_h, int seg_w)
-	float iou_thres,
-	int topk,
-	int seg_channels,
-	int seg_h,
-	int seg_w
-)
 {
-	objs.clear();
+    objs.clear();
-	auto input_h = this->input_bindings[0].dims.d[2];
+    auto input_h = this->input_bindings[0].dims.d[2];
-	auto input_w = this->input_bindings[0].dims.d[3];
+    auto input_w = this->input_bindings[0].dims.d[3];
-	int num_channels, num_anchors, num_classes;
+    int  num_channels, num_anchors, num_classes;
-	bool flag = false;
+    bool flag = false;
-	int bid;
+    int  bid;
-	int bcnt = -1;
+    int  bcnt = -1;
-	for (auto& o : this->output_bindings)
+    for (auto& o : this->output_bindings) {
-	{
+        bcnt += 1;
-		bcnt += 1;
+        if (o.dims.nbDims == 3) {
-		if (o.dims.nbDims == 3)
+            num_channels = o.dims.d[1];
-		{
+            num_anchors  = o.dims.d[2];
-			num_channels = o.dims.d[1];
+            flag         = true;
-			num_anchors = o.dims.d[2];
+            bid          = bcnt;
-			flag = true;
+        }
-			bid = bcnt;
+    }
-		}
+    assert(flag);
-	}
+    num_classes = num_channels - seg_channels - 4;
-	assert(flag);
+
-	num_classes = num_channels - seg_channels - 4;
+    auto& dw     = this->pparam.dw;
-
+    auto& dh     = this->pparam.dh;
-	auto& dw = this->pparam.dw;
+    auto& width  = this->pparam.width;
-	auto& dh = this->pparam.dh;
+    auto& height = this->pparam.height;
-	auto& width = this->pparam.width;
+    auto& ratio  = this->pparam.ratio;
-	auto& height = this->pparam.height;
+
-	auto& ratio = this->pparam.ratio;
+    cv::Mat output = cv::Mat(num_channels, num_anchors, CV_32F, static_cast<float*>(this->host_ptrs[bid]));
-
+    output         = output.t();
-	cv::Mat output = cv::Mat(num_channels, num_anchors, CV_32F,
+
-		static_cast<float*>(this->host_ptrs[bid]));
+    cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F, static_cast<float*>(this->host_ptrs[1 - bid]));
-	output = output.t();
+
-
+    std::vector<int>      labels;
-	cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F,
+    std::vector<float>    scores;
-		static_cast<float*>(this->host_ptrs[1 - bid]));
+    std::vector<cv::Rect> bboxes;
-
+    std::vector<cv::Mat>  mask_confs;
-	std::vector<int> labels;
+    std::vector<int>      indices;
-	std::vector<float> scores;
+
-	std::vector<cv::Rect> bboxes;
+    for (int i = 0; i < num_anchors; i++) {
-	std::vector<cv::Mat> mask_confs;
+        auto  row_ptr        = output.row(i).ptr<float>();
-	std::vector<int> indices;
+        auto  bboxes_ptr     = row_ptr;
-
+        auto  scores_ptr     = row_ptr + 4;
-	for (int i = 0; i < num_anchors; i++)
+        auto  mask_confs_ptr = row_ptr + 4 + num_classes;
-	{
+        auto  max_s_ptr      = std::max_element(scores_ptr, scores_ptr + num_classes);
-		auto row_ptr = output.row(i).ptr<float>();
+        float score          = *max_s_ptr;
-		auto bboxes_ptr = row_ptr;
+        if (score > score_thres) {
-		auto scores_ptr = row_ptr + 4;
+            float x = *bboxes_ptr++ - dw;
-		auto mask_confs_ptr = row_ptr + 4 + num_classes;
+            float y = *bboxes_ptr++ - dh;
-		auto max_s_ptr = std::max_element(scores_ptr, scores_ptr + num_classes);
+            float w = *bboxes_ptr++;
-		float score = *max_s_ptr;
+            float h = *bboxes_ptr;
-		if (score > score_thres)
+
-		{
+            float x0 = clamp((x - 0.5f * w) * ratio, 0.f, width);
-			float x = *bboxes_ptr++ - dw;
+            float y0 = clamp((y - 0.5f * h) * ratio, 0.f, height);
-			float y = *bboxes_ptr++ - dh;
+            float x1 = clamp((x + 0.5f * w) * ratio, 0.f, width);
-			float w = *bboxes_ptr++;
+            float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);
-			float h = *bboxes_ptr;
+
-
+            int              label = max_s_ptr - scores_ptr;
-			float x0 = clamp((x - 0.5f * w) * ratio, 0.f, width);
+            cv::Rect_<float> bbox;
-			float y0 = clamp((y - 0.5f * h) * ratio, 0.f, height);
+            bbox.x      = x0;
-			float x1 = clamp((x + 0.5f * w) * ratio, 0.f, width);
+            bbox.y      = y0;
-			float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);
+            bbox.width  = x1 - x0;
-
+            bbox.height = y1 - y0;
-			int label = max_s_ptr - scores_ptr;
+
-			cv::Rect_<float> bbox;
+            cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, mask_confs_ptr);
-			bbox.x = x0;
+
-			bbox.y = y0;
+            bboxes.push_back(bbox);
-			bbox.width = x1 - x0;
+            labels.push_back(label);
-			bbox.height = y1 - y0;
+            scores.push_back(score);
-
+            mask_confs.push_back(mask_conf);
-			cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, mask_confs_ptr);
+        }
-
+    }
-			bboxes.push_back(bbox);
-			labels.push_back(label);
-			scores.push_back(score);
-			mask_confs.push_back(mask_conf);
-		}
-	}
 
 #if defined(BATCHED_NMS)
-	cv::dnn::NMSBoxesBatched(
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		labels,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #else
-	cv::dnn::NMSBoxes(
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #endif
 
-	cv::Mat masks;
+    cv::Mat masks;
-	int cnt = 0;
+    int     cnt = 0;
-	for (auto& i : indices)
+    for (auto& i : indices) {
-	{
+        if (cnt >= topk) {
-		if (cnt >= topk)
+            break;
-		{
+        }
-			break;
+        cv::Rect tmp = bboxes[i];
-		}
+        Object   obj;
-		cv::Rect tmp = bboxes[i];
+        obj.label = labels[i];
-		Object obj;
+        obj.rect  = tmp;
-		obj.label = labels[i];
+        obj.prob  = scores[i];
-		obj.rect = tmp;
+        masks.push_back(mask_confs[i]);
-		obj.prob = scores[i];
+        objs.push_back(obj);
-		masks.push_back(mask_confs[i]);
+        cnt += 1;
-		objs.push_back(obj);
+    }
-		cnt += 1;
+
-	}
+    if (masks.empty()) {
-
+        // masks is empty
-	if(masks.empty())
+    }
-	{
+    else {
-		//masks is empty
+        cv::Mat matmulRes = (masks * protos).t();
-	}
+        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_w, seg_h});
-	else
+
-	{
+        std::vector<cv::Mat> maskChannels;
-		cv::Mat matmulRes = (masks * protos).t();
+        cv::split(maskMat, maskChannels);
-		cv::Mat maskMat = matmulRes.reshape(indices.size(), { seg_w, seg_h });
+        int scale_dw = dw / input_w * seg_w;
-
+        int scale_dh = dh / input_h * seg_h;
-		std::vector<cv::Mat> maskChannels;
+
-		cv::split(maskMat, maskChannels);
+        cv::Rect roi(scale_dw, scale_dh, seg_w - 2 * scale_dw, seg_h - 2 * scale_dh);
-		int scale_dw = dw / input_w * seg_w;
+
-		int scale_dh = dh / input_h * seg_h;
+        for (int i = 0; i < indices.size(); i++) {
-
+            cv::Mat dest, mask;
-		cv::Rect roi(
+            cv::exp(-maskChannels[i], dest);
-			scale_dw,
+            dest = 1.0 / (1.0 + dest);
-			scale_dh,
+            dest = dest(roi);
-			seg_w - 2 * scale_dw,
+            cv::resize(dest, mask, cv::Size((int)width, (int)height), cv::INTER_LINEAR);
-			seg_h - 2 * scale_dh);
+            objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-
+        }
-		for (int i = 0; i < indices.size(); i++)
+    }
-		{
-			cv::Mat dest, mask;
-			cv::exp(-maskChannels[i], dest);
-			dest = 1.0 / (1.0 + dest);
-			dest = dest(roi);
-			cv::resize(
-				dest,
-				mask,
-				cv::Size((int)width, (int)height),
-				cv::INTER_LINEAR
-			);
-			objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-		}
-	}
-
 }
 
-void YOLOv8_seg::draw_objects(const cv::Mat& image,
+void YOLOv8_seg::draw_objects(const cv::Mat&                                image,
-	cv::Mat& res,
+                              cv::Mat&                                      res,
-	const std::vector<Object>& objs,
+                              const std::vector<Object>&                    objs,
-	const std::vector<std::string>& CLASS_NAMES,
+                              const std::vector<std::string>&               CLASS_NAMES,
-	const std::vector<std::vector<unsigned int>>& COLORS,
+                              const std::vector<std::vector<unsigned int>>& COLORS,
-	const std::vector<std::vector<unsigned int>>& MASK_COLORS
+                              const std::vector<std::vector<unsigned int>>& MASK_COLORS)
-)
 {
-	res = image.clone();
+    res          = image.clone();
-	cv::Mat mask = image.clone();
+    cv::Mat mask = image.clone();
-	for (auto& obj : objs)
+    for (auto& obj : objs) {
-	{
+        int        idx   = obj.label;
-		int idx = obj.label;
+        cv::Scalar color = cv::Scalar(COLORS[idx][0], COLORS[idx][1], COLORS[idx][2]);
-		cv::Scalar color = cv::Scalar(
+        cv::Scalar mask_color =
-			COLORS[idx][0],
+            cv::Scalar(MASK_COLORS[idx % 20][0], MASK_COLORS[idx % 20][1], MASK_COLORS[idx % 20][2]);
-			COLORS[idx][1],
+        cv::rectangle(res, obj.rect, color, 2);
-			COLORS[idx][2]
+
-		);
+        char text[256];
-		cv::Scalar mask_color = cv::Scalar(
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[idx].c_str(), obj.prob * 100);
-			MASK_COLORS[idx % 20][0],
+        mask(obj.rect).setTo(mask_color, obj.boxMask);
-			MASK_COLORS[idx % 20][1],
+
-			MASK_COLORS[idx % 20][2]
+        int      baseLine   = 0;
-		);
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-		cv::rectangle(
+
-			res,
+        int x = (int)obj.rect.x;
-			obj.rect,
+        int y = (int)obj.rect.y + 1;
-			color,
+
-			2
+        if (y > res.rows)
-		);
+            y = res.rows;
-
+
-		char text[256];
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-		sprintf(
+
-			text,
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-			"%s %.1f%%",
+    }
-			CLASS_NAMES[idx].c_str(),
+    cv::addWeighted(res, 0.5, mask, 0.8, 1, res);
-			obj.prob * 100
-		);
-		mask(obj.rect).setTo(mask_color, obj.boxMask);
-
-		int baseLine = 0;
-		cv::Size label_size = cv::getTextSize(
-			text,
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			1,
-			&baseLine
-		);
-
-		int x = (int)obj.rect.x;
-		int y = (int)obj.rect.y + 1;
-
-		if (y > res.rows)
-			y = res.rows;
-
-		cv::rectangle(
-			res,
-			cv::Rect(x, y, label_size.width, label_size.height + baseLine),
-			{ 0, 0, 255 },
-			-1
-		);
-
-		cv::putText(
-			res,
-			text,
-			cv::Point(x, y + label_size.height),
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			{ 255, 255, 255 },
-			1
-		);
-	}
-	cv::addWeighted(
-		res,
-		0.5,
-		mask,
-		0.8,
-		1,
-		res
-	);
 }
-#endif //SEGMENT_NORMAL_YOLOV8_SEG_HPP
+#endif  // SEGMENT_NORMAL_YOLOV8_SEG_HPP

csrc/segment/normal/main.cpp

@@ -2,177 +2,131 @@
 // Created by ubuntu on 2/8/23.
 //
 #include "chrono"
-#include "yolov8-seg.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8-seg.hpp"
 
 const std::vector<std::string> CLASS_NAMES = {
-	"person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",
-	"train", "truck", "boat", "traffic light", "fire hydrant",
+    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",    "parking meter", "bench",
-	"stop sign", "parking meter", "bench", "bird", "cat",
+    "bird",           "cat",        "dog",           "horse",         "sheep",        "cow",           "elephant",
-	"dog", "horse", "sheep", "cow", "elephant",
+    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",     "handbag",       "tie",
-	"bear", "zebra", "giraffe", "backpack", "umbrella",
+    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball",  "kite",          "baseball bat",
-	"handbag", "tie", "suitcase", "frisbee", "skis",
+    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",       "wine glass",    "cup",
-	"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "fork",           "knife",      "spoon",         "bowl",          "banana",       "apple",         "sandwich",
-	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
+    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",        "donut",         "cake",
-	"cup", "fork", "knife", "spoon", "bowl",
+    "chair",          "couch",      "potted plant",  "bed",           "dining table", "toilet",        "tv",
-	"banana", "apple", "sandwich", "orange", "broccoli",
+    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",   "microwave",     "oven",
-	"carrot", "hot dog", "pizza", "donut", "cake",
+    "toaster",        "sink",       "refrigerator",  "book",          "clock",        "vase",          "scissors",
-	"chair", "couch", "potted plant", "bed", "dining table",
+    "teddy bear",     "hair drier", "toothbrush"};
-	"toilet", "tv", "laptop", "mouse", "remote",
-	"keyboard", "cell phone", "microwave", "oven",
-	"toaster", "sink", "refrigerator", "book", "clock", "vase",
-	"scissors", "teddy bear", "hair drier", "toothbrush" };
 
 const std::vector<std::vector<unsigned int>> COLORS = {
-	{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
+    {0, 114, 189},   {217, 83, 25},   {237, 177, 32},  {126, 47, 142},  {119, 172, 48},  {77, 190, 238},
-	{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
+    {162, 20, 47},   {76, 76, 76},    {153, 153, 153}, {255, 0, 0},     {255, 128, 0},   {191, 191, 0},
-	{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
+    {0, 255, 0},     {0, 0, 255},     {170, 0, 255},   {85, 85, 0},     {85, 170, 0},    {85, 255, 0},
-	{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
+    {170, 85, 0},    {170, 170, 0},   {170, 255, 0},   {255, 85, 0},    {255, 170, 0},   {255, 255, 0},
-	{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
+    {0, 85, 128},    {0, 170, 128},   {0, 255, 128},   {85, 0, 128},    {85, 85, 128},   {85, 170, 128},
-	{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
+    {85, 255, 128},  {170, 0, 128},   {170, 85, 128},  {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
-	{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
+    {255, 85, 128},  {255, 170, 128}, {255, 255, 128}, {0, 85, 255},    {0, 170, 255},   {0, 255, 255},
-	{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
+    {85, 0, 255},    {85, 85, 255},   {85, 170, 255},  {85, 255, 255},  {170, 0, 255},   {170, 85, 255},
-	{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
+    {170, 170, 255}, {170, 255, 255}, {255, 0, 255},   {255, 85, 255},  {255, 170, 255}, {85, 0, 0},
-	{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
+    {128, 0, 0},     {170, 0, 0},     {212, 0, 0},     {255, 0, 0},     {0, 43, 0},      {0, 85, 0},
-	{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
+    {0, 128, 0},     {0, 170, 0},     {0, 212, 0},     {0, 255, 0},     {0, 0, 43},      {0, 0, 85},
-	{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
+    {0, 0, 128},     {0, 0, 170},     {0, 0, 212},     {0, 0, 255},     {0, 0, 0},       {36, 36, 36},
-	{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
+    {73, 73, 73},    {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
-	{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
+    {80, 183, 189},  {128, 128, 0}};
-	{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
-	{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
-	{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
-	{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
-	{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
-	{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
-	{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
-	{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
-	{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
-	{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
-	{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
-	{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
-	{ 80, 183, 189 }, { 128, 128, 0 }
-};
 
 const std::vector<std::vector<unsigned int>> MASK_COLORS = {
-	{ 255, 56, 56 }, { 255, 157, 151 }, { 255, 112, 31 },
+    {255, 56, 56},  {255, 157, 151}, {255, 112, 31}, {255, 178, 29}, {207, 210, 49},  {72, 249, 10}, {146, 204, 23},
-	{ 255, 178, 29 }, { 207, 210, 49 }, { 72, 249, 10 },
+    {61, 219, 134}, {26, 147, 52},   {0, 212, 187},  {44, 153, 168}, {0, 194, 255},   {52, 69, 147}, {100, 115, 255},
-	{ 146, 204, 23 }, { 61, 219, 134 }, { 26, 147, 52 },
+    {0, 24, 236},   {132, 56, 255},  {82, 0, 133},   {203, 56, 255}, {255, 149, 200}, {255, 55, 199}};
-	{ 0, 212, 187 }, { 44, 153, 168 }, { 0, 194, 255 },
-	{ 52, 69, 147 }, { 100, 115, 255 }, { 0, 24, 236 },
-	{ 132, 56, 255 }, { 82, 0, 133 }, { 203, 56, 255 },
-	{ 255, 149, 200 }, { 255, 55, 199 }
-};
 
 int main(int argc, char** argv)
 {
-	// cuda:0
+    // cuda:0
-	cudaSetDevice(0);
+    cudaSetDevice(0);
 
-	const std::string engine_file_path{ argv[1] };
+    const std::string engine_file_path{argv[1]};
-	const std::string path{ argv[2] };
+    const std::string path{argv[2]};
 
-	std::vector<std::string> imagePathList;
+    std::vector<std::string> imagePathList;
-	bool isVideo{ false };
+    bool                     isVideo{false};
 
-	assert(argc == 3);
+    assert(argc == 3);
 
-	auto yolov8 = new YOLOv8_seg(engine_file_path);
+    auto yolov8 = new YOLOv8_seg(engine_file_path);
-	yolov8->make_pipe(true);
+    yolov8->make_pipe(true);
 
-	if (IsFile(path))
+    if (IsFile(path)) {
-	{
+        std::string suffix = path.substr(path.find_last_of('.') + 1);
-		std::string suffix = path.substr(path.find_last_of('.') + 1);
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-		if (
+            imagePathList.push_back(path);
-			suffix == "jpg" ||
+        }
-				suffix == "jpeg" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-				suffix == "png"
+                 || suffix == "mkv") {
-			)
+            isVideo = true;
-		{
+        }
-			imagePathList.push_back(path);
+        else {
-		}
+            printf("suffix %s is wrong !!!\n", suffix.c_str());
-		else if (
+            std::abort();
-			suffix == "mp4" ||
+        }
-				suffix == "avi" ||
+    }
-				suffix == "m4v" ||
+    else if (IsFolder(path)) {
-				suffix == "mpeg" ||
+        cv::glob(path + "/*.jpg", imagePathList);
-				suffix == "mov" ||
+    }
-				suffix == "mkv"
-			)
-		{
-			isVideo = true;
-		}
-		else
-		{
-			printf("suffix %s is wrong !!!\n", suffix.c_str());
-			std::abort();
-		}
-	}
-	else if (IsFolder(path))
-	{
-		cv::glob(path + "/*.jpg", imagePathList);
-	}
 
-	cv::Mat res, image;
+    cv::Mat  res, image;
-	cv::Size size = cv::Size{ 640, 640 };
+    cv::Size size         = cv::Size{640, 640};
-	int topk = 100;
+    int      topk         = 100;
-	int seg_h = 160;
+    int      seg_h        = 160;
-	int seg_w = 160;
+    int      seg_w        = 160;
-	int seg_channels = 32;
+    int      seg_channels = 32;
-	float score_thres = 0.25f;
+    float    score_thres  = 0.25f;
-	float iou_thres = 0.65f;
+    float    iou_thres    = 0.65f;
 
-	std::vector<Object> objs;
+    std::vector<Object> objs;
 
-	cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
+    cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
 
-	if (isVideo)
+    if (isVideo) {
-	{
+        cv::VideoCapture cap(path);
-		cv::VideoCapture cap(path);
 
-		if (!cap.isOpened())
+        if (!cap.isOpened()) {
-		{
+            printf("can not open %s\n", path.c_str());
-			printf("can not open %s\n", path.c_str());
+            return -1;
-			return -1;
+        }
-		}
+        while (cap.read(image)) {
-		while (cap.read(image))
+            objs.clear();
-		{
+            yolov8->copy_from_Mat(image, size);
-			objs.clear();
+            auto start = std::chrono::system_clock::now();
-			yolov8->copy_from_Mat(image, size);
+            yolov8->infer();
-			auto start = std::chrono::system_clock::now();
+            auto end = std::chrono::system_clock::now();
-			yolov8->infer();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
-			auto end = std::chrono::system_clock::now();
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
+            printf("cost %2.4lf ms\n", tc);
-			auto tc = (double)
+            cv::imshow("result", res);
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+            if (cv::waitKey(10) == 'q') {
-			printf("cost %2.4lf ms\n", tc);
+                break;
-			cv::imshow("result", res);
+            }
-			if (cv::waitKey(10) == 'q')
+        }
-			{
+    }
-				break;
+    else {
-			}
+        for (auto& path : imagePathList) {
-		}
+            objs.clear();
-	}
+            image = cv::imread(path);
-	else
+            yolov8->copy_from_Mat(image, size);
-	{
+            auto start = std::chrono::system_clock::now();
-		for (auto& path : imagePathList)
+            yolov8->infer();
-		{
+            auto end = std::chrono::system_clock::now();
-			objs.clear();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
-			image = cv::imread(path);
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
-			yolov8->copy_from_Mat(image, size);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			auto start = std::chrono::system_clock::now();
+            printf("cost %2.4lf ms\n", tc);
-			yolov8->infer();
+            cv::imshow("result", res);
-			auto end = std::chrono::system_clock::now();
+            cv::waitKey(0);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
+        }
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
+    }
-			auto tc = (double)
+    cv::destroyAllWindows();
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+    delete yolov8;
-			printf("cost %2.4lf ms\n", tc);
+    return 0;
-			cv::imshow("result", res);
-			cv::waitKey(0);
-		}
-	}
-	cv::destroyAllWindows();
-	delete yolov8;
-	return 0;
 }

csrc/segment/simple/include/common.hpp

@@ -4,154 +4,139 @@
 
 #ifndef SEGMENT_SIMPLE_COMMON_HPP
 #define SEGMENT_SIMPLE_COMMON_HPP
+#include "NvInfer.h"
 #include "opencv2/opencv.hpp"
 #include <sys/stat.h>
 #include <unistd.h>
-#include "NvInfer.h"
 
-#define CHECK(call)                                   \
+#define CHECK(call)                                                                                                    \
-do                                                    \
+    do {                                                                                                               \
-{                                                     \
+        const cudaError_t error_code = call;                                                                           \
-    const cudaError_t error_code = call;              \
+        if (error_code != cudaSuccess) {                                                                               \
-    if (error_code != cudaSuccess)                    \
+            printf("CUDA Error:\n");                                                                                   \
-    {                                                 \
+            printf("    File:       %s\n", __FILE__);                                                                  \
-        printf("CUDA Error:\n");                      \
+            printf("    Line:       %d\n", __LINE__);                                                                  \
-        printf("    File:       %s\n", __FILE__);     \
+            printf("    Error code: %d\n", error_code);                                                                \
-        printf("    Line:       %d\n", __LINE__);     \
+            printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \
-        printf("    Error code: %d\n", error_code);   \
+            exit(1);                                                                                                   \
-        printf("    Error text: %s\n",                \
+        }                                                                                                              \
-            cudaGetErrorString(error_code));          \
+    } while (0)
-        exit(1);                                      \
-    }                                                 \
-} while (0)
 
-class Logger : public nvinfer1::ILogger
+class Logger: public nvinfer1::ILogger {
-{
 public:
-	nvinfer1::ILogger::Severity reportableSeverity;
+    nvinfer1::ILogger::Severity reportableSeverity;
 
-	explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO) :
+    explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):
-		reportableSeverity(severity)
+        reportableSeverity(severity)
-	{
+    {
-	}
+    }
 
-	void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
+    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override
-	{
+    {
-		if (severity > reportableSeverity)
+        if (severity > reportableSeverity) {
-		{
+            return;
-			return;
+        }
-		}
+        switch (severity) {
-		switch (severity)
+            case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
-		{
+                std::cerr << "INTERNAL_ERROR: ";
-		case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:
+                break;
-			std::cerr << "INTERNAL_ERROR: ";
+            case nvinfer1::ILogger::Severity::kERROR:
-			break;
+                std::cerr << "ERROR: ";
-		case nvinfer1::ILogger::Severity::kERROR:
+                break;
-			std::cerr << "ERROR: ";
+            case nvinfer1::ILogger::Severity::kWARNING:
-			break;
+                std::cerr << "WARNING: ";
-		case nvinfer1::ILogger::Severity::kWARNING:
+                break;
-			std::cerr << "WARNING: ";
+            case nvinfer1::ILogger::Severity::kINFO:
-			break;
+                std::cerr << "INFO: ";
-		case nvinfer1::ILogger::Severity::kINFO:
+                break;
-			std::cerr << "INFO: ";
+            default:
-			break;
+                std::cerr << "VERBOSE: ";
-		default:
+                break;
-			std::cerr << "VERBOSE: ";
+        }
-			break;
+        std::cerr << msg << std::endl;
-		}
+    }
-		std::cerr << msg << std::endl;
-	}
 };
 
 inline int get_size_by_dims(const nvinfer1::Dims& dims)
 {
-	int size = 1;
+    int size = 1;
-	for (int i = 0; i < dims.nbDims; i++)
+    for (int i = 0; i < dims.nbDims; i++) {
-	{
+        size *= dims.d[i];
-		size *= dims.d[i];
+    }
-	}
+    return size;
-	return size;
 }
 
 inline int type_to_size(const nvinfer1::DataType& dataType)
 {
-	switch (dataType)
+    switch (dataType) {
-	{
+        case nvinfer1::DataType::kFLOAT:
-	case nvinfer1::DataType::kFLOAT:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kHALF:
-	case nvinfer1::DataType::kHALF:
+            return 2;
-		return 2;
+        case nvinfer1::DataType::kINT32:
-	case nvinfer1::DataType::kINT32:
+            return 4;
-		return 4;
+        case nvinfer1::DataType::kINT8:
-	case nvinfer1::DataType::kINT8:
+            return 1;
-		return 1;
+        case nvinfer1::DataType::kBOOL:
-	case nvinfer1::DataType::kBOOL:
+            return 1;
-		return 1;
+        default:
-	default:
+            return 4;
-		return 4;
+    }
-	}
 }
 
 inline static float clamp(float val, float min, float max)
 {
-	return val > min ? (val < max ? val : max) : min;
+    return val > min ? (val < max ? val : max) : min;
 }
 
 inline bool IsPathExist(const std::string& path)
 {
-	if (access(path.c_str(), 0) == F_OK)
+    if (access(path.c_str(), 0) == F_OK) {
-	{
+        return true;
-		return true;
+    }
-	}
+    return false;
-	return false;
 }
 
 inline bool IsFile(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
-		printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
 }
 
 inline bool IsFolder(const std::string& path)
 {
-	if (!IsPathExist(path))
+    if (!IsPathExist(path)) {
-	{
+        return false;
-		return false;
+    }
-	}
+    struct stat buffer;
-	struct stat buffer;
+    return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
-	return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-namespace seg
+namespace seg {
-{
+struct Binding {
-	struct Binding
+    size_t         size  = 1;
-	{
+    size_t         dsize = 1;
-		size_t size = 1;
+    nvinfer1::Dims dims;
-		size_t dsize = 1;
+    std::string    name;
-		nvinfer1::Dims dims;
+};
-		std::string name;
-	};
 
-	struct Object
+struct Object {
-	{
+    cv::Rect_<float> rect;
-		cv::Rect_<float> rect;
+    int              label = 0;
-		int label = 0;
+    float            prob  = 0.0;
-		float prob = 0.0;
+    cv::Mat          boxMask;
-		cv::Mat boxMask;
+};
-	};
 
-	struct PreParam
+struct PreParam {
-	{
+    float ratio  = 1.0f;
-		float ratio = 1.0f;
+    float dw     = 0.0f;
-		float dw = 0.0f;
+    float dh     = 0.0f;
-		float dh = 0.0f;
+    float height = 0;
-		float height = 0;
+    float width  = 0;
-		float width = 0;
+};
-	};
+}  // namespace seg
-}
+#endif  // SEGMENT_SIMPLE_COMMON_HPP
-#endif //SEGMENT_SIMPLE_COMMON_HPP

csrc/segment/simple/include/yolov8-seg.hpp

@@ -3,549 +3,356 @@
 //
 #ifndef SEGMENT_SIMPLE_YOLOV8_SEG_HPP
 #define SEGMENT_SIMPLE_YOLOV8_SEG_HPP
-#include <fstream>
-#include "common.hpp"
 #include "NvInferPlugin.h"
+#include "common.hpp"
+#include <fstream>
 
 using namespace seg;
 
-class YOLOv8_seg
+class YOLOv8_seg {
-{
 public:
-	explicit YOLOv8_seg(const std::string& engine_file_path);
+    explicit YOLOv8_seg(const std::string& engine_file_path);
-	~YOLOv8_seg();
+    ~YOLOv8_seg();
-
+
-	void make_pipe(bool warmup = true);
+    void                 make_pipe(bool warmup = true);
-	void copy_from_Mat(const cv::Mat& image);
+    void                 copy_from_Mat(const cv::Mat& image);
-	void copy_from_Mat(const cv::Mat& image, cv::Size& size);
+    void                 copy_from_Mat(const cv::Mat& image, cv::Size& size);
-	void letterbox(
+    void                 letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size);
-		const cv::Mat& image,
+    void                 infer();
-		cv::Mat& out,
+    void                 postprocess(std::vector<Object>& objs,
-		cv::Size& size
+                                     float                score_thres  = 0.25f,
-	);
+                                     float                iou_thres    = 0.65f,
-	void infer();
+                                     int                  topk         = 100,
-	void postprocess(
+                                     int                  seg_channels = 32,
-		std::vector<Object>& objs,
+                                     int                  seg_h        = 160,
-		float score_thres = 0.25f,
+                                     int                  seg_w        = 160);
-		float iou_thres = 0.65f,
+    static void          draw_objects(const cv::Mat&                                image,
-		int topk = 100,
+                                      cv::Mat&                                      res,
-		int seg_channels = 32,
+                                      const std::vector<Object>&                    objs,
-		int seg_h = 160,
+                                      const std::vector<std::string>&               CLASS_NAMES,
-		int seg_w = 160
+                                      const std::vector<std::vector<unsigned int>>& COLORS,
-	);
+                                      const std::vector<std::vector<unsigned int>>& MASK_COLORS);
-	static void draw_objects(
+    int                  num_bindings;
-		const cv::Mat& image,
+    int                  num_inputs  = 0;
-		cv::Mat& res,
+    int                  num_outputs = 0;
-		const std::vector<Object>& objs,
+    std::vector<Binding> input_bindings;
-		const std::vector<std::string>& CLASS_NAMES,
+    std::vector<Binding> output_bindings;
-		const std::vector<std::vector<unsigned int>>& COLORS,
+    std::vector<void*>   host_ptrs;
-		const std::vector<std::vector<unsigned int>>& MASK_COLORS
+    std::vector<void*>   device_ptrs;
-	);
+
-	int num_bindings;
+    PreParam pparam;
-	int num_inputs = 0;
-	int num_outputs = 0;
-	std::vector<Binding> input_bindings;
-	std::vector<Binding> output_bindings;
-	std::vector<void*> host_ptrs;
-	std::vector<void*> device_ptrs;
-
-	PreParam pparam;
-private:
-	nvinfer1::ICudaEngine* engine = nullptr;
-	nvinfer1::IRuntime* runtime = nullptr;
-	nvinfer1::IExecutionContext* context = nullptr;
-	cudaStream_t stream = nullptr;
-	Logger gLogger{ nvinfer1::ILogger::Severity::kERROR };
 
+private:
+    nvinfer1::ICudaEngine*       engine  = nullptr;
+    nvinfer1::IRuntime*          runtime = nullptr;
+    nvinfer1::IExecutionContext* context = nullptr;
+    cudaStream_t                 stream  = nullptr;
+    Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
 };
 
 YOLOv8_seg::YOLOv8_seg(const std::string& engine_file_path)
 {
-	std::ifstream file(engine_file_path, std::ios::binary);
+    std::ifstream file(engine_file_path, std::ios::binary);
-	assert(file.good());
+    assert(file.good());
-	file.seekg(0, std::ios::end);
+    file.seekg(0, std::ios::end);
-	auto size = file.tellg();
+    auto size = file.tellg();
-	file.seekg(0, std::ios::beg);
+    file.seekg(0, std::ios::beg);
-	char* trtModelStream = new char[size];
+    char* trtModelStream = new char[size];
-	assert(trtModelStream);
+    assert(trtModelStream);
-	file.read(trtModelStream, size);
+    file.read(trtModelStream, size);
-	file.close();
+    file.close();
-	initLibNvInferPlugins(&this->gLogger, "");
+    initLibNvInferPlugins(&this->gLogger, "");
-	this->runtime = nvinfer1::createInferRuntime(this->gLogger);
+    this->runtime = nvinfer1::createInferRuntime(this->gLogger);
-	assert(this->runtime != nullptr);
+    assert(this->runtime != nullptr);
-
+
-	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
+    this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
-	assert(this->engine != nullptr);
+    assert(this->engine != nullptr);
-	delete[] trtModelStream;
+    delete[] trtModelStream;
-	this->context = this->engine->createExecutionContext();
+    this->context = this->engine->createExecutionContext();
-
+
-	assert(this->context != nullptr);
+    assert(this->context != nullptr);
-	cudaStreamCreate(&this->stream);
+    cudaStreamCreate(&this->stream);
-	this->num_bindings = this->engine->getNbBindings();
+    this->num_bindings = this->engine->getNbBindings();
-
+
-	for (int i = 0; i < this->num_bindings; ++i)
+    for (int i = 0; i < this->num_bindings; ++i) {
-	{
+        Binding            binding;
-		Binding binding;
+        nvinfer1::Dims     dims;
-		nvinfer1::Dims dims;
+        nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
-		nvinfer1::DataType dtype = this->engine->getBindingDataType(i);
+        std::string        name  = this->engine->getBindingName(i);
-		std::string name = this->engine->getBindingName(i);
+        binding.name             = name;
-		binding.name = name;
+        binding.dsize            = type_to_size(dtype);
-		binding.dsize = type_to_size(dtype);
+
-
+        bool IsInput = engine->bindingIsInput(i);
-		bool IsInput = engine->bindingIsInput(i);
+        if (IsInput) {
-		if (IsInput)
+            this->num_inputs += 1;
-		{
+            dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);
-			this->num_inputs += 1;
+            binding.size = get_size_by_dims(dims);
-			dims = this->engine->getProfileDimensions(
+            binding.dims = dims;
-				i,
+            this->input_bindings.push_back(binding);
-				0,
+            // set max opt shape
-				nvinfer1::OptProfileSelector::kMAX);
+            this->context->setBindingDimensions(i, dims);
-			binding.size = get_size_by_dims(dims);
+        }
-			binding.dims = dims;
+        else {
-			this->input_bindings.push_back(binding);
+            dims         = this->context->getBindingDimensions(i);
-			// set max opt shape
+            binding.size = get_size_by_dims(dims);
-			this->context->setBindingDimensions(i, dims);
+            binding.dims = dims;
-
+            this->output_bindings.push_back(binding);
-		}
+            this->num_outputs += 1;
-		else
+        }
-		{
+    }
-			dims = this->context->getBindingDimensions(i);
-			binding.size = get_size_by_dims(dims);
-			binding.dims = dims;
-			this->output_bindings.push_back(binding);
-			this->num_outputs += 1;
-		}
-	}
-
 }
 
 YOLOv8_seg::~YOLOv8_seg()
 {
-	this->context->destroy();
+    this->context->destroy();
-	this->engine->destroy();
+    this->engine->destroy();
-	this->runtime->destroy();
+    this->runtime->destroy();
-	cudaStreamDestroy(this->stream);
+    cudaStreamDestroy(this->stream);
-	for (auto& ptr : this->device_ptrs)
+    for (auto& ptr : this->device_ptrs) {
-	{
+        CHECK(cudaFree(ptr));
-		CHECK(cudaFree(ptr));
+    }
-	}
+
-
+    for (auto& ptr : this->host_ptrs) {
-	for (auto& ptr : this->host_ptrs)
+        CHECK(cudaFreeHost(ptr));
-	{
+    }
-		CHECK(cudaFreeHost(ptr));
-	}
 }
 
 void YOLOv8_seg::make_pipe(bool warmup)
 {
 
-	for (auto& bindings : this->input_bindings)
+    for (auto& bindings : this->input_bindings) {
-	{
+        void* d_ptr;
-		void* d_ptr;
+        CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream));
-		CHECK(cudaMallocAsync(
+        this->device_ptrs.push_back(d_ptr);
-			&d_ptr,
+    }
-			bindings.size * bindings.dsize,
+
-			this->stream)
+    for (auto& bindings : this->output_bindings) {
-		);
+        void * d_ptr, *h_ptr;
-		this->device_ptrs.push_back(d_ptr);
+        size_t size = bindings.size * bindings.dsize;
-	}
+        CHECK(cudaMallocAsync(&d_ptr, size, this->stream));
-
+        CHECK(cudaHostAlloc(&h_ptr, size, 0));
-	for (auto& bindings : this->output_bindings)
+        this->device_ptrs.push_back(d_ptr);
-	{
+        this->host_ptrs.push_back(h_ptr);
-		void* d_ptr, * h_ptr;
+    }
-		size_t size = bindings.size * bindings.dsize;
+
-		CHECK(cudaMallocAsync(
+    if (warmup) {
-			&d_ptr,
+        for (int i = 0; i < 10; i++) {
-			size,
+            for (auto& bindings : this->input_bindings) {
-			this->stream)
+                size_t size  = bindings.size * bindings.dsize;
-		);
+                void*  h_ptr = malloc(size);
-		CHECK(cudaHostAlloc(
+                memset(h_ptr, 0, size);
-			&h_ptr,
+                CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));
-			size,
+                free(h_ptr);
-			0)
+            }
-		);
+            this->infer();
-		this->device_ptrs.push_back(d_ptr);
+        }
-		this->host_ptrs.push_back(h_ptr);
+        printf("model warmup 10 times\n");
-	}
+    }
-
-	if (warmup)
-	{
-		for (int i = 0; i < 10; i++)
-		{
-			for (auto& bindings : this->input_bindings)
-			{
-				size_t size = bindings.size * bindings.dsize;
-				void* h_ptr = malloc(size);
-				memset(h_ptr, 0, size);
-				CHECK(cudaMemcpyAsync(
-					this->device_ptrs[0],
-					h_ptr,
-					size,
-					cudaMemcpyHostToDevice,
-					this->stream)
-				);
-				free(h_ptr);
-			}
-			this->infer();
-		}
-		printf("model warmup 10 times\n");
-
-	}
 }
 
-void YOLOv8_seg::letterbox(
+void YOLOv8_seg::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
-	const cv::Mat& image,
-	cv::Mat& out,
-	cv::Size& size
-)
 {
-	const float inp_h = size.height;
+    const float inp_h  = size.height;
-	const float inp_w = size.width;
+    const float inp_w  = size.width;
-	float height = image.rows;
+    float       height = image.rows;
-	float width = image.cols;
+    float       width  = image.cols;
-
+
-	float r = std::min(inp_h / height, inp_w / width);
+    float r    = std::min(inp_h / height, inp_w / width);
-	int padw = std::round(width * r);
+    int   padw = std::round(width * r);
-	int padh = std::round(height * r);
+    int   padh = std::round(height * r);
-
+
-	cv::Mat tmp;
+    cv::Mat tmp;
-	if ((int)width != padw || (int)height != padh)
+    if ((int)width != padw || (int)height != padh) {
-	{
+        cv::resize(image, tmp, cv::Size(padw, padh));
-		cv::resize(
+    }
-			image,
+    else {
-			tmp,
+        tmp = image.clone();
-			cv::Size(padw, padh)
+    }
-		);
+
-	}
+    float dw = inp_w - padw;
-	else
+    float dh = inp_h - padh;
-	{
+
-		tmp = image.clone();
+    dw /= 2.0f;
-	}
+    dh /= 2.0f;
-
+    int top    = int(std::round(dh - 0.1f));
-	float dw = inp_w - padw;
+    int bottom = int(std::round(dh + 0.1f));
-	float dh = inp_h - padh;
+    int left   = int(std::round(dw - 0.1f));
-
+    int right  = int(std::round(dw + 0.1f));
-	dw /= 2.0f;
+
-	dh /= 2.0f;
+    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});
-	int top = int(std::round(dh - 0.1f));
+
-	int bottom = int(std::round(dh + 0.1f));
+    cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
-	int left = int(std::round(dw - 0.1f));
+    this->pparam.ratio  = 1 / r;
-	int right = int(std::round(dw + 0.1f));
+    this->pparam.dw     = dw;
-
+    this->pparam.dh     = dh;
-	cv::copyMakeBorder(
+    this->pparam.height = height;
-		tmp,
+    this->pparam.width  = width;
-		tmp,
+    ;
-		top,
-		bottom,
-		left,
-		right,
-		cv::BORDER_CONSTANT,
-		{ 114, 114, 114 }
-	);
-
-	cv::dnn::blobFromImage(tmp,
-		out,
-		1 / 255.f,
-		cv::Size(),
-		cv::Scalar(0, 0, 0),
-		true,
-		false,
-		CV_32F
-	);
-	this->pparam.ratio = 1 / r;
-	this->pparam.dw = dw;
-	this->pparam.dh = dh;
-	this->pparam.height = height;
-	this->pparam.width = width;;
 }
 
 void YOLOv8_seg::copy_from_Mat(const cv::Mat& image)
 {
-	cv::Mat nchw;
+    cv::Mat  nchw;
-	auto& in_binding = this->input_bindings[0];
+    auto&    in_binding = this->input_bindings[0];
-	auto width = in_binding.dims.d[3];
+    auto     width      = in_binding.dims.d[3];
-	auto height = in_binding.dims.d[2];
+    auto     height     = in_binding.dims.d[2];
-	cv::Size size{ width, height };
+    cv::Size size{width, height};
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+
-		nchw,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});
-		size
+
-	);
+    CHECK(cudaMemcpyAsync(
-
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{
-				4,
-				{ 1, 3, height, width }
-			}
-	);
-
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8_seg::copy_from_Mat(const cv::Mat& image, cv::Size& size)
 {
-	cv::Mat nchw;
+    cv::Mat nchw;
-	this->letterbox(
+    this->letterbox(image, nchw, size);
-		image,
+    this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});
-		nchw,
+    CHECK(cudaMemcpyAsync(
-		size
+        this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
-	);
-	this->context->setBindingDimensions(
-		0,
-		nvinfer1::Dims
-			{ 4,
-			  { 1, 3, size.height, size.width }
-			}
-	);
-	CHECK(cudaMemcpyAsync(
-		this->device_ptrs[0],
-		nchw.ptr<float>(),
-		nchw.total() * nchw.elemSize(),
-		cudaMemcpyHostToDevice,
-		this->stream)
-	);
 }
 
 void YOLOv8_seg::infer()
 {
 
-	this->context->enqueueV2(
+    this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);
-		this->device_ptrs.data(),
+    for (int i = 0; i < this->num_outputs; i++) {
-		this->stream,
+        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		nullptr
+        CHECK(cudaMemcpyAsync(
-	);
+            this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));
-	for (int i = 0; i < this->num_outputs; i++)
+    }
-	{
+    cudaStreamSynchronize(this->stream);
-		size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
-		CHECK(cudaMemcpyAsync(this->host_ptrs[i],
-			this->device_ptrs[i + this->num_inputs],
-			osize,
-			cudaMemcpyDeviceToHost,
-			this->stream)
-		);
-
-	}
-	cudaStreamSynchronize(this->stream);
-
 }
 
-void YOLOv8_seg::postprocess(std::vector<Object>& objs,
+void YOLOv8_seg::postprocess(
-	float score_thres,
+    std::vector<Object>& objs, float score_thres, float iou_thres, int topk, int seg_channels, int seg_h, int seg_w)
-	float iou_thres,
-	int topk,
-	int seg_channels,
-	int seg_h,
-	int seg_w
-)
 {
-	objs.clear();
+    objs.clear();
-	auto input_h = this->input_bindings[0].dims.d[2];
+    auto input_h      = this->input_bindings[0].dims.d[2];
-	auto input_w = this->input_bindings[0].dims.d[3];
+    auto input_w      = this->input_bindings[0].dims.d[3];
-	auto num_anchors = this->output_bindings[0].dims.d[1];
+    auto num_anchors  = this->output_bindings[0].dims.d[1];
-	auto num_channels = this->output_bindings[0].dims.d[2];
+    auto num_channels = this->output_bindings[0].dims.d[2];
-
+
-	auto& dw = this->pparam.dw;
+    auto& dw     = this->pparam.dw;
-	auto& dh = this->pparam.dh;
+    auto& dh     = this->pparam.dh;
-	auto& width = this->pparam.width;
+    auto& width  = this->pparam.width;
-	auto& height = this->pparam.height;
+    auto& height = this->pparam.height;
-	auto& ratio = this->pparam.ratio;
+    auto& ratio  = this->pparam.ratio;
-
+
-	auto* output = static_cast<float*>(this->host_ptrs[0]);
+    auto*   output = static_cast<float*>(this->host_ptrs[0]);
-	cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F,
+    cv::Mat protos = cv::Mat(seg_channels, seg_h * seg_w, CV_32F, static_cast<float*>(this->host_ptrs[1]));
-		static_cast<float*>(this->host_ptrs[1]));
+
-
+    std::vector<int>      labels;
-	std::vector<int> labels;
+    std::vector<float>    scores;
-	std::vector<float> scores;
+    std::vector<cv::Rect> bboxes;
-	std::vector<cv::Rect> bboxes;
+    std::vector<cv::Mat>  mask_confs;
-	std::vector<cv::Mat> mask_confs;
+    std::vector<int>      indices;
-	std::vector<int> indices;
+
-
+    for (int i = 0; i < num_anchors; i++) {
-	for (int i = 0; i < num_anchors; i++)
+        float* ptr   = output + i * num_channels;
-	{
+        float  score = *(ptr + 4);
-		float* ptr = output + i * num_channels;
+        if (score > score_thres) {
-		float score = *(ptr + 4);
+            float x0 = *ptr++ - dw;
-		if (score > score_thres)
+            float y0 = *ptr++ - dh;
-		{
+            float x1 = *ptr++ - dw;
-			float x0 = *ptr++ - dw;
+            float y1 = *ptr++ - dh;
-			float y0 = *ptr++ - dh;
+
-			float x1 = *ptr++ - dw;
+            x0 = clamp(x0 * ratio, 0.f, width);
-			float y1 = *ptr++ - dh;
+            y0 = clamp(y0 * ratio, 0.f, height);
-
+            x1 = clamp(x1 * ratio, 0.f, width);
-			x0 = clamp(x0 * ratio, 0.f, width);
+            y1 = clamp(y1 * ratio, 0.f, height);
-			y0 = clamp(y0 * ratio, 0.f, height);
+
-			x1 = clamp(x1 * ratio, 0.f, width);
+            int     label     = *(++ptr);
-			y1 = clamp(y1 * ratio, 0.f, height);
+            cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, ++ptr);
-
+            mask_confs.push_back(mask_conf);
-			int label = *(++ptr);
+            labels.push_back(label);
-			cv::Mat mask_conf = cv::Mat(1, seg_channels, CV_32F, ++ptr);
+            scores.push_back(score);
-			mask_confs.push_back(mask_conf);
+            bboxes.push_back(cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0));
-			labels.push_back(label);
+        }
-			scores.push_back(score);
+    }
-			bboxes.push_back(cv::Rect_<float>(x0, y0, x1 - x0, y1 - y0));
-
-		}
-	}
 
 #if defined(BATCHED_NMS)
-	cv::dnn::NMSBoxesBatched(
+    cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		labels,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #else
-	cv::dnn::NMSBoxes(
+    cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
-		bboxes,
-		scores,
-		score_thres,
-		iou_thres,
-		indices
-	);
 #endif
 
-	cv::Mat masks;
+    cv::Mat masks;
-	int cnt = 0;
+    int     cnt = 0;
-	for (auto& i : indices)
+    for (auto& i : indices) {
-	{
+        if (cnt >= topk) {
-		if (cnt >= topk)
+            break;
-		{
+        }
-			break;
+        cv::Rect tmp = bboxes[i];
-		}
+        Object   obj;
-		cv::Rect tmp = bboxes[i];
+        obj.label = labels[i];
-		Object obj;
+        obj.rect  = tmp;
-		obj.label = labels[i];
+        obj.prob  = scores[i];
-		obj.rect = tmp;
+        masks.push_back(mask_confs[i]);
-		obj.prob = scores[i];
+        objs.push_back(obj);
-		masks.push_back(mask_confs[i]);
+        cnt += 1;
-		objs.push_back(obj);
+    }
-		cnt += 1;
+
-	}
+    if (masks.empty()) {
-
+        // masks is empty
-	if(masks.empty())
+    }
-	{
+    else {
-		//masks is empty
+        cv::Mat matmulRes = (masks * protos).t();
-	}
+        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_w, seg_h});
-	else
+
-	{
+        std::vector<cv::Mat> maskChannels;
-		cv::Mat matmulRes = (masks * protos).t();
+        cv::split(maskMat, maskChannels);
-		cv::Mat maskMat = matmulRes.reshape(indices.size(), { seg_w, seg_h });
+        int scale_dw = dw / input_w * seg_w;
-
+        int scale_dh = dh / input_h * seg_h;
-		std::vector<cv::Mat> maskChannels;
+
-		cv::split(maskMat, maskChannels);
+        cv::Rect roi(scale_dw, scale_dh, seg_w - 2 * scale_dw, seg_h - 2 * scale_dh);
-		int scale_dw = dw / input_w * seg_w;
+
-		int scale_dh = dh / input_h * seg_h;
+        for (int i = 0; i < indices.size(); i++) {
-
+            cv::Mat dest, mask;
-		cv::Rect roi(
+            cv::exp(-maskChannels[i], dest);
-			scale_dw,
+            dest = 1.0 / (1.0 + dest);
-			scale_dh,
+            dest = dest(roi);
-			seg_w - 2 * scale_dw,
+            cv::resize(dest, mask, cv::Size((int)width, (int)height), cv::INTER_LINEAR);
-			seg_h - 2 * scale_dh);
+            objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-
+        }
-		for (int i = 0; i < indices.size(); i++)
+    }
-		{
-			cv::Mat dest, mask;
-			cv::exp(-maskChannels[i], dest);
-			dest = 1.0 / (1.0 + dest);
-			dest = dest(roi);
-			cv::resize(
-				dest,
-				mask,
-				cv::Size((int)width, (int)height),
-				cv::INTER_LINEAR
-			);
-			objs[i].boxMask = mask(objs[i].rect) > 0.5f;
-		}
-	}
 }
 
-void YOLOv8_seg::draw_objects(const cv::Mat& image,
+void YOLOv8_seg::draw_objects(const cv::Mat&                                image,
-	cv::Mat& res,
+                              cv::Mat&                                      res,
-	const std::vector<Object>& objs,
+                              const std::vector<Object>&                    objs,
-	const std::vector<std::string>& CLASS_NAMES,
+                              const std::vector<std::string>&               CLASS_NAMES,
-	const std::vector<std::vector<unsigned int>>& COLORS,
+                              const std::vector<std::vector<unsigned int>>& COLORS,
-	const std::vector<std::vector<unsigned int>>& MASK_COLORS
+                              const std::vector<std::vector<unsigned int>>& MASK_COLORS)
-)
 {
-	res = image.clone();
+    res          = image.clone();
-	cv::Mat mask = image.clone();
+    cv::Mat mask = image.clone();
-	for (auto& obj : objs)
+    for (auto& obj : objs) {
-	{
+        int        idx   = obj.label;
-		int idx = obj.label;
+        cv::Scalar color = cv::Scalar(COLORS[idx][0], COLORS[idx][1], COLORS[idx][2]);
-		cv::Scalar color = cv::Scalar(
+        cv::Scalar mask_color =
-			COLORS[idx][0],
+            cv::Scalar(MASK_COLORS[idx % 20][0], MASK_COLORS[idx % 20][1], MASK_COLORS[idx % 20][2]);
-			COLORS[idx][1],
+        cv::rectangle(res, obj.rect, color, 2);
-			COLORS[idx][2]
+
-		);
+        char text[256];
-		cv::Scalar mask_color = cv::Scalar(
+        sprintf(text, "%s %.1f%%", CLASS_NAMES[idx].c_str(), obj.prob * 100);
-			MASK_COLORS[idx % 20][0],
+        mask(obj.rect).setTo(mask_color, obj.boxMask);
-			MASK_COLORS[idx % 20][1],
+
-			MASK_COLORS[idx % 20][2]
+        int      baseLine   = 0;
-		);
+        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);
-		cv::rectangle(
+
-			res,
+        int x = (int)obj.rect.x;
-			obj.rect,
+        int y = (int)obj.rect.y + 1;
-			color,
+
-			2
+        if (y > res.rows)
-		);
+            y = res.rows;
-
+
-		char text[256];
+        cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);
-		sprintf(
+
-			text,
+        cv::putText(res, text, cv::Point(x, y + label_size.height), cv::FONT_HERSHEY_SIMPLEX, 0.4, {255, 255, 255}, 1);
-			"%s %.1f%%",
+    }
-			CLASS_NAMES[idx].c_str(),
+    cv::addWeighted(res, 0.5, mask, 0.8, 1, res);
-			obj.prob * 100
-		);
-		mask(obj.rect).setTo(mask_color, obj.boxMask);
-
-		int baseLine = 0;
-		cv::Size label_size = cv::getTextSize(
-			text,
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			1,
-			&baseLine
-		);
-
-		int x = (int)obj.rect.x;
-		int y = (int)obj.rect.y + 1;
-
-		if (y > res.rows)
-			y = res.rows;
-
-		cv::rectangle(
-			res,
-			cv::Rect(x, y, label_size.width, label_size.height + baseLine),
-			{ 0, 0, 255 },
-			-1
-		);
-
-		cv::putText(
-			res,
-			text,
-			cv::Point(x, y + label_size.height),
-			cv::FONT_HERSHEY_SIMPLEX,
-			0.4,
-			{ 255, 255, 255 },
-			1
-		);
-	}
-	cv::addWeighted(
-		res,
-		0.5,
-		mask,
-		0.8,
-		1,
-		res
-	);
 }
-#endif //SEGMENT_SIMPLE_YOLOV8_SEG_HPP
+#endif  // SEGMENT_SIMPLE_YOLOV8_SEG_HPP

csrc/segment/simple/main.cpp

@@ -2,177 +2,131 @@
 // Created by ubuntu on 1/20/23.
 //
 #include "chrono"
-#include "yolov8-seg.hpp"
 #include "opencv2/opencv.hpp"
+#include "yolov8-seg.hpp"
 
 const std::vector<std::string> CLASS_NAMES = {
-	"person", "bicycle", "car", "motorcycle", "airplane", "bus",
+    "person",         "bicycle",    "car",           "motorcycle",    "airplane",     "bus",           "train",
-	"train", "truck", "boat", "traffic light", "fire hydrant",
+    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",    "parking meter", "bench",
-	"stop sign", "parking meter", "bench", "bird", "cat",
+    "bird",           "cat",        "dog",           "horse",         "sheep",        "cow",           "elephant",
-	"dog", "horse", "sheep", "cow", "elephant",
+    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",     "handbag",       "tie",
-	"bear", "zebra", "giraffe", "backpack", "umbrella",
+    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball",  "kite",          "baseball bat",
-	"handbag", "tie", "suitcase", "frisbee", "skis",
+    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",       "wine glass",    "cup",
-	"snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "fork",           "knife",      "spoon",         "bowl",          "banana",       "apple",         "sandwich",
-	"skateboard", "surfboard", "tennis racket", "bottle", "wine glass",
+    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",        "donut",         "cake",
-	"cup", "fork", "knife", "spoon", "bowl",
+    "chair",          "couch",      "potted plant",  "bed",           "dining table", "toilet",        "tv",
-	"banana", "apple", "sandwich", "orange", "broccoli",
+    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",   "microwave",     "oven",
-	"carrot", "hot dog", "pizza", "donut", "cake",
+    "toaster",        "sink",       "refrigerator",  "book",          "clock",        "vase",          "scissors",
-	"chair", "couch", "potted plant", "bed", "dining table",
+    "teddy bear",     "hair drier", "toothbrush"};
-	"toilet", "tv", "laptop", "mouse", "remote",
-	"keyboard", "cell phone", "microwave", "oven",
-	"toaster", "sink", "refrigerator", "book", "clock", "vase",
-	"scissors", "teddy bear", "hair drier", "toothbrush" };
 
 const std::vector<std::vector<unsigned int>> COLORS = {
-	{ 0, 114, 189 }, { 217, 83, 25 }, { 237, 177, 32 },
+    {0, 114, 189},   {217, 83, 25},   {237, 177, 32},  {126, 47, 142},  {119, 172, 48},  {77, 190, 238},
-	{ 126, 47, 142 }, { 119, 172, 48 }, { 77, 190, 238 },
+    {162, 20, 47},   {76, 76, 76},    {153, 153, 153}, {255, 0, 0},     {255, 128, 0},   {191, 191, 0},
-	{ 162, 20, 47 }, { 76, 76, 76 }, { 153, 153, 153 },
+    {0, 255, 0},     {0, 0, 255},     {170, 0, 255},   {85, 85, 0},     {85, 170, 0},    {85, 255, 0},
-	{ 255, 0, 0 }, { 255, 128, 0 }, { 191, 191, 0 },
+    {170, 85, 0},    {170, 170, 0},   {170, 255, 0},   {255, 85, 0},    {255, 170, 0},   {255, 255, 0},
-	{ 0, 255, 0 }, { 0, 0, 255 }, { 170, 0, 255 },
+    {0, 85, 128},    {0, 170, 128},   {0, 255, 128},   {85, 0, 128},    {85, 85, 128},   {85, 170, 128},
-	{ 85, 85, 0 }, { 85, 170, 0 }, { 85, 255, 0 },
+    {85, 255, 128},  {170, 0, 128},   {170, 85, 128},  {170, 170, 128}, {170, 255, 128}, {255, 0, 128},
-	{ 170, 85, 0 }, { 170, 170, 0 }, { 170, 255, 0 },
+    {255, 85, 128},  {255, 170, 128}, {255, 255, 128}, {0, 85, 255},    {0, 170, 255},   {0, 255, 255},
-	{ 255, 85, 0 }, { 255, 170, 0 }, { 255, 255, 0 },
+    {85, 0, 255},    {85, 85, 255},   {85, 170, 255},  {85, 255, 255},  {170, 0, 255},   {170, 85, 255},
-	{ 0, 85, 128 }, { 0, 170, 128 }, { 0, 255, 128 },
+    {170, 170, 255}, {170, 255, 255}, {255, 0, 255},   {255, 85, 255},  {255, 170, 255}, {85, 0, 0},
-	{ 85, 0, 128 }, { 85, 85, 128 }, { 85, 170, 128 },
+    {128, 0, 0},     {170, 0, 0},     {212, 0, 0},     {255, 0, 0},     {0, 43, 0},      {0, 85, 0},
-	{ 85, 255, 128 }, { 170, 0, 128 }, { 170, 85, 128 },
+    {0, 128, 0},     {0, 170, 0},     {0, 212, 0},     {0, 255, 0},     {0, 0, 43},      {0, 0, 85},
-	{ 170, 170, 128 }, { 170, 255, 128 }, { 255, 0, 128 },
+    {0, 0, 128},     {0, 0, 170},     {0, 0, 212},     {0, 0, 255},     {0, 0, 0},       {36, 36, 36},
-	{ 255, 85, 128 }, { 255, 170, 128 }, { 255, 255, 128 },
+    {73, 73, 73},    {109, 109, 109}, {146, 146, 146}, {182, 182, 182}, {219, 219, 219}, {0, 114, 189},
-	{ 0, 85, 255 }, { 0, 170, 255 }, { 0, 255, 255 },
+    {80, 183, 189},  {128, 128, 0}};
-	{ 85, 0, 255 }, { 85, 85, 255 }, { 85, 170, 255 },
-	{ 85, 255, 255 }, { 170, 0, 255 }, { 170, 85, 255 },
-	{ 170, 170, 255 }, { 170, 255, 255 }, { 255, 0, 255 },
-	{ 255, 85, 255 }, { 255, 170, 255 }, { 85, 0, 0 },
-	{ 128, 0, 0 }, { 170, 0, 0 }, { 212, 0, 0 },
-	{ 255, 0, 0 }, { 0, 43, 0 }, { 0, 85, 0 },
-	{ 0, 128, 0 }, { 0, 170, 0 }, { 0, 212, 0 },
-	{ 0, 255, 0 }, { 0, 0, 43 }, { 0, 0, 85 },
-	{ 0, 0, 128 }, { 0, 0, 170 }, { 0, 0, 212 },
-	{ 0, 0, 255 }, { 0, 0, 0 }, { 36, 36, 36 },
-	{ 73, 73, 73 }, { 109, 109, 109 }, { 146, 146, 146 },
-	{ 182, 182, 182 }, { 219, 219, 219 }, { 0, 114, 189 },
-	{ 80, 183, 189 }, { 128, 128, 0 }
-};
 
 const std::vector<std::vector<unsigned int>> MASK_COLORS = {
-	{ 255, 56, 56 }, { 255, 157, 151 }, { 255, 112, 31 },
+    {255, 56, 56},  {255, 157, 151}, {255, 112, 31}, {255, 178, 29}, {207, 210, 49},  {72, 249, 10}, {146, 204, 23},
-	{ 255, 178, 29 }, { 207, 210, 49 }, { 72, 249, 10 },
+    {61, 219, 134}, {26, 147, 52},   {0, 212, 187},  {44, 153, 168}, {0, 194, 255},   {52, 69, 147}, {100, 115, 255},
-	{ 146, 204, 23 }, { 61, 219, 134 }, { 26, 147, 52 },
+    {0, 24, 236},   {132, 56, 255},  {82, 0, 133},   {203, 56, 255}, {255, 149, 200}, {255, 55, 199}};
-	{ 0, 212, 187 }, { 44, 153, 168 }, { 0, 194, 255 },
-	{ 52, 69, 147 }, { 100, 115, 255 }, { 0, 24, 236 },
-	{ 132, 56, 255 }, { 82, 0, 133 }, { 203, 56, 255 },
-	{ 255, 149, 200 }, { 255, 55, 199 }
-};
 
 int main(int argc, char** argv)
 {
-	// cuda:0
+    // cuda:0
-	cudaSetDevice(0);
+    cudaSetDevice(0);
 
-	const std::string engine_file_path{ argv[1] };
+    const std::string engine_file_path{argv[1]};
-	const std::string path{ argv[2] };
+    const std::string path{argv[2]};
 
-	std::vector<std::string> imagePathList;
+    std::vector<std::string> imagePathList;
-	bool isVideo{ false };
+    bool                     isVideo{false};
 
-	assert(argc == 3);
+    assert(argc == 3);
 
-	auto yolov8 = new YOLOv8_seg(engine_file_path);
+    auto yolov8 = new YOLOv8_seg(engine_file_path);
-	yolov8->make_pipe(true);
+    yolov8->make_pipe(true);
 
-	if (IsFile(path))
+    if (IsFile(path)) {
-	{
+        std::string suffix = path.substr(path.find_last_of('.') + 1);
-		std::string suffix = path.substr(path.find_last_of('.') + 1);
+        if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {
-		if (
+            imagePathList.push_back(path);
-			suffix == "jpg" ||
+        }
-				suffix == "jpeg" ||
+        else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"
-				suffix == "png"
+                 || suffix == "mkv") {
-			)
+            isVideo = true;
-		{
+        }
-			imagePathList.push_back(path);
+        else {
-		}
+            printf("suffix %s is wrong !!!\n", suffix.c_str());
-		else if (
+            std::abort();
-			suffix == "mp4" ||
+        }
-				suffix == "avi" ||
+    }
-				suffix == "m4v" ||
+    else if (IsFolder(path)) {
-				suffix == "mpeg" ||
+        cv::glob(path + "/*.jpg", imagePathList);
-				suffix == "mov" ||
+    }
-				suffix == "mkv"
-			)
-		{
-			isVideo = true;
-		}
-		else
-		{
-			printf("suffix %s is wrong !!!\n", suffix.c_str());
-			std::abort();
-		}
-	}
-	else if (IsFolder(path))
-	{
-		cv::glob(path + "/*.jpg", imagePathList);
-	}
 
-	cv::Mat res, image;
+    cv::Mat  res, image;
-	cv::Size size = cv::Size{ 640, 640 };
+    cv::Size size         = cv::Size{640, 640};
-	int topk = 100;
+    int      topk         = 100;
-	int seg_h = 160;
+    int      seg_h        = 160;
-	int seg_w = 160;
+    int      seg_w        = 160;
-	int seg_channels = 32;
+    int      seg_channels = 32;
-	float score_thres = 0.25f;
+    float    score_thres  = 0.25f;
-	float iou_thres = 0.65f;
+    float    iou_thres    = 0.65f;
 
-	std::vector<Object> objs;
+    std::vector<Object> objs;
 
-	cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
+    cv::namedWindow("result", cv::WINDOW_AUTOSIZE);
 
-	if (isVideo)
+    if (isVideo) {
-	{
+        cv::VideoCapture cap(path);
-		cv::VideoCapture cap(path);
 
-		if (!cap.isOpened())
+        if (!cap.isOpened()) {
-		{
+            printf("can not open %s\n", path.c_str());
-			printf("can not open %s\n", path.c_str());
+            return -1;
-			return -1;
+        }
-		}
+        while (cap.read(image)) {
-		while (cap.read(image))
+            objs.clear();
-		{
+            yolov8->copy_from_Mat(image, size);
-			objs.clear();
+            auto start = std::chrono::system_clock::now();
-			yolov8->copy_from_Mat(image, size);
+            yolov8->infer();
-			auto start = std::chrono::system_clock::now();
+            auto end = std::chrono::system_clock::now();
-			yolov8->infer();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
-			auto end = std::chrono::system_clock::now();
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
+            printf("cost %2.4lf ms\n", tc);
-			auto tc = (double)
+            cv::imshow("result", res);
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+            if (cv::waitKey(10) == 'q') {
-			printf("cost %2.4lf ms\n", tc);
+                break;
-			cv::imshow("result", res);
+            }
-			if (cv::waitKey(10) == 'q')
+        }
-			{
+    }
-				break;
+    else {
-			}
+        for (auto& path : imagePathList) {
-		}
+            objs.clear();
-	}
+            image = cv::imread(path);
-	else
+            yolov8->copy_from_Mat(image, size);
-	{
+            auto start = std::chrono::system_clock::now();
-		for (auto& path : imagePathList)
+            yolov8->infer();
-		{
+            auto end = std::chrono::system_clock::now();
-			objs.clear();
+            yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
-			image = cv::imread(path);
+            yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
-			yolov8->copy_from_Mat(image, size);
+            auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
-			auto start = std::chrono::system_clock::now();
+            printf("cost %2.4lf ms\n", tc);
-			yolov8->infer();
+            cv::imshow("result", res);
-			auto end = std::chrono::system_clock::now();
+            cv::waitKey(0);
-			yolov8->postprocess(objs, score_thres, iou_thres, topk, seg_channels, seg_h, seg_w);
+        }
-			yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS, MASK_COLORS);
+    }
-			auto tc = (double)
+    cv::destroyAllWindows();
-				std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;
+    delete yolov8;
-			printf("cost %2.4lf ms\n", tc);
+    return 0;
-			cv::imshow("result", res);
-			cv::waitKey(0);
-		}
-	}
-	cv::destroyAllWindows();
-	delete yolov8;
-	return 0;
 }

docs/Jetson.md

@@ -3,6 +3,7 @@
 Only test on `Jetson-NX 4GB`
 
 ENVS:
+
 - Jetpack 4.6.3
 - CUDA-10.2
 - CUDNN-8.2.1
@@ -20,7 +21,8 @@ If you have other environment-related issues, please discuss in issue.
 `yolov8s.pt` is your trained pytorch model, or the official pre-trained model.
 
 Do not use any model other than pytorch model.
-Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on jetson.
+Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on
+jetson.
 
 ***!!! Please use the PC to execute the following script !!!***
 
@@ -79,7 +81,8 @@ Usage:
 `yolov8s-seg.pt` is your trained pytorch model, or the official pre-trained model.
 
 Do not use any model other than pytorch model.
-Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on jetson.
+Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on
+jetson.
 
 ***!!! Please use the PC to execute the following script !!!***
 
@@ -106,7 +109,8 @@ Here is a demo: [`csrc/jetson/segment`](../csrc/jetson/segment) .
 
 #### Build:
 
-Please modify `CLASS_NAMES` and `COLORS` and postprocess parameters in [`main.cpp`](../csrc/jetson/segment/main.cpp) for yourself.
+Please modify `CLASS_NAMES` and `COLORS` and postprocess parameters in [`main.cpp`](../csrc/jetson/segment/main.cpp) for
+yourself.
 
 ```c++
 int topk = 100;
@@ -140,8 +144,6 @@ Usage:
 ./yolov8-seg yolov8s-seg.engine data/test.mp4 # the video path
 ```
 
-
-
 ## Normal Posture
 
 ### 1. Export Posture Normal ONNX
@@ -149,7 +151,8 @@ Usage:
 `yolov8s-pose.pt` is your trained pytorch model, or the official pre-trained model.
 
 Do not use any model other than pytorch model.
-Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on jetson.
+Do not use [`build.py`](../build.py) to export engine if you don't know how to install pytorch and other environments on
+jetson.
 
 ***!!! Please use the PC to execute the following script !!!***
 
@@ -176,7 +179,8 @@ Here is a demo: [`csrc/jetson/pose`](../csrc/jetson/pose) .
 
 #### Build:
 
-Please modify `KPS_COLORS` and `SKELETON` and `LIMB_COLORS` and postprocess parameters in [`main.cpp`](../csrc/jetson/pose/main.cpp) for yourself.
+Please modify `KPS_COLORS` and `SKELETON` and `LIMB_COLORS` and postprocess parameters
+in [`main.cpp`](../csrc/jetson/pose/main.cpp) for yourself.
 
 ```c++
 int topk = 100;

docs/Normal.md

@@ -2,23 +2,39 @@
 
 ## Export TensorRT Engine
 
-### 1. Python script
+### 1. ONNX -> TensorRT
 
-Usage:
+You can export your onnx model by `ultralytics` API.
 
-```python
+``` shell
+yolo export model=yolov8s.pt format=onnx opset=11 simplify=True
+```
 
+or run this python script:
+
+```python
 from ultralytics import YOLO
 
 # Load a model
 model = YOLO("yolov8s.pt")  # load a pretrained model (recommended for training)
-success = model.export(format="engine", device=0)  # export the model to engine format
+success = model.export(format="onnx", opset=11, simplify=True)  # export the model to onnx format
 assert success
 ```
 
-After executing the above script, you will get an engine named `yolov8s.engine` .
+Then build engine by Trtexec Tools.
+
+You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools.
+
+Usage:
+
+``` shell
+/usr/src/tensorrt/bin/trtexec \
+--onnx=yolov8s.onnx \
+--saveEngine=yolov8s.engine \
+--fp16
+```
 
-### 2. CLI tools
+### 2. Direct to TensorRT (NOT RECOMMAND!!)
 
 Usage:
 
@@ -26,7 +42,19 @@ Usage:
 yolo export model=yolov8s.pt format=engine device=0
 ```
 
-After executing the above command, you will get an engine named `yolov8s.engine` too.
+or run python script:
+
+```python
+
+from ultralytics import YOLO
+
+# Load a model
+model = YOLO("yolov8s.pt")  # load a pretrained model (recommended for training)
+success = model.export(format="engine", device=0)  # export the model to engine format
+assert success
+```
+
+After executing the above script, you will get an engine named `yolov8s.engine` .
 
 ## Inference with c++
 
@@ -34,9 +62,11 @@ You can infer with c++ in [`csrc/detect/normal`](../csrc/detect/normal) .
 
 ### Build:
 
-Please set you own librarys in [`CMakeLists.txt`](../csrc/detect/normal/CMakeLists.txt) and modify `CLASS_NAMES` and `COLORS` in [`main.cpp`](../csrc/detect/normal/main.cpp).
+Please set you own librarys in [`CMakeLists.txt`](../csrc/detect/normal/CMakeLists.txt) and modify `CLASS_NAMES`
+and `COLORS` in [`main.cpp`](../csrc/detect/normal/main.cpp).
 
-Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk` in [`main.cpp`](../csrc/detect/normal/main.cpp).
+Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk`
+in [`main.cpp`](../csrc/detect/normal/main.cpp).
 
 ```c++
 int num_labels = 80;

docs/Pose.md

@@ -1,7 +1,7 @@
 # YOLOv8-pose Model with TensorRT
 
 The yolov8-pose model conversion route is :
-    YOLOv8 PyTorch model -> ONNX -> TensorRT Engine
+YOLOv8 PyTorch model -> ONNX -> TensorRT Engine
 
 ***Notice !!!*** We don't support TensorRT API building !!!
 
@@ -9,10 +9,48 @@ The yolov8-pose model conversion route is :
 
 You can leave this repo and use the original `ultralytics` repo for onnx export.
 
-### 1. Python script
+### 1. ONNX -> TensorRT
+
+You can export your onnx model by `ultralytics` API.
+
+``` shell
+yolo export model=yolov8s-pose.pt format=onnx opset=11 simplify=True
+```
+
+or run this python script:
+
+```python
+from ultralytics import YOLO
+
+# Load a model
+model = YOLO("yolov8s-pose.pt")  # load a pretrained model (recommended for training)
+success = model.export(format="onnx", opset=11, simplify=True)  # export the model to onnx format
+assert success
+```
+
+Then build engine by Trtexec Tools.
+
+You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools.
+
+Usage:
+
+``` shell
+/usr/src/tensorrt/bin/trtexec \
+--onnx=yolov8s-pose.onnx \
+--saveEngine=yolov8s-pose.engine \
+--fp16
+```
+
+### 2. Direct to TensorRT (NOT RECOMMAND!!)
 
 Usage:
 
+```shell
+yolo export model=yolov8s-pose.pt format=engine device=0
+```
+
+or run python script:
+
 ```python
 from ultralytics import YOLO
 
@@ -24,15 +62,31 @@ assert success
 
 After executing the above script, you will get an engine named `yolov8s-pose.engine` .
 
-### 2. CLI tools
+# Inference
+
+## Infer with python script
+
+You can infer images with the engine by [`infer-pose.py`](../infer-pose.py) .
 
 Usage:
 
-```shell
+``` shell
-yolo export model=yolov8s-pose.pt format=engine device=0
+python3 infer-pose.py \
+--engine yolov8s-pose.engine \
+--imgs data \
+--show \
+--out-dir outputs \
+--device cuda:0
 ```
 
-After executing the above command, you will get an engine named `yolov8s-pose.engine` too.
+#### Description of all arguments
+
+- `--engine` : The Engine you export.
+- `--imgs` : The images path you want to detect.
+- `--show` : Whether to show detection results.
+- `--out-dir` : Where to save detection results images. It will not work when use `--show` flag.
+- `--device` : The CUDA deivce you use.
+- `--profile` : Profile the TensorRT engine.
 
 ## Inference with c++
 
@@ -40,9 +94,11 @@ You can infer with c++ in [`csrc/pose/normal`](../csrc/pose/normal) .
 
 ### Build:
 
-Please set you own librarys in [`CMakeLists.txt`](../csrc/pose/normal/CMakeLists.txt) and modify `KPS_COLORS` and `SKELETON` and  `LIMB_COLORS`  in [`main.cpp`](../csrc/pose/normal/main.cpp).
+Please set you own librarys in [`CMakeLists.txt`](../csrc/pose/normal/CMakeLists.txt) and modify `KPS_COLORS`
+and `SKELETON` and  `LIMB_COLORS`  in [`main.cpp`](../csrc/pose/normal/main.cpp).
 
-Besides, you can modify the postprocess parameters such as `score_thres` and `iou_thres` and `topk` in [`main.cpp`](../csrc/pose/normal/main.cpp).
+Besides, you can modify the postprocess parameters such as `score_thres` and `iou_thres` and `topk`
+in [`main.cpp`](../csrc/pose/normal/main.cpp).
 
 ```c++
 int topk = 100;

docs/Segment.md

@@ -1,7 +1,7 @@
 # YOLOv8-seg Model with TensorRT
 
 The yolov8-seg model conversion route is :
-    YOLOv8 PyTorch model -> ONNX -> TensorRT Engine
+YOLOv8 PyTorch model -> ONNX -> TensorRT Engine
 
 ***Notice !!!*** We don't support TensorRT API building !!!
 
@@ -96,7 +96,9 @@ You can infer segment engine with c++ in [`csrc/segment/simple`](../csrc/segment
 
 ### Build:
 
-Please set you own librarys in [`CMakeLists.txt`](../csrc/segment/simple/CMakeLists.txt) and modify you own config in [`main.cpp`](../csrc/segment/simple/main.cpp) such as `CLASS_NAMES`, `COLORS`, `MASK_COLORS` and postprocess parameters .
+Please set you own librarys in [`CMakeLists.txt`](../csrc/segment/simple/CMakeLists.txt) and modify you own config
+in [`main.cpp`](../csrc/segment/simple/main.cpp) such as `CLASS_NAMES`, `COLORS`, `MASK_COLORS` and postprocess
+parameters .
 
 ```c++
 int topk = 100;
@@ -119,7 +121,8 @@ cd ${root}
 
 ***Notice !!!***
 
-If you have build OpenCV(>=4.7.0) by yourself, it provides a new api [`cv::dnn::NMSBoxesBatched`](https://docs.opencv.org/4.x/d6/d0f/group__dnn.html#ga977aae09fbf7c804e003cfea1d4e928c) .
+If you have build OpenCV(>=4.7.0) by yourself, it provides a new
+api [`cv::dnn::NMSBoxesBatched`](https://docs.opencv.org/4.x/d6/d0f/group__dnn.html#ga977aae09fbf7c804e003cfea1d4e928c) .
 It is a gread api about efficient in-class nms . It will be used by default!
 
 ***!!!***
@@ -139,22 +142,39 @@ Usage:
 
 You can leave this repo and use the original `ultralytics` repo for onnx export.
 
-### 1. Python script
+### 1. ONNX -> TensorRT
 
-Usage:
+You can export your onnx model by `ultralytics` API.
+
+``` shell
+yolo export model=yolov8s-seg.pt format=onnx opset=11 simplify=True
+```
+
+or run this python script:
 
 ```python
 from ultralytics import YOLO
 
 # Load a model
 model = YOLO("yolov8s-seg.pt")  # load a pretrained model (recommended for training)
-success = model.export(format="engine", device=0)  # export the model to engine format
+success = model.export(format="onnx", opset=11, simplify=True)  # export the model to onnx format
 assert success
 ```
 
-After executing the above script, you will get an engine named `yolov8s-seg.engine` .
+Then build engine by Trtexec Tools.
 
-### 2. CLI tools
+You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools.
+
+Usage:
+
+``` shell
+/usr/src/tensorrt/bin/trtexec \
+--onnx=yolov8s-seg.onnx \
+--saveEngine=yolov8s-seg.engine \
+--fp16
+```
+
+### 2. Direct to TensorRT (NOT RECOMMAND!!)
 
 Usage:
 
@@ -162,7 +182,18 @@ Usage:
 yolo export model=yolov8s-seg.pt format=engine device=0
 ```
 
-After executing the above command, you will get an engine named `yolov8s-seg.engine` too.
+or run python script:
+
+```python
+from ultralytics import YOLO
+
+# Load a model
+model = YOLO("yolov8s-seg.pt")  # load a pretrained model (recommended for training)
+success = model.export(format="engine", device=0)  # export the model to engine format
+assert success
+```
+
+After executing the above script, you will get an engine named `yolov8s-seg.engine` .
 
 ## Inference with c++
 
@@ -170,9 +201,11 @@ You can infer with c++ in [`csrc/segment/normal`](../csrc/segment/normal) .
 
 ### Build:
 
-Please set you own librarys in [`CMakeLists.txt`](../csrc/segment/normal/CMakeLists.txt) and modify `CLASS_NAMES` and `COLORS` in [`main.cpp`](../csrc/segment/normal/main.cpp).
+Please set you own librarys in [`CMakeLists.txt`](../csrc/segment/normal/CMakeLists.txt) and modify `CLASS_NAMES`
+and `COLORS` in [`main.cpp`](../csrc/segment/normal/main.cpp).
 
-Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk` in [`main.cpp`](../csrc/segment/normal/main.cpp).
+Besides, you can modify the postprocess parameters such as `num_labels` and `score_thres` and `iou_thres` and `topk`
+in [`main.cpp`](../csrc/segment/normal/main.cpp).
 
 ```c++
 int topk = 100;

infer-det-without-torch.py

@@ -38,6 +38,10 @@ def main(args: argparse.Namespace) -> None:
         data = Engine(tensor)
 
         bboxes, scores, labels = det_postprocess(data)
+        if bboxes.size == 0:
+            # if no bounding box
+            print(f'{image}: no object!')
+            continue
         bboxes -= dwdh
         bboxes /= ratio
 

infer-det.py

@@ -37,6 +37,10 @@ def main(args: argparse.Namespace) -> None:
         data = Engine(tensor)
 
         bboxes, scores, labels = det_postprocess(data)
+        if bboxes.numel() == 0:
+            # if no bounding box
+            print(f'{image}: no object!')
+            continue
         bboxes -= dwdh
         bboxes /= ratio
 

infer-pose-without-torch.py

@@ -0,0 +1,116 @@
+import argparse
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from config import COLORS, KPS_COLORS, LIMB_COLORS, SKELETON
+from models.utils import blob, letterbox, path_to_list, pose_postprocess
+
+
+def main(args: argparse.Namespace) -> None:
+    if args.method == 'cudart':
+        from models.cudart_api import TRTEngine
+    elif args.method == 'pycuda':
+        from models.pycuda_api import TRTEngine
+    else:
+        raise NotImplementedError
+
+    Engine = TRTEngine(args.engine)
+    H, W = Engine.inp_info[0].shape[-2:]
+
+    images = path_to_list(args.imgs)
+    save_path = Path(args.out_dir)
+
+    if not args.show and not save_path.exists():
+        save_path.mkdir(parents=True, exist_ok=True)
+
+    for image in images:
+        save_image = save_path / image.name
+        bgr = cv2.imread(str(image))
+        draw = bgr.copy()
+        bgr, ratio, dwdh = letterbox(bgr, (W, H))
+        dw, dh = int(dwdh[0]), int(dwdh[1])
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        tensor = blob(rgb, return_seg=False)
+        dwdh = np.array(dwdh * 2, dtype=np.float32)
+        tensor = np.ascontiguousarray(tensor)
+        # inference
+        data = Engine(tensor)
+
+        bboxes, scores, kpts = pose_postprocess(data, args.conf_thres,
+                                                args.iou_thres)
+        if bboxes.size == 0:
+            # if no bounding box
+            print(f'{image}: no object!')
+            continue
+        bboxes -= dwdh
+        bboxes /= ratio
+
+        for (bbox, score, kpt) in zip(bboxes, scores, kpts):
+            bbox = bbox.round().astype(np.int32).tolist()
+            color = COLORS['person']
+            cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2)
+            cv2.putText(draw,
+                        f'person:{score:.3f}', (bbox[0], bbox[1] - 2),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.75, [225, 255, 255],
+                        thickness=2)
+            for i in range(19):
+                if i < 17:
+                    px, py, ps = kpt[i]
+                    if ps > 0.5:
+                        kcolor = KPS_COLORS[i]
+                        px = round(float(px - dw) / ratio)
+                        py = round(float(py - dh) / ratio)
+                        cv2.circle(draw, (px, py), 5, kcolor, -1)
+                xi, yi = SKELETON[i]
+                pos1_s = kpt[xi - 1][2]
+                pos2_s = kpt[yi - 1][2]
+                if pos1_s > 0.5 and pos2_s > 0.5:
+                    limb_color = LIMB_COLORS[i]
+                    pos1_x = round(float(kpt[xi - 1][0] - dw) / ratio)
+                    pos1_y = round(float(kpt[xi - 1][1] - dh) / ratio)
+
+                    pos2_x = round(float(kpt[yi - 1][0] - dw) / ratio)
+                    pos2_y = round(float(kpt[yi - 1][1] - dh) / ratio)
+
+                    cv2.line(draw, (pos1_x, pos1_y), (pos2_x, pos2_y),
+                             limb_color, 2)
+        if args.show:
+            cv2.imshow('result', draw)
+            cv2.waitKey(0)
+        else:
+            cv2.imwrite(str(save_image), draw)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--engine', type=str, help='Engine file')
+    parser.add_argument('--imgs', type=str, help='Images file')
+    parser.add_argument('--show',
+                        action='store_true',
+                        help='Show the detection results')
+    parser.add_argument('--out-dir',
+                        type=str,
+                        default='./output',
+                        help='Path to output file')
+    parser.add_argument('--conf-thres',
+                        type=float,
+                        default=0.25,
+                        help='Confidence threshold')
+    parser.add_argument('--iou-thres',
+                        type=float,
+                        default=0.65,
+                        help='Confidence threshold')
+    parser.add_argument('--method',
+                        type=str,
+                        default='cudart',
+                        help='CUDART pipeline')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

infer-pose.py

@@ -0,0 +1,112 @@
+from models import TRTModule  # isort:skip
+import argparse
+from pathlib import Path
+
+import cv2
+import torch
+
+from config import COLORS, KPS_COLORS, LIMB_COLORS, SKELETON
+from models.torch_utils import pose_postprocess
+from models.utils import blob, letterbox, path_to_list
+
+
+def main(args: argparse.Namespace) -> None:
+    device = torch.device(args.device)
+    Engine = TRTModule(args.engine, device)
+    H, W = Engine.inp_info[0].shape[-2:]
+
+    images = path_to_list(args.imgs)
+    save_path = Path(args.out_dir)
+
+    if not args.show and not save_path.exists():
+        save_path.mkdir(parents=True, exist_ok=True)
+
+    for image in images:
+        save_image = save_path / image.name
+        bgr = cv2.imread(str(image))
+        draw = bgr.copy()
+        bgr, ratio, dwdh = letterbox(bgr, (W, H))
+        dw, dh = int(dwdh[0]), int(dwdh[1])
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        tensor = blob(rgb, return_seg=False)
+        dwdh = torch.asarray(dwdh * 2, dtype=torch.float32, device=device)
+        tensor = torch.asarray(tensor, device=device)
+        # inference
+        data = Engine(tensor)
+
+        bboxes, scores, kpts = pose_postprocess(data, args.conf_thres,
+                                                args.iou_thres)
+        if bboxes.numel() == 0:
+            # if no bounding box
+            print(f'{image}: no object!')
+            continue
+        bboxes -= dwdh
+        bboxes /= ratio
+
+        for (bbox, score, kpt) in zip(bboxes, scores, kpts):
+            bbox = bbox.round().int().tolist()
+            color = COLORS['person']
+            cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2)
+            cv2.putText(draw,
+                        f'person:{score:.3f}', (bbox[0], bbox[1] - 2),
+                        cv2.FONT_HERSHEY_SIMPLEX,
+                        0.75, [225, 255, 255],
+                        thickness=2)
+            for i in range(19):
+                if i < 17:
+                    px, py, ps = kpt[i]
+                    if ps > 0.5:
+                        kcolor = KPS_COLORS[i]
+                        px = round(float(px - dw) / ratio)
+                        py = round(float(py - dh) / ratio)
+                        cv2.circle(draw, (px, py), 5, kcolor, -1)
+                xi, yi = SKELETON[i]
+                pos1_s = kpt[xi - 1][2]
+                pos2_s = kpt[yi - 1][2]
+                if pos1_s > 0.5 and pos2_s > 0.5:
+                    limb_color = LIMB_COLORS[i]
+                    pos1_x = round(float(kpt[xi - 1][0] - dw) / ratio)
+                    pos1_y = round(float(kpt[xi - 1][1] - dh) / ratio)
+
+                    pos2_x = round(float(kpt[yi - 1][0] - dw) / ratio)
+                    pos2_y = round(float(kpt[yi - 1][1] - dh) / ratio)
+
+                    cv2.line(draw, (pos1_x, pos1_y), (pos2_x, pos2_y),
+                             limb_color, 2)
+        if args.show:
+            cv2.imshow('result', draw)
+            cv2.waitKey(0)
+        else:
+            cv2.imwrite(str(save_image), draw)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--engine', type=str, help='Engine file')
+    parser.add_argument('--imgs', type=str, help='Images file')
+    parser.add_argument('--show',
+                        action='store_true',
+                        help='Show the detection results')
+    parser.add_argument('--out-dir',
+                        type=str,
+                        default='./output',
+                        help='Path to output file')
+    parser.add_argument('--conf-thres',
+                        type=float,
+                        default=0.25,
+                        help='Confidence threshold')
+    parser.add_argument('--iou-thres',
+                        type=float,
+                        default=0.65,
+                        help='Confidence threshold')
+    parser.add_argument('--device',
+                        type=str,
+                        default='cuda:0',
+                        help='TensorRT infer device')
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

infer-seg-without-torch.py

@@ -41,6 +41,10 @@ def main(args: argparse.Namespace) -> None:
         seg_img = seg_img[dh:H - dh, dw:W - dw, [2, 1, 0]]
         bboxes, scores, labels, masks = seg_postprocess(
             data, bgr.shape[:2], args.conf_thres, args.iou_thres)
+        if bboxes.size == 0:
+            # if no bounding box
+            print(f'{image}: no object!')
+            continue
         masks = masks[:, dh:H - dh, dw:W - dw, :]
         mask_colors = MASK_COLORS[labels % len(MASK_COLORS)]
         mask_colors = mask_colors.reshape(-1, 1, 1, 3) * ALPHA

infer-seg.py

@@ -42,10 +42,9 @@ def main(args: argparse.Namespace) -> None:
                                 device=device)
         bboxes, scores, labels, masks = seg_postprocess(
             data, bgr.shape[:2], args.conf_thres, args.iou_thres)
-        if bboxes is None:
+        if bboxes.numel() == 0:
-            # if no bounding box or others save original image
+            # if no bounding box
-            if not args.show:
+            print(f'{image}: no object!')
-                cv2.imwrite(str(save_image), draw)
             continue
         masks = masks[:, dh:H - dh, dw:W - dw, :]
         indices = (labels % len(MASK_COLORS)).long()

models/torch_utils.py

@@ -3,7 +3,7 @@ from typing import List, Tuple, Union
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-from torchvision.ops import batched_nms
+from torchvision.ops import batched_nms, nms
 
 
 def seg_postprocess(
@@ -14,12 +14,13 @@ def seg_postprocess(
         -> Tuple[Tensor, Tensor, Tensor, Tensor]:
     assert len(data) == 2
     h, w = shape[0] // 4, shape[1] // 4  # 4x downsampling
-    outputs, proto = (i[0] for i in data)
+    outputs, proto = data[0][0], data[1][0]
     bboxes, scores, labels, maskconf = outputs.split([4, 1, 1, 32], 1)
     scores, labels = scores.squeeze(), labels.squeeze()
     idx = scores > conf_thres
-    if idx.sum() == 0:  # no bounding boxes or seg were created
+    if not idx.any():  # no bounding boxes or seg were created
-        return None, None, None, None
+        return bboxes.new_zeros((0, 4)), scores.new_zeros(
+            (0, )), labels.new_zeros((0, )), bboxes.new_zeros((0, 0, 0, 0))
     bboxes, scores, labels, maskconf = \
         bboxes[idx], scores[idx], labels[idx], maskconf[idx]
     idx = batched_nms(bboxes, scores, labels, iou_thres)
@@ -35,10 +36,37 @@ def seg_postprocess(
     return bboxes, scores, labels, masks
 
 
+def pose_postprocess(
+        data: Union[Tuple, Tensor],
+        conf_thres: float = 0.25,
+        iou_thres: float = 0.65) \
+        -> Tuple[Tensor, Tensor, Tensor]:
+    if isinstance(data, tuple):
+        assert len(data) == 1
+        data = data[0]
+    outputs = torch.transpose(data[0], 0, 1).contiguous()
+    bboxes, scores, kpts = outputs.split([4, 1, 51], 1)
+    scores, kpts = scores.squeeze(), kpts.squeeze()
+    idx = scores > conf_thres
+    if not idx.any():  # no bounding boxes or seg were created
+        return bboxes.new_zeros((0, 4)), scores.new_zeros(
+            (0, )), bboxes.new_zeros((0, 0, 0))
+    bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx]
+    xycenter, wh = bboxes.chunk(2, -1)
+    bboxes = torch.cat([xycenter - 0.5 * wh, xycenter + 0.5 * wh], -1)
+    idx = nms(bboxes, scores, iou_thres)
+    bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx]
+    return bboxes, scores, kpts.reshape(idx.shape[0], -1, 3)
+
+
 def det_postprocess(data: Tuple[Tensor, Tensor, Tensor, Tensor]):
     assert len(data) == 4
-    num_dets, bboxes, scores, labels = (i[0] for i in data)
+    num_dets, bboxes, scores, labels = data[0][0], data[1][0], data[2][
+        0], data[3][0]
     nums = num_dets.item()
+    if nums == 0:
+        return bboxes.new_zeros((0, 4)), scores.new_zeros(
+            (0, )), labels.new_zeros((0, ))
     bboxes = bboxes[:nums]
     scores = scores[:nums]
     labels = labels[:nums]

models/utils.py

@@ -45,6 +45,7 @@ def letterbox(im: ndarray,
 
 
 def blob(im: ndarray, return_seg: bool = False) -> Union[ndarray, Tuple]:
+    seg = None
     if return_seg:
         seg = im.astype(np.float32) / 255
     im = im.transpose([2, 0, 1])
@@ -88,6 +89,9 @@ def det_postprocess(data: Tuple[ndarray, ndarray, ndarray, ndarray]):
     assert len(data) == 4
     num_dets, bboxes, scores, labels = (i[0] for i in data)
     nums = num_dets.item()
+    if nums == 0:
+        return np.empty((0, 4), dtype=np.float32), np.empty(
+            (0, ), dtype=np.float32), np.empty((0, ), dtype=np.int32)
     bboxes = bboxes[:nums]
     scores = scores[:nums]
     labels = labels[:nums]
@@ -106,6 +110,12 @@ def seg_postprocess(
     bboxes, scores, labels, maskconf = np.split(outputs, [4, 5, 6], 1)
     scores, labels = scores.squeeze(), labels.squeeze()
     idx = scores > conf_thres
+    if not idx.any():  # no bounding boxes or seg were created
+        return np.empty((0, 4), dtype=np.float32), \
+            np.empty((0,), dtype=np.float32), \
+            np.empty((0,), dtype=np.int32), \
+            np.empty((0, 0, 0, 0), dtype=np.int32)
+
     bboxes, scores, labels, maskconf = \
         bboxes[idx], scores[idx], labels[idx], maskconf[idx]
     cvbboxes = np.concatenate([bboxes[:, :2], bboxes[:, 2:] - bboxes[:, :2]],
@@ -128,3 +138,29 @@ def seg_postprocess(
     masks = masks.transpose(2, 0, 1)
     masks = np.ascontiguousarray((masks > 0.5)[..., None], dtype=np.float32)
     return bboxes, scores, labels, masks
+
+
+def pose_postprocess(
+        data: Union[Tuple, ndarray],
+        conf_thres: float = 0.25,
+        iou_thres: float = 0.65) \
+        -> Tuple[ndarray, ndarray, ndarray]:
+    if isinstance(data, tuple):
+        assert len(data) == 1
+        data = data[0]
+    outputs = np.transpose(data[0], (1, 0))
+    bboxes, scores, kpts = np.split(outputs, [4, 5], 1)
+    scores, kpts = scores.squeeze(), kpts.squeeze()
+    idx = scores > conf_thres
+    if not idx.any():  # no bounding boxes or seg were created
+        return np.empty((0, 4), dtype=np.float32), np.empty(
+            (0, ), dtype=np.float32), np.empty((0, 0, 0), dtype=np.float32)
+    bboxes, scores, kpts = bboxes[idx], scores[idx], kpts[idx]
+    xycenter, wh = np.split(bboxes, [
+        2,
+    ], -1)
+    cvbboxes = np.concatenate([xycenter - 0.5 * wh, wh], -1)
+    idx = cv2.dnn.NMSBoxes(cvbboxes, scores, conf_thres, iou_thres)
+    cvbboxes, scores, kpts = cvbboxes[idx], scores[idx], kpts[idx]
+    cvbboxes[:, 2:] += cvbboxes[:, :2]
+    return cvbboxes, scores, kpts.reshape(idx.shape[0], -1, 3)