main.cc 14.1 KB
Newer Older
Q
qingqing01 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
//   Copyright (c) 2020 PaddlePaddle Authors. 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.

#include <glog/logging.h>

G
Guanghua Yu 已提交
17
#include <dirent.h>
Q
qingqing01 已提交
18 19 20
#include <iostream>
#include <string>
#include <vector>
G
Guanghua Yu 已提交
21
#include <numeric>
Q
qingqing01 已提交
22 23
#include <sys/types.h>
#include <sys/stat.h>
C
cnn 已提交
24
#include <math.h>
G
Guanghua Yu 已提交
25
#include <algorithm>
Q
qingqing01 已提交
26 27 28 29 30 31 32 33 34 35

#ifdef _WIN32
#include <direct.h>
#include <io.h>
#elif LINUX
#include <stdarg.h>
#include <sys/stat.h>
#endif

#include "include/object_detector.h"
36
#include <gflags/gflags.h>
Q
qingqing01 已提交
37 38 39


DEFINE_string(model_dir, "", "Path of inference model");
G
Guanghua Yu 已提交
40 41
DEFINE_string(image_file, "", "Path of input image");
DEFINE_string(image_dir, "", "Dir of input image, `image_file` has a higher priority.");
C
cnn 已提交
42
DEFINE_int32(batch_size, 1, "batch_size");
G
Guanghua Yu 已提交
43 44
DEFINE_string(video_file, "", "Path of input video, `video_file` or `camera_id` has a highest priority.");
DEFINE_int32(camera_id, -1, "Device id of camera to predict");
G
Guanghua Yu 已提交
45 46
DEFINE_bool(use_gpu, false, "Deprecated, please use `--device` to set the device you want to run.");
DEFINE_string(device, "CPU", "Choose the device you want to run, it can be: CPU/GPU/XPU, default is CPU.");
G
Guanghua Yu 已提交
47 48
DEFINE_double(threshold, 0.5, "Threshold of score.");
DEFINE_string(output_dir, "output", "Directory of output visualization files.");
49
DEFINE_string(run_mode, "fluid", "Mode of running(fluid/trt_fp32/trt_fp16/trt_int8)");
Q
qingqing01 已提交
50 51
DEFINE_int32(gpu_id, 0, "Device id of GPU to execute");
DEFINE_bool(run_benchmark, false, "Whether to predict a image_file repeatedly for benchmark");
G
Guanghua Yu 已提交
52 53
DEFINE_bool(use_mkldnn, false, "Whether use mkldnn with CPU");
DEFINE_int32(cpu_threads, 1, "Num of threads with CPU");
54 55 56
DEFINE_int32(trt_min_shape, 1, "Min shape of TRT DynamicShapeI");
DEFINE_int32(trt_max_shape, 1280, "Max shape of TRT DynamicShapeI");
DEFINE_int32(trt_opt_shape, 640, "Opt shape of TRT DynamicShapeI");
G
Guanghua Yu 已提交
57 58 59 60
DEFINE_bool(trt_calib_mode, false, "If the model is produced by TRT offline quantitative calibration, trt_calib_mode need to set True");

void PrintBenchmarkLog(std::vector<double> det_time, int img_num){
  LOG(INFO) << "----------------------- Config info -----------------------";
G
Guanghua Yu 已提交
61
  LOG(INFO) << "runtime_device: " << FLAGS_device;
G
Guanghua Yu 已提交
62 63 64 65 66 67 68 69 70 71 72 73 74 75
  LOG(INFO) << "ir_optim: " << "True";
  LOG(INFO) << "enable_memory_optim: " << "True";
  int has_trt = FLAGS_run_mode.find("trt");
  if (has_trt >= 0) {
    LOG(INFO) << "enable_tensorrt: " << "True";
    std::string precision = FLAGS_run_mode.substr(4, 8);
    LOG(INFO) << "precision: " << precision;
  } else {
    LOG(INFO) << "enable_tensorrt: " << "False";
    LOG(INFO) << "precision: " << "fp32";
  }
  LOG(INFO) << "enable_mkldnn: " << (FLAGS_use_mkldnn ? "True" : "False");
  LOG(INFO) << "cpu_math_library_num_threads: " << FLAGS_cpu_threads;
  LOG(INFO) << "----------------------- Data info -----------------------";
76
  LOG(INFO) << "batch_size: " << FLAGS_batch_size;
G
Guanghua Yu 已提交
77 78 79 80 81 82
  LOG(INFO) << "input_shape: " << "dynamic shape";
  LOG(INFO) << "----------------------- Model info -----------------------";
  FLAGS_model_dir.erase(FLAGS_model_dir.find_last_not_of("/") + 1);
  LOG(INFO) << "model_name: " << FLAGS_model_dir.substr(FLAGS_model_dir.find_last_of('/') + 1);
  LOG(INFO) << "----------------------- Perf info ------------------------";
  LOG(INFO) << "Total number of predicted data: " << img_num
G
Guanghua Yu 已提交
83
            << " and total time spent(ms): "
G
Guanghua Yu 已提交
84 85 86 87 88
            << std::accumulate(det_time.begin(), det_time.end(), 0);
  LOG(INFO) << "preproce_time(ms): " << det_time[0] / img_num
            << ", inference_time(ms): " << det_time[1] / img_num
            << ", postprocess_time(ms): " << det_time[2];
}
Q
qingqing01 已提交
89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130

static std::string DirName(const std::string &filepath) {
  auto pos = filepath.rfind(OS_PATH_SEP);
  if (pos == std::string::npos) {
    return "";
  }
  return filepath.substr(0, pos);
}

static bool PathExists(const std::string& path){
#ifdef _WIN32
  struct _stat buffer;
  return (_stat(path.c_str(), &buffer) == 0);
#else
  struct stat buffer;
  return (stat(path.c_str(), &buffer) == 0);
#endif  // !_WIN32
}

static void MkDir(const std::string& path) {
  if (PathExists(path)) return;
  int ret = 0;
#ifdef _WIN32
  ret = _mkdir(path.c_str());
#else
  ret = mkdir(path.c_str(), 0755);
#endif  // !_WIN32
  if (ret != 0) {
    std::string path_error(path);
    path_error += " mkdir failed!";
    throw std::runtime_error(path_error);
  }
}

static void MkDirs(const std::string& path) {
  if (path.empty()) return;
  if (PathExists(path)) return;

  MkDirs(DirName(path));
  MkDir(path);
}

G
Guanghua Yu 已提交
131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161
void GetAllFiles(const char *dir_name,
                          std::vector<std::string> &all_inputs) {
  if (NULL == dir_name) {
    std::cout << " dir_name is null ! " << std::endl;
    return;
  }
  struct stat s;
  lstat(dir_name, &s);
  if (!S_ISDIR(s.st_mode)) {
    std::cout << "dir_name is not a valid directory !" << std::endl;
    all_inputs.push_back(dir_name);
    return;
  } else {
    struct dirent *filename; // return value for readdir()
    DIR *dir;                // return value for opendir()
    dir = opendir(dir_name);
    if (NULL == dir) {
      std::cout << "Can not open dir " << dir_name << std::endl;
      return;
    }
    std::cout << "Successfully opened the dir !" << std::endl;
    while ((filename = readdir(dir)) != NULL) {
      if (strcmp(filename->d_name, ".") == 0 ||
          strcmp(filename->d_name, "..") == 0)
        continue;
      all_inputs.push_back(dir_name + std::string("/") +
                           std::string(filename->d_name));
    }
  }
}

Q
qingqing01 已提交
162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194
void PredictVideo(const std::string& video_path,
                  PaddleDetection::ObjectDetector* det) {
  // Open video
  cv::VideoCapture capture;
  if (FLAGS_camera_id != -1){
    capture.open(FLAGS_camera_id);
  }else{
    capture.open(video_path.c_str());
  }
  if (!capture.isOpened()) {
    printf("can not open video : %s\n", video_path.c_str());
    return;
  }

  // Get Video info : resolution, fps
  int video_width = static_cast<int>(capture.get(CV_CAP_PROP_FRAME_WIDTH));
  int video_height = static_cast<int>(capture.get(CV_CAP_PROP_FRAME_HEIGHT));
  int video_fps = static_cast<int>(capture.get(CV_CAP_PROP_FPS));

  // Create VideoWriter for output
  cv::VideoWriter video_out;
  std::string video_out_path = "output.mp4";
  video_out.open(video_out_path.c_str(),
                 0x00000021,
                 video_fps,
                 cv::Size(video_width, video_height),
                 true);
  if (!video_out.isOpened()) {
    printf("create video writer failed!\n");
    return;
  }

  std::vector<PaddleDetection::ObjectResult> result;
C
cnn 已提交
195
  std::vector<int> bbox_num;
G
Guanghua Yu 已提交
196
  std::vector<double> det_times;
Q
qingqing01 已提交
197 198 199 200 201
  auto labels = det->GetLabelList();
  auto colormap = PaddleDetection::GenerateColorMap(labels.size());
  // Capture all frames and do inference
  cv::Mat frame;
  int frame_id = 0;
C
cnn 已提交
202
  bool is_rbox = false;
Q
qingqing01 已提交
203 204 205 206
  while (capture.read(frame)) {
    if (frame.empty()) {
      break;
    }
C
cnn 已提交
207 208 209
    std::vector<cv::Mat> imgs;
    imgs.push_back(frame);
    det->Predict(imgs, 0.5, 0, 1, &result, &bbox_num, &det_times);
Q
qingqing01 已提交
210
    for (const auto& item : result) {
C
cnn 已提交
211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238
      if (item.rect.size() > 6){
      is_rbox = true;
      printf("class=%d confidence=%.4f rect=[%d %d %d %d %d %d %d %d]\n",
          item.class_id,
          item.confidence,
          item.rect[0],
          item.rect[1],
          item.rect[2],
          item.rect[3],
          item.rect[4],
          item.rect[5],
          item.rect[6],
          item.rect[7]);
      }
      else{
        printf("class=%d confidence=%.4f rect=[%d %d %d %d]\n",
          item.class_id,
          item.confidence,
          item.rect[0],
          item.rect[1],
          item.rect[2],
          item.rect[3]);
      }
   }

   cv::Mat out_im = PaddleDetection::VisualizeResult(
        frame, result, labels, colormap, is_rbox);

Q
qingqing01 已提交
239 240 241 242 243 244 245
    video_out.write(out_im);
    frame_id += 1;
  }
  capture.release();
  video_out.release();
}

C
cnn 已提交
246 247
void PredictImage(const std::vector<std::string> all_img_paths,
                  const int batch_size,
Q
qingqing01 已提交
248 249 250 251
                  const double threshold,
                  const bool run_benchmark,
                  PaddleDetection::ObjectDetector* det,
                  const std::string& output_dir = "output") {
G
Guanghua Yu 已提交
252
  std::vector<double> det_t = {0, 0, 0};
C
cnn 已提交
253 254 255 256 257 258 259 260 261 262 263 264 265 266 267
  int steps = ceil(float(all_img_paths.size()) / batch_size);
  printf("total images = %d, batch_size = %d, total steps = %d\n",
                all_img_paths.size(), batch_size, steps);
  for (int idx = 0; idx < steps; idx++) {
    std::vector<cv::Mat> batch_imgs;
    int left_image_cnt = all_img_paths.size() - idx * batch_size;
    if (left_image_cnt > batch_size) {
      left_image_cnt = batch_size;
    }
    for (int bs = 0; bs < left_image_cnt; bs++) {
      std::string image_file_path = all_img_paths.at(idx * batch_size+bs);
      cv::Mat im = cv::imread(image_file_path, 1);
      batch_imgs.insert(batch_imgs.end(), im);
    }
    
G
Guanghua Yu 已提交
268 269
    // Store all detected result
    std::vector<PaddleDetection::ObjectResult> result;
C
cnn 已提交
270
    std::vector<int> bbox_num;
G
Guanghua Yu 已提交
271
    std::vector<double> det_times;
C
cnn 已提交
272
    bool is_rbox = false;
G
Guanghua Yu 已提交
273
    if (run_benchmark) {
C
cnn 已提交
274
      det->Predict(batch_imgs, threshold, 10, 10, &result, &bbox_num,  &det_times);
G
Guanghua Yu 已提交
275
    } else {
C
cnn 已提交
276 277 278 279 280 281 282 283 284 285
      det->Predict(batch_imgs, 0.5, 0, 1, &result, &bbox_num, &det_times);
      // get labels and colormap
      auto labels = det->GetLabelList();
      auto colormap = PaddleDetection::GenerateColorMap(labels.size());

      int item_start_idx = 0;
      for (int i = 0; i < left_image_cnt; i++) {
        std::cout << all_img_paths.at(idx * batch_size + i) << "result" << std::endl;
        if (bbox_num[i] <= 1) {
            continue;
C
cnn 已提交
286
        }
C
cnn 已提交
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311
        for (int j = 0; j < bbox_num[i]; j++) {
          PaddleDetection::ObjectResult item = result[item_start_idx + j];
          if (item.rect.size() > 6){
            is_rbox = true;
            printf("class=%d confidence=%.4f rect=[%d %d %d %d %d %d %d %d]\n",
              item.class_id,
              item.confidence,
              item.rect[0],
              item.rect[1],
              item.rect[2],
              item.rect[3],
              item.rect[4],
              item.rect[5],
              item.rect[6],
              item.rect[7]);
          }
          else{
            printf("class=%d confidence=%.4f rect=[%d %d %d %d]\n",
              item.class_id,
              item.confidence,
              item.rect[0],
              item.rect[1],
              item.rect[2],
              item.rect[3]);
          }
C
cnn 已提交
312
        }
C
cnn 已提交
313
        item_start_idx = item_start_idx + bbox_num[i];
G
Guanghua Yu 已提交
314 315
      }
      // Visualization result
C
cnn 已提交
316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335
      int bbox_idx = 0;
      for (int bs = 0; bs < batch_imgs.size(); bs++) {
        if (bbox_num[bs] <= 1) {
            continue;
        }
        cv::Mat im = batch_imgs[bs];
        std::vector<PaddleDetection::ObjectResult> im_result;
        for (int k = 0; k < bbox_num[bs]; k++) {
          im_result.push_back(result[bbox_idx+k]);
        }
        bbox_idx += bbox_num[bs];
        cv::Mat vis_img = PaddleDetection::VisualizeResult(
            im, im_result, labels, colormap, is_rbox);
        std::vector<int> compression_params;
        compression_params.push_back(CV_IMWRITE_JPEG_QUALITY);
        compression_params.push_back(95);
        std::string output_path(output_dir);
        if (output_dir.rfind(OS_PATH_SEP) != output_dir.size() - 1) {
          output_path += OS_PATH_SEP;
        }
336
        std::string image_file_path = all_img_paths.at(idx * batch_size + bs);
C
cnn 已提交
337 338 339
        output_path += image_file_path.substr(image_file_path.find_last_of('/') + 1);
        cv::imwrite(output_path, vis_img, compression_params);
        printf("Visualized output saved as %s\n", output_path.c_str());
G
Guanghua Yu 已提交
340
      }
Q
qingqing01 已提交
341
    }
G
Guanghua Yu 已提交
342 343 344
    det_t[0] += det_times[0];
    det_t[1] += det_times[1];
    det_t[2] += det_times[2];
Q
qingqing01 已提交
345
  }
C
cnn 已提交
346
  PrintBenchmarkLog(det_t, all_img_paths.size());
Q
qingqing01 已提交
347 348 349 350 351 352
}

int main(int argc, char** argv) {
  // Parsing command-line
  google::ParseCommandLineFlags(&argc, &argv, true);
  if (FLAGS_model_dir.empty()
G
Guanghua Yu 已提交
353
      || (FLAGS_image_file.empty() && FLAGS_image_dir.empty() && FLAGS_video_file.empty())) {
Q
qingqing01 已提交
354
    std::cout << "Usage: ./main --model_dir=/PATH/TO/INFERENCE_MODEL/ "
G
Guanghua Yu 已提交
355
                << "--image_file=/PATH/TO/INPUT/IMAGE/" << std::endl;
Q
qingqing01 已提交
356 357 358
    return -1;
  }
  if (!(FLAGS_run_mode == "fluid" || FLAGS_run_mode == "trt_fp32"
359 360
      || FLAGS_run_mode == "trt_fp16" || FLAGS_run_mode == "trt_int8")) {
    std::cout << "run_mode should be 'fluid', 'trt_fp32', 'trt_fp16' or 'trt_int8'.";
Q
qingqing01 已提交
361 362
    return -1;
  }
G
Guanghua Yu 已提交
363 364 365 366 367 368 369 370 371
  transform(FLAGS_device.begin(),FLAGS_device.end(),FLAGS_device.begin(),::toupper);
  if (!(FLAGS_device == "CPU" || FLAGS_device == "GPU" || FLAGS_device == "XPU")) {
    std::cout << "device should be 'CPU', 'GPU' or 'XPU'.";
    return -1;
  }
  if (FLAGS_use_gpu) {
    std::cout << "Deprecated, please use `--device` to set the device you want to run.";
    return -1;
  }
Q
qingqing01 已提交
372
  // Load model and create a object detector
G
Guanghua Yu 已提交
373
  PaddleDetection::ObjectDetector det(FLAGS_model_dir, FLAGS_device, FLAGS_use_mkldnn,
374 375 376
                        FLAGS_cpu_threads, FLAGS_run_mode, FLAGS_batch_size,FLAGS_gpu_id,
                        FLAGS_trt_min_shape, FLAGS_trt_max_shape, FLAGS_trt_opt_shape,
			FLAGS_trt_calib_mode);
Q
qingqing01 已提交
377
  // Do inference on input video or image
G
Guanghua Yu 已提交
378 379 380
  if (!FLAGS_video_file.empty() || FLAGS_camera_id != -1) {
    PredictVideo(FLAGS_video_file, &det);
  } else if (!FLAGS_image_file.empty() || !FLAGS_image_dir.empty()) {
Q
qingqing01 已提交
381 382 383
    if (!PathExists(FLAGS_output_dir)) {
      MkDirs(FLAGS_output_dir);
    }
C
cnn 已提交
384
    std::vector<std::string> all_imgs;
G
Guanghua Yu 已提交
385
    if (!FLAGS_image_file.empty()) {
C
cnn 已提交
386 387
      all_imgs.push_back(FLAGS_image_file);
      if (FLAGS_batch_size > 1) {
388 389
        std::cout << "batch_size should be 1, when set `image_file`." << std::endl;
	return -1;
C
cnn 已提交
390
      }
G
Guanghua Yu 已提交
391
    } else {
C
cnn 已提交
392
      GetAllFiles((char *)FLAGS_image_dir.c_str(), all_imgs);
G
Guanghua Yu 已提交
393
    }
394 395
    PredictImage(all_imgs, FLAGS_batch_size, FLAGS_threshold,
		 FLAGS_run_benchmark, &det, FLAGS_output_dir);
Q
qingqing01 已提交
396 397 398
  }
  return 0;
}