Tensorrt model implementation

.vscode/settings.json

@@ -5,6 +5,41 @@
         "limits": "cpp",
         "type_traits": "cpp",
         "iosfwd": "cpp",
-        "fstream": "cpp"
+        "fstream": "cpp",
+        "array": "cpp",
+        "atomic": "cpp",
+        "*.tcc": "cpp",
+        "cctype": "cpp",
+        "clocale": "cpp",
+        "cmath": "cpp",
+        "cstdarg": "cpp",
+        "cstddef": "cpp",
+        "cstdint": "cpp",
+        "cstdio": "cpp",
+        "cstdlib": "cpp",
+        "cwchar": "cpp",
+        "cwctype": "cpp",
+        "deque": "cpp",
+        "unordered_map": "cpp",
+        "vector": "cpp",
+        "exception": "cpp",
+        "algorithm": "cpp",
+        "memory": "cpp",
+        "memory_resource": "cpp",
+        "optional": "cpp",
+        "string": "cpp",
+        "string_view": "cpp",
+        "system_error": "cpp",
+        "tuple": "cpp",
+        "utility": "cpp",
+        "initializer_list": "cpp",
+        "iostream": "cpp",
+        "istream": "cpp",
+        "new": "cpp",
+        "ostream": "cpp",
+        "sstream": "cpp",
+        "stdexcept": "cpp",
+        "streambuf": "cpp",
+        "typeinfo": "cpp"
     }
 }

Makefile

@@ -11,6 +11,8 @@ TARGET := $(BUILD_DIR)/pipeline
 
 LIB := tensorflow
 
+LEAK := FALSE
+
 # Conditionally add TFLiteModel.o if target is tensorflow
 ifeq ($(LIB),tensorflow)
   OBJS += $(BUILD_DIR)/TFLiteModel.o
@@ -22,19 +24,29 @@ ifeq ($(LIB),nvinfer)
 endif
 
 # Compiler flags
-CFLAGS := -Wall -Werror -Wpedantic
+CFLAGS := -Wall -Werror -Wpedantic 
+
+LINKERFLAGS := -lstdc++ 
+
+# Conditionally add leak sanitizer
+ifeq ($(LEAK),TRUE)
+  CFLAGS += -fsanitize=leak
+endif
+
 
 # Include paths
 ifeq ($(LIB),nvinfer)
-  CFLAGS += -I/usr/local/cuda/include
+  CFLAGS += -I/usr/local/cuda/include 
+  LINKERFLAGS += -lcudart -lcuda -L/usr/local/cuda/lib64
+
 endif
 
 # Library to link against (default to tensorflow)
 
 
 # Executable
 $(TARGET): $(OBJS)
-	$(CC) $(CFLAGS) $(OBJS) -l$(LIB) -o $(TARGET) -lstdc++
+	$(CC) $(CFLAGS) $(OBJS) -l$(LIB) -o $(TARGET) $(LINKERFLAGS)
 
 # Object file rules
 $(BUILD_DIR)/main.o: $(SRC_DIR)/main.cpp $(SRC_DIR)/filter/IFilter.hpp $(SRC_DIR)/filter/segfilter.hpp

src/filter/segfilter.cpp

@@ -13,7 +13,7 @@ void SegFilter::doDecision() {
 
 SegFilter::SegFilter(const char *modelPath) {
 
-    IModel * model = new TFLiteModel();
+    IModel * model = new TensorRTModel();
     model -> loadModel(modelPath);
     SegFilter::model = model;
     printf("SegFilter initialized!\n");

src/filter/segfilter.hpp

@@ -2,7 +2,7 @@
 #define SEG_FILTER_H
 
 #include "IFilter.hpp"
-#include "../model/TFLiteModel.hpp"
+#include "../model/TensorRTModel.hpp"
 #include <stdio.h>
 
 class SegFilter: public IFilter{

src/main.cpp

@@ -10,4 +10,6 @@ int main() {
     stateManager -> runStateProcess();
     stateManager -> setState(new ProcessingState());
     stateManager -> runStateProcess();
+
+    delete stateManager;
 }

src/model/TensorRTModel.cpp

@@ -37,15 +37,67 @@ std::vector<char> readEngine(const std::string& enginePath)
 
 void TensorRTModel::loadModel(const char *modelPath)
 {
+    std::cout << "TensorRT version: "
+          << NV_TENSORRT_MAJOR << "." 
+          << NV_TENSORRT_MINOR << "." 
+          << NV_TENSORRT_PATCH << "." 
+          << NV_TENSORRT_BUILD << std::endl;
         // Create a TensorRT runtime
+    if (runtime != nullptr) {
+        printf("Runtime already initialized\n");
+        return;
+    }
     runtime = nvinfer1::createInferRuntime(gLogger);
 
-    std::vector<char> engineData = readEngine("your_model_path");
+    std::vector<char> engineData = readEngine(modelPath);
+
+    if (engine != nullptr) {
+        printf("Engine already initialized\n");
+        return;
+    }
 
     engine = runtime->deserializeCudaEngine(engineData.data(), engineData.size());
+
+    engineData.clear();
+
+    if (engine == nullptr) {
+        printf("Error loading engine\n");
+        return;
+    }
 }
 
 void TensorRTModel::predict(unsigned char *image, int height, int width, int channels)
 {
+    if (engine == nullptr) {
+        printf("Engine not initialized\n");
+        return;
+    }
+    printf("Running TensorRT inference on GPU\n");
+        // Allocate GPU memory for the input and output buffers
+    float* gpu_input;
+    float* gpu_output;
+    cudaMalloc((void**)&gpu_input, sizeof(float) * height * width * channels);
+    cudaMalloc((void**)&gpu_output, sizeof(float) * height * width * channels);
+
+        // Create an execution context
+    nvinfer1::IExecutionContext* context = engine->createExecutionContext();
+
+    // Set the input and output buffers for the execution context
+    void* buffers[2] = { gpu_input, gpu_output };
+
+    // Perform inference
+    context->execute(1, buffers);
+
+    // Copy the output data to the CPU memory
+    float* cpu_output = new float[height * width * channels];
+    cudaMemcpy(cpu_output, gpu_output, sizeof(float) * height * width * channels, cudaMemcpyDeviceToHost);
 
+    // Clean up
+    //!TODO: This is possibly the image output from the model so it has to be returned
+    delete[] cpu_output;
+
+    cudaFree(gpu_input);
+    cudaFree(gpu_output);
+    delete context;
+    printf("TensorRT inference done!\n");
 }

src/state/ProcessingState.cpp

@@ -28,15 +28,15 @@ int ProcessingState::runStateProcess(){
 
 
     // data processing
-    SegFilter segfilter("model.pb");
+    SegFilter segfilter("model.engine");
     segfilter.doProcessing(image, width, height, channels);
     segfilter.doDecision();
     printf("processing done!\n");
     // image compression
 
     // IO writing
     printf("freeing image\n");
-
+    stbi_image_free(image);
     printf("image freed\n");
     return 0;
 }