59 lines
1.9 KiB
C++
59 lines
1.9 KiB
C++
#include <iostream>
|
|
#include "structs.h"
|
|
#include <chrono>
|
|
#include "cuda_runtime_api.h"
|
|
#include "CropLP_logging.h"
|
|
#include "CropLP_common.hpp"
|
|
#include <math.h>
|
|
#define DEVICE 0
|
|
#define NET s // s m l x
|
|
#define NETSTRUCT(str) createEngine_##str
|
|
#define CREATENET(net) NETSTRUCT(net)
|
|
#define STR1(x) #x
|
|
#define STR2(x) STR1(x)
|
|
// #define USE_FP16 // comment out this if want to use FP16
|
|
#define CONF_THRESH 0.51
|
|
#define BATCH_SIZE 1
|
|
// stuff we know about the network and the input/output blobs
|
|
|
|
static const int INPUT_H_UNET = 224;
|
|
static const int INPUT_W_UNET = 224;
|
|
static const int OUTPUT_SIZE_UNET = 224*224;
|
|
|
|
using namespace nvinfer1;
|
|
static Logger gLogger;
|
|
|
|
struct DetectionCropLP{
|
|
float mask[INPUT_W_UNET * INPUT_H_UNET * 1];
|
|
};
|
|
|
|
class CropLP{
|
|
private:
|
|
std::string model_path = "/home/thai/py_script/MTA/unet-segment/unet/license_ckpt_crop/22-0.9823-0.0604.trt";
|
|
char *trtModelStream{nullptr};
|
|
size_t size{0};
|
|
IRuntime* runtime{nullptr};
|
|
ICudaEngine* engine{nullptr};
|
|
IExecutionContext* context{nullptr};
|
|
const char* INPUT_BLOB_NAME_UNET = "input_1";
|
|
const char* OUTPUT_BLOB_NAME_UNET = "conv2d_23";
|
|
|
|
|
|
public:
|
|
static CropLP &getInst()
|
|
{
|
|
static CropLP instance;
|
|
return instance;
|
|
};
|
|
CropLP();
|
|
void _init();
|
|
float eculidDistance(cv::Point p1, cv::Point p2);
|
|
void validate_point(cv::Point2f &tl, cv::Point2f &tr, cv::Point2f &br, cv::Point2f &bl, cv::Mat img_ori);
|
|
cv::Mat preprocess_img(cv::Mat& img);
|
|
void doInference(IExecutionContext& context, float* input, float* output, int batchSize);
|
|
void process_cls_result(DetectionCropLP &res, float *output);
|
|
float sigmoid(float x);
|
|
cv::Mat segmentation(cv::Mat img);
|
|
cv::Mat Crop(cv::Mat mask, cv::Mat img_ori, std::string mode);
|
|
std::vector<cv::Mat> predict(cv::Mat inputImage, std::string mode);
|
|
}; |