utils.h 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. #ifndef TRTX_YOLOV5_UTILS_H_
  2. #define TRTX_YOLOV5_UTILS_H_
  3. #include <dirent.h>
  4. #include <opencv2/opencv.hpp>
  5. static inline cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h) {
  6. int w, h, x, y;
  7. float r_w = input_w / (img.cols*1.0);
  8. float r_h = input_h / (img.rows*1.0);
  9. if (r_h > r_w) {
  10. w = input_w;
  11. h = r_w * img.rows;
  12. x = 0;
  13. y = (input_h - h) / 2;
  14. } else {
  15. w = r_h * img.cols;
  16. h = input_h;
  17. x = (input_w - w) / 2;
  18. y = 0;
  19. }
  20. cv::Mat re(h, w, CV_8UC3);
  21. cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR);
  22. cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(128, 128, 128));
  23. re.copyTo(out(cv::Rect(x, y, re.cols, re.rows)));
  24. return out;
  25. }
  26. static inline int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_names) {
  27. DIR *p_dir = opendir(p_dir_name);
  28. if (p_dir == nullptr) {
  29. return -1;
  30. }
  31. struct dirent* p_file = nullptr;
  32. while ((p_file = readdir(p_dir)) != nullptr) {
  33. if (strcmp(p_file->d_name, ".") != 0 &&
  34. strcmp(p_file->d_name, "..") != 0) {
  35. //std::string cur_file_name(p_dir_name);
  36. //cur_file_name += "/";
  37. //cur_file_name += p_file->d_name;
  38. std::string cur_file_name(p_file->d_name);
  39. file_names.push_back(cur_file_name);
  40. }
  41. }
  42. closedir(p_dir);
  43. return 0;
  44. }
  45. #endif // TRTX_YOLOV5_UTILS_H_