diff --git a/csrnet/CMakeLists.txt b/csrnet/CMakeLists.txt
new file mode 100644
index 0000000..55bb75b
--- /dev/null
+++ b/csrnet/CMakeLists.txt
@@ -0,0 +1,26 @@
+cmake_minimum_required(VERSION 3.10)
+
+project(csrnet)
+
+add_definitions(-std=c++11)
+add_definitions(-DAPI_EXPORTS)
+option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
+set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_BUILD_TYPE Debug)
+
+# cuda
+include_directories(/usr/local/cuda/targets/x86_64-linux/include )
+link_directories(/usr/local/cuda/targets/x86_64-linux/lib)
+
+# tensorrt
+include_directories(/usr/include/x86_64-linux-gnu/)
+link_directories(/usr/lib/x86_64-linux-gnu/)
+
+# opencv
+find_package(OpenCV)
+include_directories(${OpenCV_INCLUDE_DIRS})
+
+include_directories(${PROJECT_SOURCE_DIR}/)
+
+add_executable(csrnet csrnet.cpp)
+target_link_libraries(csrnet nvinfer cudart ${OpenCV_LIBS})
\ No newline at end of file
diff --git a/csrnet/README.md b/csrnet/README.md
new file mode 100644
index 0000000..54cfc4a
--- /dev/null
+++ b/csrnet/README.md
@@ -0,0 +1,58 @@
+# csrnet
+
+The Pytorch implementation is [leeyeehoo/CSRNet-pytorch](https://github.com/leeyeehoo/CSRNet-pytorch).
+
+This repo is a TensorRT implementation of CSRNet.
+
+paper : [CSRNet: Dilated Convolutional Neural Networks for Understanding the Highly Congested Scenes](https://arxiv.org/abs/1802.10062)
+
+Dev environment:
+- Ubuntu 22.04
+- TensorRT 8.6
+- OpenCV 4.5.4
+- CMake 3.24
+- GPU Driver 535.113.01
+- CUDA 12.2
+- RTX3080
+
+
+# how to run
+
+```bash
+1. generate csrnet engine
+git clone https://github.com/leeyeehoo/CSRNet-pytorch.git
+git clone https://github.com/wang-xinyu/tensorrtx.git
+// copy gen_wts.py to CSRNet-pytorch
+// generate wts file
+python gen_wts.py
+// csrnet wts will be generated in CSRNet-pytorch
+
+2. build csrnet.engine
+// mv CSRNet-pytorch/csrnet.engine to tensorrtx/csrnet
+mv CSRNet-pytorch/csrnet.wts tensorrtx/csrnet
+// build
+mkdir build
+cmake ..
+make
+sudo ./csrnet -s  ./csrnet.wts
+
+Loading weights: ./csrnet.wts
+build engine successfully : ./csrnet.engine
+
+// download images https://github.com/wang-xinyu/tensorrtx/assets/46584679/46bc4def-e573-44ae-996d-5d68927c78ff and copy to images
+sudo ./csrnet -d  ./images
+
+// output e.g
+// enqueueV2 time: 0.0323869s
+// detect time:44ms
+// people num :22.9101 write_path: ../images/data.jpg
+```
+
+
+# result 
+
+inference people num: 22.9101
+
+<p align="center">
+<img src= https://raw.githubusercontent.com/wang-xinyu/tensorrtx/dbf857d25f77bf64113fc99a745ccf4973bdd44e/Density_Plot.jpg>
+</p>
diff --git a/csrnet/config.h b/csrnet/config.h
new file mode 100644
index 0000000..1bf0a17
--- /dev/null
+++ b/csrnet/config.h
@@ -0,0 +1,16 @@
+#pragma once
+
+const static char *kInputTensorName = "data";
+const static char *kOutputTensorName = "prob";
+const static char *kEngineFile = "./csrnet.engine";
+
+const static int kBatchSize = 1;
+
+const static int MAX_INPUT_SIZE = 1440; // 32x
+const static int MIN_INPUT_SIZE = 608;
+const static int OPT_INPUT_W = 1152;
+const static int OPT_INPUT_H = 640;
+
+constexpr static int kMaxInputImageSize = MAX_INPUT_SIZE * MAX_INPUT_SIZE * 3;
+constexpr static int kMaxOutputProbSize =
+    (MAX_INPUT_SIZE * MAX_INPUT_SIZE) >> 6;
\ No newline at end of file
diff --git a/csrnet/csrnet.cpp b/csrnet/csrnet.cpp
new file mode 100644
index 0000000..7ef08bd
--- /dev/null
+++ b/csrnet/csrnet.cpp
@@ -0,0 +1,536 @@
+#include "NvInfer.h"
+#include "cuda_runtime_api.h"
+#include <chrono>
+#include <config.h>
+#include <cstring>
+#include <dirent.h>
+#include <fstream>
+#include <iostream>
+#include <logging.h>
+#include <map>
+#include <numeric>
+#include <opencv2/opencv.hpp>
+#include <vector>
+using namespace nvinfer1;
+
+#define CHECK(status)                                                          \
+  do {                                                                         \
+    auto ret = (status);                                                       \
+    if (ret != 0) {                                                            \
+      std::cerr << "Cuda failure: " << ret << std::endl;                       \
+      abort();                                                                 \
+    }                                                                          \
+  } while (0)
+
+static Logger gLogger;
+static char *kWTSFile = "";
+std::map<std::string, Weights> loadWeights(const std::string file) {
+  std::cout << "Loading weights: " << file << std::endl;
+  std::map<std::string, Weights> weightMap;
+
+  // Open weights file
+  std::ifstream input(file);
+  assert(input.is_open() && "Unable to load weight file.");
+
+  // Read number of weight blobs
+  int32_t count;
+  input >> count;
+  assert(count > 0 && "Invalid weight map file.");
+
+  while (count--) {
+    Weights wt{DataType::kFLOAT, nullptr, 0};
+    uint32_t size;
+
+    // Read name and type of blob
+    std::string name;
+    input >> name >> std::dec >> size;
+    wt.type = DataType::kFLOAT;
+
+    // Load blob
+    uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
+    for (uint32_t x = 0, y = size; x < y; ++x) {
+      input >> std::hex >> val[x];
+    }
+    wt.values = val;
+
+    wt.count = size;
+    weightMap[name] = wt;
+  }
+
+  return weightMap;
+}
+// clang-format off
+/*
+CSRNet(
+ (frontend): Sequential(
+    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (1): ReLU(inplace=True)
+    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (3): ReLU(inplace=True)
+    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (6): ReLU(inplace=True)
+    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (8): ReLU(inplace=True)
+    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (11): ReLU(inplace=True)
+    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (13): ReLU(inplace=True)
+    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (15): ReLU(inplace=True)
+    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (18): ReLU(inplace=True)
+    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (20): ReLU(inplace=True)
+    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (22): ReLU(inplace=True)
+  )
+  (backend): Sequential(
+    (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2),
+    dilation=(2, 2)) (1): ReLU(inplace=True) (2): Conv2d(512, 512,
+    kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2)) (3):
+    ReLU(inplace=True) (4): Conv2d(512, 512, kernel_size=(3, 3), stride=(1,
+    1), padding=(2, 2), dilation=(2, 2)) (5): ReLU(inplace=True) (6):
+    Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(2, 2),
+    dilation=(2, 2)) (7): ReLU(inplace=True) (8): Conv2d(256, 128,
+    kernel_size=(3, 3), stride=(1, 1), padding=(2, 2), dilation=(2, 2)) (9):
+    ReLU(inplace=True) (10): Conv2d(128, 64, kernel_size=(3, 3), stride=(1,
+    1), padding=(2, 2), dilation=(2, 2)) (11): ReLU(inplace=True)
+  )
+  (output_layer): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))
+)
+*/
+// clang-format on
+void doInference(IExecutionContext &context, float *input, float *output,
+                 int input_h, int input_w) {
+  const ICudaEngine &engine = context.getEngine();
+
+  uint64_t input_size = 3 * input_h * input_w * sizeof(float);
+  uint64_t output_size = ((input_h * input_w) >> 6) * sizeof(float);
+
+  // Pointers to input and output device buffers to pass to engine.
+  // Engine requires exactly IEngine::getNbBindings() number of buffers.
+  assert(engine.getNbBindings() == 2);
+  void *buffers[2];
+
+  // In order to bind the buffers, we need to know the names of the input and
+  // output tensors. Note that indices are guaranteed to be less than
+  // IEngine::getNbBindings()
+  const int inputIndex = engine.getBindingIndex(kInputTensorName);
+  const int outputIndex = engine.getBindingIndex(kOutputTensorName);
+  context.setBindingDimensions(inputIndex, Dims4(1, 3, input_h, input_w));
+
+  // Create GPU buffers on device
+  CHECK(cudaMalloc(&buffers[inputIndex], input_size));
+  CHECK(cudaMalloc(&buffers[outputIndex], output_size));
+
+  // Create stream
+  cudaStream_t stream;
+  CHECK(cudaStreamCreate(&stream));
+
+  // DMA input batch data to device, infer on the batch asynchronously, and DMA
+  // output back to host
+  CHECK(cudaMemcpyAsync(buffers[inputIndex], input, input_size,
+                        cudaMemcpyHostToDevice, stream));
+  auto t1 = std::chrono::high_resolution_clock::now();
+  context.enqueueV2(buffers, stream, nullptr);
+  std::cout << "enqueueV2 time: "
+            << std::chrono::duration<float>(
+                   std::chrono::high_resolution_clock::now() - t1)
+                   .count()
+            << "s" << std::endl;
+  CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size,
+                        cudaMemcpyDeviceToHost, stream));
+  cudaStreamSynchronize(stream);
+
+  // Release stream and buffers
+  cudaStreamDestroy(stream);
+  CHECK(cudaFree(buffers[inputIndex]));
+  CHECK(cudaFree(buffers[outputIndex]));
+}
+ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder,
+                          IBuilderConfig *config, DataType dt) {
+
+  //   INetworkDefinition *network = builder->createNetworkV2(0U);
+  const auto explicitBatch =
+      1U << static_cast<uint32_t>(
+          NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
+  INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
+  ITensor *data = network->addInput(kInputTensorName, dt, Dims4{1, 3, -1, -1});
+  assert(data);
+  std::map<std::string, Weights> weightMap = loadWeights(kWTSFile);
+
+  IConvolutionLayer *conv1 = network->addConvolutionNd(
+      *data, 64, DimsHW{3, 3}, weightMap["frontend.0.weight"],
+      weightMap["frontend.0.bias"]);
+  assert(conv1);
+  conv1->setStrideNd(DimsHW{1, 1});
+  conv1->setPaddingNd(DimsHW{1, 1});
+
+  IActivationLayer *relu1 =
+      network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
+
+  assert(relu1);
+
+  auto conv2 = network->addConvolutionNd(*relu1->getOutput(0), 64, DimsHW{3, 3},
+                                         weightMap["frontend.2.weight"],
+                                         weightMap["frontend.2.bias"]);
+  assert(conv2);
+  conv2->setStrideNd(DimsHW{1, 1});
+  conv2->setPaddingNd(DimsHW{1, 1});
+  auto relu2 =
+      network->addActivation(*conv2->getOutput(0), ActivationType::kRELU);
+  assert(relu2);
+  auto pool1 = network->addPoolingNd(*relu2->getOutput(0), PoolingType::kMAX,
+                                     DimsHW{2, 2});
+  assert(pool1);
+  pool1->setStrideNd(DimsHW{2, 2});
+  auto conv3 = network->addConvolutionNd(
+      *pool1->getOutput(0), 128, DimsHW{3, 3}, weightMap["frontend.5.weight"],
+      weightMap["frontend.5.bias"]);
+  assert(conv3);
+  conv3->setStrideNd(DimsHW{1, 1});
+
+  conv3->setPaddingNd(DimsHW{1, 1});
+  auto relu3 =
+      network->addActivation(*conv3->getOutput(0), ActivationType::kRELU);
+  assert(relu3);
+
+  auto conv4 = network->addConvolutionNd(
+      *relu3->getOutput(0), 128, DimsHW{3, 3}, weightMap["frontend.7.weight"],
+      weightMap["frontend.7.bias"]);
+  assert(conv4);
+  conv4->setStrideNd(DimsHW{1, 1});
+  conv4->setPaddingNd(DimsHW{1, 1});
+  auto relu4 =
+      network->addActivation(*conv4->getOutput(0), ActivationType::kRELU);
+  assert(relu4);
+
+  auto pool2 = network->addPoolingNd(*relu4->getOutput(0), PoolingType::kMAX,
+                                     DimsHW{2, 2});
+  assert(pool2);
+  pool2->setStrideNd(DimsHW{2, 2});
+
+  auto conv5 = network->addConvolutionNd(
+      *pool2->getOutput(0), 256, DimsHW{3, 3}, weightMap["frontend.10.weight"],
+      weightMap["frontend.10.bias"]);
+  assert(conv5);
+  conv5->setStrideNd(DimsHW{1, 1});
+  conv5->setPaddingNd(DimsHW{1, 1});
+  auto relu5 =
+      network->addActivation(*conv5->getOutput(0), ActivationType::kRELU);
+  assert(relu5);
+
+  auto conv6 = network->addConvolutionNd(
+      *relu5->getOutput(0), 256, DimsHW{3, 3}, weightMap["frontend.12.weight"],
+      weightMap["frontend.12.bias"]);
+  assert(conv6);
+  conv6->setStrideNd(DimsHW{1, 1});
+  conv6->setPaddingNd(DimsHW{1, 1});
+  auto relu6 =
+      network->addActivation(*conv6->getOutput(0), ActivationType::kRELU);
+  assert(relu6);
+  auto conv7 = network->addConvolutionNd(
+      *relu6->getOutput(0), 256, DimsHW{3, 3}, weightMap["frontend.14.weight"],
+      weightMap["frontend.14.bias"]);
+  assert(conv7);
+  conv7->setStrideNd(DimsHW{1, 1});
+  conv7->setPaddingNd(DimsHW{1, 1});
+  auto relu7 =
+      network->addActivation(*conv7->getOutput(0), ActivationType::kRELU);
+  assert(relu7);
+  auto pool3 = network->addPoolingNd(*relu7->getOutput(0), PoolingType::kMAX,
+                                     DimsHW{2, 2});
+  assert(pool3);
+  pool3->setStrideNd(DimsHW{2, 2});
+  auto conv8 = network->addConvolutionNd(
+      *pool3->getOutput(0), 512, DimsHW{3, 3}, weightMap["frontend.17.weight"],
+      weightMap["frontend.17.bias"]);
+  assert(conv8);
+  conv8->setStrideNd(DimsHW{1, 1});
+  conv8->setPaddingNd(DimsHW{1, 1});
+  auto relu8 =
+      network->addActivation(*conv8->getOutput(0), ActivationType::kRELU);
+  assert(relu8);
+  auto conv9 = network->addConvolutionNd(
+      *relu8->getOutput(0), 512, DimsHW{3, 3}, weightMap["frontend.19.weight"],
+      weightMap["frontend.19.bias"]);
+  assert(conv9);
+  conv9->setStrideNd(DimsHW{1, 1});
+  conv9->setPaddingNd(DimsHW{1, 1});
+  auto relu9 =
+      network->addActivation(*conv9->getOutput(0), ActivationType::kRELU);
+  assert(relu9);
+  auto conv10 = network->addConvolutionNd(
+      *relu9->getOutput(0), 512, DimsHW{3, 3}, weightMap["frontend.21.weight"],
+      weightMap["frontend.21.bias"]);
+  assert(conv10);
+  conv10->setStrideNd(DimsHW{1, 1});
+  conv10->setPaddingNd(DimsHW{1, 1});
+  auto relu10 =
+      network->addActivation(*conv10->getOutput(0), ActivationType::kRELU);
+  assert(relu10);
+  // backend
+  auto conv11 = network->addConvolutionNd(
+      *relu10->getOutput(0), 512, DimsHW{3, 3}, weightMap["backend.0.weight"],
+      weightMap["backend.0.bias"]);
+  assert(conv11);
+  conv11->setPaddingNd(DimsHW{2, 2});
+  conv11->setStrideNd(DimsHW{1, 1});
+  conv11->setDilationNd(DimsHW{2, 2});
+  auto relu11 =
+      network->addActivation(*conv11->getOutput(0), ActivationType::kRELU);
+
+  assert(relu11);
+  auto conv12 = network->addConvolutionNd(
+      *relu11->getOutput(0), 512, DimsHW{3, 3}, weightMap["backend.2.weight"],
+      weightMap["backend.2.bias"]);
+  assert(conv12);
+  conv12->setPaddingNd(DimsHW{2, 2});
+  conv12->setStrideNd(DimsHW{1, 1});
+  conv12->setDilationNd(DimsHW{2, 2});
+  auto relu12 =
+      network->addActivation(*conv12->getOutput(0), ActivationType::kRELU);
+  assert(relu12);
+
+  auto conv13 = network->addConvolutionNd(
+      *relu12->getOutput(0), 512, DimsHW{3, 3}, weightMap["backend.4.weight"],
+      weightMap["backend.4.bias"]);
+  assert(conv13);
+  conv13->setPaddingNd(DimsHW{2, 2});
+  conv13->setStrideNd(DimsHW{1, 1});
+  conv13->setDilationNd(DimsHW{2, 2});
+  auto relu13 =
+      network->addActivation(*conv13->getOutput(0), ActivationType::kRELU);
+  assert(relu13);
+
+  auto conv14 = network->addConvolutionNd(
+      *relu13->getOutput(0), 256, DimsHW{3, 3}, weightMap["backend.6.weight"],
+      weightMap["backend.6.bias"]);
+  assert(conv14);
+  conv14->setPaddingNd(DimsHW{2, 2});
+  conv14->setStrideNd(DimsHW{1, 1});
+  conv14->setDilationNd(DimsHW{2, 2});
+  auto relu14 =
+      network->addActivation(*conv14->getOutput(0), ActivationType::kRELU);
+  assert(relu14);
+  auto conv15 = network->addConvolutionNd(
+      *relu14->getOutput(0), 128, DimsHW{3, 3}, weightMap["backend.8.weight"],
+      weightMap["backend.8.bias"]);
+  assert(conv15);
+  conv15->setPaddingNd(DimsHW{2, 2});
+  conv15->setStrideNd(DimsHW{1, 1});
+  conv15->setDilationNd(DimsHW{2, 2});
+  auto relu15 =
+      network->addActivation(*conv15->getOutput(0), ActivationType::kRELU);
+  assert(relu15);
+  auto conv16 = network->addConvolutionNd(
+      *relu15->getOutput(0), 64, DimsHW{3, 3}, weightMap["backend.10.weight"],
+      weightMap["backend.10.bias"]);
+  assert(conv16);
+  conv16->setPaddingNd(DimsHW{2, 2});
+  conv16->setStrideNd(DimsHW{1, 1});
+  conv16->setDilationNd(DimsHW{2, 2});
+  auto relu16 =
+      network->addActivation(*conv16->getOutput(0), ActivationType::kRELU);
+
+  assert(relu16);
+
+  auto conv17 = network->addConvolutionNd(
+      *relu16->getOutput(0), 1, DimsHW{1, 1}, weightMap["output_layer.weight"],
+      weightMap["output_layer.bias"]);
+  assert(conv17);
+
+  conv17->setStrideNd(DimsHW{1, 1});
+  conv17->getOutput(0)->setName(kOutputTensorName);
+  network->markOutput(*conv17->getOutput(0));
+
+  IOptimizationProfile *profile = builder->createOptimizationProfile();
+  profile->setDimensions(kInputTensorName, OptProfileSelector::kMIN,
+                         Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE));
+  profile->setDimensions(kInputTensorName, OptProfileSelector::kOPT,
+                         Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W));
+  profile->setDimensions(kInputTensorName, OptProfileSelector::kMAX,
+                         Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE));
+  config->addOptimizationProfile(profile);
+
+  builder->setMaxBatchSize(kBatchSize);
+  config->setMaxWorkspaceSize(16 << 20);
+#ifdef USE_FP16
+  config->setFlag(BuilderFlag::kFP16);
+#endif
+  ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
+
+  printf("build engine successfully : %s\n", kEngineFile);
+  // Don't need the network any more
+  network->destroy();
+
+  // Release host memory
+  for (auto &mem : weightMap) {
+    free((void *)(mem.second.values));
+  }
+
+  return engine;
+}
+void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream) {
+  // Create builder
+  IBuilder *builder = createInferBuilder(gLogger);
+  IBuilderConfig *config = builder->createBuilderConfig();
+
+  // Create model to populate the network, then set the outputs and create an
+  // engine
+  ICudaEngine *engine =
+      createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
+  assert(engine != nullptr);
+
+  // Serialize the engine
+  (*modelStream) = engine->serialize();
+
+  // Close everything down
+  engine->destroy();
+  config->destroy();
+  builder->destroy();
+}
+
+int read_files_in_dir(const char *p_dir_name,
+                      std::vector<std::string> &file_names) {
+  DIR *p_dir = opendir(p_dir_name);
+  if (p_dir == nullptr) {
+    return -1;
+  }
+
+  struct dirent *p_file = nullptr;
+  while ((p_file = readdir(p_dir)) != nullptr) {
+    if (strcmp(p_file->d_name, ".") != 0 && strcmp(p_file->d_name, "..") != 0) {
+      std::string cur_file_name(p_file->d_name);
+      file_names.push_back(cur_file_name);
+    }
+  }
+  closedir(p_dir);
+  return 0;
+}
+
+int main(int argc, char **argv) {
+
+  if (argc != 3) {
+    std::cerr << "arguments not right!" << std::endl;
+    std::cerr << "./csrnet -s  ./csrnet.wts // serialize model to plan file"
+              << std::endl;
+    std::cerr
+        << "./csrnet -d  ../images  // deserialize plan file and run inference"
+        << std::endl;
+    return -1;
+  }
+  char *trtModelStream{nullptr};
+  size_t size{0};
+
+  if (std::string(argv[1]) == "-s") {
+    IHostMemory *modelStream{nullptr};
+    kWTSFile = argv[2];
+    APIToModel(kBatchSize, &modelStream);
+    assert(modelStream != nullptr);
+
+    std::ofstream p(kEngineFile, std::ios::binary);
+    if (!p) {
+      std::cerr << "could not open plan output file" << std::endl;
+      return -1;
+    }
+    p.write(reinterpret_cast<const char *>(modelStream->data()),
+            modelStream->size());
+    modelStream->destroy();
+    return 1;
+  } else if (std::string(argv[1]) == "-d") {
+    std::ifstream file(kEngineFile, std::ios::binary);
+    if (file.good()) {
+      file.seekg(0, file.end);
+      size = file.tellg();
+      file.seekg(0, file.beg);
+      trtModelStream = new char[size];
+      assert(trtModelStream);
+      file.read(trtModelStream, size);
+      file.close();
+    }
+  } else {
+    return -1;
+  }
+  IRuntime *runtime = createInferRuntime(gLogger);
+  assert(runtime != nullptr);
+  ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size);
+  assert(engine != nullptr);
+  IExecutionContext *context = engine->createExecutionContext();
+  assert(context != nullptr);
+  delete[] trtModelStream;
+
+  std::vector<std::string> file_names;
+  if (read_files_in_dir(argv[2], file_names) < 0) {
+    std::cout << "read_files_in_dir failed." << std::endl;
+    return -1;
+  }
+
+  std::vector<float> mean_value{0.406, 0.456, 0.485}; // BGR
+  std::vector<float> std_value{0.225, 0.224, 0.229};
+
+  int fcount = 0;
+
+  float *data = new float[kMaxInputImageSize];
+  float *prob = new float[kMaxOutputProbSize];
+
+  for (auto f : file_names) {
+    fcount++;
+    cv::Mat src_img = cv::imread(std::string(argv[2]) + "/" + f);
+    if (src_img.empty())
+      continue;
+
+    int i = 0;
+    for (int row = 0; row < src_img.rows; ++row) {
+      uchar *uc_pixel = src_img.data + row * src_img.step;
+      for (int col = 0; col < src_img.cols; ++col) {
+        data[i] = (uc_pixel[2] / 255.0 - mean_value[2]) / std_value[2];
+        data[i + src_img.rows * src_img.cols] =
+            (uc_pixel[1] / 255.0 - mean_value[1]) / std_value[1];
+        data[i + 2 * src_img.rows * src_img.cols] =
+            (uc_pixel[0] / 255.0 - mean_value[0]) / std_value[0];
+        uc_pixel += 3;
+        ++i;
+      }
+    }
+    // Run inference
+    auto start = std::chrono::system_clock::now();
+    doInference(*context, data, prob, src_img.rows, src_img.cols);
+    auto end = std::chrono::system_clock::now();
+    std::cout << "detect time:"
+              << std::chrono::duration_cast<std::chrono::milliseconds>(end -
+                                                                       start)
+                     .count()
+              << "ms" << std::endl;
+    float num = std::accumulate(
+        prob, prob + ((src_img.rows * src_img.cols) >> 6), 0.0f);
+
+    cv::Mat densityMap(src_img.rows >> 3, src_img.cols >> 3, CV_32FC1,
+                       (void *)prob);
+
+    cv::Mat densityMapScaled;
+    cv::normalize(densityMap, densityMapScaled, 0, 255, cv::NORM_MINMAX,
+                  CV_8UC1);
+    cv::Mat densityColorMap;
+    cv::applyColorMap(densityMapScaled, densityColorMap, cv::COLORMAP_VIRIDIS);
+
+    cv::resize(densityColorMap, densityColorMap, src_img.size());
+    cv::addWeighted(densityColorMap, 0.5, src_img, 0.5, 0, src_img);
+
+    // write to jpg
+    cv::putText(src_img, std::string("people num: ") + std::to_string(num),
+                cv::Point(10, 50), cv::FONT_HERSHEY_SIMPLEX, 0.5,
+                cv::Scalar(255, 255, 255), 1);
+    std::string write_path = std::string(argv[2]) + "result_" + f;
+    std::cout << "people num :" << num << " write_path: " << write_path
+              << std::endl;
+    cv::imwrite(write_path, src_img);
+  }
+  delete[] data;
+  delete[] prob;
+
+  return 0;
+}
\ No newline at end of file
diff --git a/csrnet/gen_wts.py b/csrnet/gen_wts.py
new file mode 100644
index 0000000..68ba51f
--- /dev/null
+++ b/csrnet/gen_wts.py
@@ -0,0 +1,31 @@
+from torch.nn.modules import module
+from model import CSRNet
+import torch
+import os
+import struct
+
+
+save_path = os.path.join(os.path.dirname(
+    __file__), "output", os.path.basename(__file__).split('.')[0])
+os.makedirs(save_path, exist_ok=True)
+wts_file = os.path.join(save_path, "csrnet.wts")
+
+
+# load model
+model_path = "partBmodel_best.pth.tar"
+model = CSRNet()
+checkpoint = torch.load(model_path)
+model.load_state_dict(checkpoint['state_dict'])
+
+
+# save to wts
+print(f'Writing into {wts_file}')
+with open(wts_file, 'w') as f:
+    f.write('{}\n'.format(len(model.state_dict().keys())))
+    for k, v in model.state_dict().items():
+        vr = v.reshape(-1).cpu().numpy()
+        f.write('{} {} '.format(k, len(vr)))
+        for vv in vr:
+            f.write(' ')
+            f.write(struct.pack('>f', float(vv)).hex())
+        f.write('\n')
\ No newline at end of file
diff --git a/csrnet/logging.h b/csrnet/logging.h
new file mode 100644
index 0000000..fb09a20
--- /dev/null
+++ b/csrnet/logging.h
@@ -0,0 +1,502 @@
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TENSORRT_LOGGING_H
+#define TENSORRT_LOGGING_H
+
+#include "NvInferRuntimeCommon.h"
+#include "macros.h"
+#include <cassert>
+#include <ctime>
+#include <iomanip>
+#include <iostream>
+#include <ostream>
+#include <sstream>
+#include <string>
+
+using Severity = nvinfer1::ILogger::Severity;
+
+class LogStreamConsumerBuffer : public std::stringbuf {
+public:
+  LogStreamConsumerBuffer(std::ostream &stream, const std::string &prefix,
+                          bool shouldLog)
+      : mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {}
+
+  LogStreamConsumerBuffer(LogStreamConsumerBuffer &&other)
+      : mOutput(other.mOutput) {}
+
+  ~LogStreamConsumerBuffer() {
+    // std::streambuf::pbase() gives a pointer to the beginning of the buffered
+    // part of the output sequence std::streambuf::pptr() gives a pointer to the
+    // current position of the output sequence if the pointer to the beginning
+    // is not equal to the pointer to the current position, call putOutput() to
+    // log the output to the stream
+    if (pbase() != pptr()) {
+      putOutput();
+    }
+  }
+
+  // synchronizes the stream buffer and returns 0 on success
+  // synchronizing the stream buffer consists of inserting the buffer contents
+  // into the stream, resetting the buffer and flushing the stream
+  virtual int sync() {
+    putOutput();
+    return 0;
+  }
+
+  void putOutput() {
+    if (mShouldLog) {
+      // prepend timestamp
+      std::time_t timestamp = std::time(nullptr);
+      tm *tm_local = std::localtime(&timestamp);
+      std::cout << "[";
+      std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon
+                << "/";
+      std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday
+                << "/";
+      std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year
+                << "-";
+      std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour
+                << ":";
+      std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
+      std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec
+                << "] ";
+      // std::stringbuf::str() gets the string contents of the buffer
+      // insert the buffer contents pre-appended by the appropriate prefix into
+      // the stream
+      mOutput << mPrefix << str();
+      // set the buffer to empty
+      str("");
+      // flush the stream
+      mOutput.flush();
+    }
+  }
+
+  void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; }
+
+private:
+  std::ostream &mOutput;
+  std::string mPrefix;
+  bool mShouldLog;
+};
+
+//!
+//! \class LogStreamConsumerBase
+//! \brief Convenience object used to initialize LogStreamConsumerBuffer before
+//! std::ostream in LogStreamConsumer
+//!
+class LogStreamConsumerBase {
+public:
+  LogStreamConsumerBase(std::ostream &stream, const std::string &prefix,
+                        bool shouldLog)
+      : mBuffer(stream, prefix, shouldLog) {}
+
+protected:
+  LogStreamConsumerBuffer mBuffer;
+};
+
+//!
+//! \class LogStreamConsumer
+//! \brief Convenience object used to facilitate use of C++ stream syntax when
+//! logging messages.
+//!  Order of base classes is LogStreamConsumerBase and then std::ostream.
+//!  This is because the LogStreamConsumerBase class is used to initialize the
+//!  LogStreamConsumerBuffer member field in LogStreamConsumer and then the
+//!  address of the buffer is passed to std::ostream. This is necessary to
+//!  prevent the address of an uninitialized buffer from being passed to
+//!  std::ostream. Please do not change the order of the parent classes.
+//!
+class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
+public:
+  //! \brief Creates a LogStreamConsumer which logs messages with level
+  //! severity.
+  //!  Reportable severity determines if the messages are severe enough to be
+  //!  logged.
+  LogStreamConsumer(Severity reportableSeverity, Severity severity)
+      : LogStreamConsumerBase(severityOstream(severity),
+                              severityPrefix(severity),
+                              severity <= reportableSeverity),
+        std::ostream(&mBuffer) // links the stream buffer with the stream
+        ,
+        mShouldLog(severity <= reportableSeverity), mSeverity(severity) {}
+
+  LogStreamConsumer(LogStreamConsumer &&other)
+      : LogStreamConsumerBase(severityOstream(other.mSeverity),
+                              severityPrefix(other.mSeverity),
+                              other.mShouldLog),
+        std::ostream(&mBuffer) // links the stream buffer with the stream
+        ,
+        mShouldLog(other.mShouldLog), mSeverity(other.mSeverity) {}
+
+  void setReportableSeverity(Severity reportableSeverity) {
+    mShouldLog = mSeverity <= reportableSeverity;
+    mBuffer.setShouldLog(mShouldLog);
+  }
+
+private:
+  static std::ostream &severityOstream(Severity severity) {
+    return severity >= Severity::kINFO ? std::cout : std::cerr;
+  }
+
+  static std::string severityPrefix(Severity severity) {
+    switch (severity) {
+    case Severity::kINTERNAL_ERROR:
+      return "[F] ";
+    case Severity::kERROR:
+      return "[E] ";
+    case Severity::kWARNING:
+      return "[W] ";
+    case Severity::kINFO:
+      return "[I] ";
+    case Severity::kVERBOSE:
+      return "[V] ";
+    default:
+      assert(0);
+      return "";
+    }
+  }
+
+  bool mShouldLog;
+  Severity mSeverity;
+};
+
+//! \class Logger
+//!
+//! \brief Class which manages logging of TensorRT tools and samples
+//!
+//! \details This class provides a common interface for TensorRT tools and
+//! samples to log information to the console, and supports logging two types of
+//! messages:
+//!
+//! - Debugging messages with an associated severity (info, warning, error, or
+//! internal error/fatal)
+//! - Test pass/fail messages
+//!
+//! The advantage of having all samples use this class for logging as opposed to
+//! emitting directly to stdout/stderr is that the logic for controlling the
+//! verbosity and formatting of sample output is centralized in one location.
+//!
+//! In the future, this class could be extended to support dumping test results
+//! to a file in some standard format (for example, JUnit XML), and providing
+//! additional metadata (e.g. timing the duration of a test run).
+//!
+//! TODO: For backwards compatibility with existing samples, this class inherits
+//! directly from the nvinfer1::ILogger interface, which is problematic since
+//! there isn't a clean separation between messages coming from the TensorRT
+//! library and messages coming from the sample.
+//!
+//! In the future (once all samples are updated to use Logger::getTRTLogger() to
+//! access the ILogger) we can refactor the class to eliminate the inheritance
+//! and instead make the nvinfer1::ILogger implementation a member of the Logger
+//! object.
+
+class Logger : public nvinfer1::ILogger {
+public:
+  Logger(Severity severity = Severity::kWARNING)
+      : mReportableSeverity(severity) {}
+
+  //!
+  //! \enum TestResult
+  //! \brief Represents the state of a given test
+  //!
+  enum class TestResult {
+    kRUNNING, //!< The test is running
+    kPASSED,  //!< The test passed
+    kFAILED,  //!< The test failed
+    kWAIVED   //!< The test was waived
+  };
+
+  //!
+  //! \brief Forward-compatible method for retrieving the nvinfer::ILogger
+  //! associated with this Logger \return The nvinfer1::ILogger associated with
+  //! this Logger
+  //!
+  //! TODO Once all samples are updated to use this method to register the
+  //! logger with TensorRT, we can eliminate the inheritance of Logger from
+  //! ILogger
+  //!
+  nvinfer1::ILogger &getTRTLogger() { return *this; }
+
+  //!
+  //! \brief Implementation of the nvinfer1::ILogger::log() virtual method
+  //!
+  //! Note samples should not be calling this function directly; it will
+  //! eventually go away once we eliminate the inheritance from
+  //! nvinfer1::ILogger
+  //!
+  void log(Severity severity, const char *msg) TRT_NOEXCEPT override {
+    LogStreamConsumer(mReportableSeverity, severity)
+        << "[TRT] " << std::string(msg) << std::endl;
+  }
+
+  //!
+  //! \brief Method for controlling the verbosity of logging output
+  //!
+  //! \param severity The logger will only emit messages that have severity of
+  //! this level or higher.
+  //!
+  void setReportableSeverity(Severity severity) {
+    mReportableSeverity = severity;
+  }
+
+  //!
+  //! \brief Opaque handle that holds logging information for a particular test
+  //!
+  //! This object is an opaque handle to information used by the Logger to print
+  //! test results. The sample must call Logger::defineTest() in order to obtain
+  //! a TestAtom that can be used with Logger::reportTest{Start,End}().
+  //!
+  class TestAtom {
+  public:
+    TestAtom(TestAtom &&) = default;
+
+  private:
+    friend class Logger;
+
+    TestAtom(bool started, const std::string &name, const std::string &cmdline)
+        : mStarted(started), mName(name), mCmdline(cmdline) {}
+
+    bool mStarted;
+    std::string mName;
+    std::string mCmdline;
+  };
+
+  //!
+  //! \brief Define a test for logging
+  //!
+  //! \param[in] name The name of the test.  This should be a string starting
+  //! with
+  //!                  "TensorRT" and containing dot-separated strings
+  //!                  containing the characters [A-Za-z0-9_]. For example,
+  //!                  "TensorRT.sample_googlenet"
+  //! \param[in] cmdline The command line used to reproduce the test
+  //
+  //! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
+  //!
+  static TestAtom defineTest(const std::string &name,
+                             const std::string &cmdline) {
+    return TestAtom(false, name, cmdline);
+  }
+
+  //!
+  //! \brief A convenience overloaded version of defineTest() that accepts an
+  //! array of command-line arguments
+  //!        as input
+  //!
+  //! \param[in] name The name of the test
+  //! \param[in] argc The number of command-line arguments
+  //! \param[in] argv The array of command-line arguments (given as C strings)
+  //!
+  //! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
+  static TestAtom defineTest(const std::string &name, int argc,
+                             char const *const *argv) {
+    auto cmdline = genCmdlineString(argc, argv);
+    return defineTest(name, cmdline);
+  }
+
+  //!
+  //! \brief Report that a test has started.
+  //!
+  //! \pre reportTestStart() has not been called yet for the given testAtom
+  //!
+  //! \param[in] testAtom The handle to the test that has started
+  //!
+  static void reportTestStart(TestAtom &testAtom) {
+    reportTestResult(testAtom, TestResult::kRUNNING);
+    assert(!testAtom.mStarted);
+    testAtom.mStarted = true;
+  }
+
+  //!
+  //! \brief Report that a test has ended.
+  //!
+  //! \pre reportTestStart() has been called for the given testAtom
+  //!
+  //! \param[in] testAtom The handle to the test that has ended
+  //! \param[in] result The result of the test. Should be one of
+  //! TestResult::kPASSED,
+  //!                   TestResult::kFAILED, TestResult::kWAIVED
+  //!
+  static void reportTestEnd(const TestAtom &testAtom, TestResult result) {
+    assert(result != TestResult::kRUNNING);
+    assert(testAtom.mStarted);
+    reportTestResult(testAtom, result);
+  }
+
+  static int reportPass(const TestAtom &testAtom) {
+    reportTestEnd(testAtom, TestResult::kPASSED);
+    return EXIT_SUCCESS;
+  }
+
+  static int reportFail(const TestAtom &testAtom) {
+    reportTestEnd(testAtom, TestResult::kFAILED);
+    return EXIT_FAILURE;
+  }
+
+  static int reportWaive(const TestAtom &testAtom) {
+    reportTestEnd(testAtom, TestResult::kWAIVED);
+    return EXIT_SUCCESS;
+  }
+
+  static int reportTest(const TestAtom &testAtom, bool pass) {
+    return pass ? reportPass(testAtom) : reportFail(testAtom);
+  }
+
+  Severity getReportableSeverity() const { return mReportableSeverity; }
+
+private:
+  //!
+  //! \brief returns an appropriate string for prefixing a log message with the
+  //! given severity
+  //!
+  static const char *severityPrefix(Severity severity) {
+    switch (severity) {
+    case Severity::kINTERNAL_ERROR:
+      return "[F] ";
+    case Severity::kERROR:
+      return "[E] ";
+    case Severity::kWARNING:
+      return "[W] ";
+    case Severity::kINFO:
+      return "[I] ";
+    case Severity::kVERBOSE:
+      return "[V] ";
+    default:
+      assert(0);
+      return "";
+    }
+  }
+
+  //!
+  //! \brief returns an appropriate string for prefixing a test result message
+  //! with the given result
+  //!
+  static const char *testResultString(TestResult result) {
+    switch (result) {
+    case TestResult::kRUNNING:
+      return "RUNNING";
+    case TestResult::kPASSED:
+      return "PASSED";
+    case TestResult::kFAILED:
+      return "FAILED";
+    case TestResult::kWAIVED:
+      return "WAIVED";
+    default:
+      assert(0);
+      return "";
+    }
+  }
+
+  //!
+  //! \brief returns an appropriate output stream (cout or cerr) to use with the
+  //! given severity
+  //!
+  static std::ostream &severityOstream(Severity severity) {
+    return severity >= Severity::kINFO ? std::cout : std::cerr;
+  }
+
+  //!
+  //! \brief method that implements logging test results
+  //!
+  static void reportTestResult(const TestAtom &testAtom, TestResult result) {
+    severityOstream(Severity::kINFO)
+        << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
+        << testAtom.mCmdline << std::endl;
+  }
+
+  //!
+  //! \brief generate a command line string from the given (argc, argv) values
+  //!
+  static std::string genCmdlineString(int argc, char const *const *argv) {
+    std::stringstream ss;
+    for (int i = 0; i < argc; i++) {
+      if (i > 0)
+        ss << " ";
+      ss << argv[i];
+    }
+    return ss.str();
+  }
+
+  Severity mReportableSeverity;
+};
+
+namespace {
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages
+//! of severity kVERBOSE
+//!
+//! Example usage:
+//!
+//!     LOG_VERBOSE(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_VERBOSE(const Logger &logger) {
+  return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages
+//! of severity kINFO
+//!
+//! Example usage:
+//!
+//!     LOG_INFO(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_INFO(const Logger &logger) {
+  return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages
+//! of severity kWARNING
+//!
+//! Example usage:
+//!
+//!     LOG_WARN(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_WARN(const Logger &logger) {
+  return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages
+//! of severity kERROR
+//!
+//! Example usage:
+//!
+//!     LOG_ERROR(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_ERROR(const Logger &logger) {
+  return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages
+//! of severity kINTERNAL_ERROR
+//         ("fatal" severity)
+//!
+//! Example usage:
+//!
+//!     LOG_FATAL(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_FATAL(const Logger &logger) {
+  return LogStreamConsumer(logger.getReportableSeverity(),
+                           Severity::kINTERNAL_ERROR);
+}
+
+} // anonymous namespace
+
+#endif // TENSORRT_LOGGING_H
diff --git a/csrnet/macros.h b/csrnet/macros.h
new file mode 100644
index 0000000..0555103
--- /dev/null
+++ b/csrnet/macros.h
@@ -0,0 +1,12 @@
+#ifndef __MACROS_H
+#define __MACROS_H
+
+#if NV_TENSORRT_MAJOR >= 8
+#define TRT_NOEXCEPT noexcept
+#define TRT_CONST_ENQUEUE const
+#else
+#define TRT_NOEXCEPT
+#define TRT_CONST_ENQUEUE
+#endif
+
+#endif  // __MACROS_H