main-yolov3.cpp 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. #include <rtthread.h>
  2. #include <rtdevice.h>
  3. #include <board.h>
  4. #include "layer.h"
  5. #include "net.h"
  6. #include "simpleocv.h"
  7. #include <float.h>
  8. #include <stdio.h>
  9. #include <vector>
  10. #include <stdio.h>
  11. #include <vector>
  12. struct Object
  13. {
  14. cv::Rect_<float> rect;
  15. int label;
  16. float prob;
  17. };
  18. static int detect_yolov3(const cv::Mat& bgr, std::vector<Object>& objects)
  19. {
  20. ncnn::Net yolov3;
  21. yolov3.opt.use_vulkan_compute = true;
  22. // original pretrained model from https://github.com/eric612/MobileNet-YOLO
  23. // param : https://drive.google.com/open?id=1V9oKHP6G6XvXZqhZbzNKL6FI_clRWdC-
  24. // bin : https://drive.google.com/open?id=1DBcuFCr-856z3FRQznWL_S5h-Aj3RawA
  25. // the ncnn model https://github.com/nihui/ncnn-assets/tree/master/models
  26. if (yolov3.load_param("mobilenetv2_yolov3.param"))
  27. exit(-1);
  28. if (yolov3.load_model("mobilenetv2_yolov3.bin"))
  29. exit(-1);
  30. const int target_size = 352;
  31. int img_w = bgr.cols;
  32. int img_h = bgr.rows;
  33. ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR, bgr.cols, bgr.rows, target_size, target_size);
  34. const float mean_vals[3] = {127.5f, 127.5f, 127.5f};
  35. const float norm_vals[3] = {0.007843f, 0.007843f, 0.007843f};
  36. in.substract_mean_normalize(mean_vals, norm_vals);
  37. ncnn::Extractor ex = yolov3.create_extractor();
  38. ex.input("data", in);
  39. ncnn::Mat out;
  40. ex.extract("detection_out", out);
  41. // printf("%d %d %d\n", out.w, out.h, out.c);
  42. objects.clear();
  43. for (int i = 0; i < out.h; i++)
  44. {
  45. const float* values = out.row(i);
  46. Object object;
  47. object.label = values[0];
  48. object.prob = values[1];
  49. object.rect.x = values[2] * img_w;
  50. object.rect.y = values[3] * img_h;
  51. object.rect.width = values[4] * img_w - object.rect.x;
  52. object.rect.height = values[5] * img_h - object.rect.y;
  53. objects.push_back(object);
  54. }
  55. return 0;
  56. }
  57. static void draw_objects(const cv::Mat& bgr, const std::vector<Object>& objects)
  58. {
  59. static const char* class_names[] = {"background",
  60. "aeroplane", "bicycle", "bird", "boat",
  61. "bottle", "bus", "car", "cat", "chair",
  62. "cow", "diningtable", "dog", "horse",
  63. "motorbike", "person", "pottedplant",
  64. "sheep", "sofa", "train", "tvmonitor"
  65. };
  66. cv::Mat image = bgr.clone();
  67. for (size_t i = 0; i < objects.size(); i++)
  68. {
  69. const Object& obj = objects[i];
  70. fprintf(stderr, "%d = %.5f at %.2f %.2f %.2f x %.2f\n", obj.label, obj.prob,
  71. obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height);
  72. cv::rectangle(image, obj.rect, cv::Scalar(255, 0, 0));
  73. char text[256];
  74. sprintf(text, "%s %.1f%%", class_names[obj.label], obj.prob * 100);
  75. int baseLine = 0;
  76. cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 1, 1, &baseLine);
  77. int x = obj.rect.x;
  78. int y = obj.rect.y - label_size.height - baseLine;
  79. if (y < 0)
  80. y = 0;
  81. if (x + label_size.width > image.cols)
  82. x = image.cols - label_size.width;
  83. cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
  84. cv::Scalar(255, 0, 0), -1);
  85. cv::putText(image, text, cv::Point(x, y + label_size.height),
  86. cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 0, 0));
  87. }
  88. cv::imwrite("image1.png", image);
  89. }
  90. void mnet_yolov3_test()
  91. {
  92. rt_kprintf("Hello RT-Thread NCNN\n");
  93. cv::Mat m = cv::imread("bus.jpg", 1);
  94. if (m.empty())
  95. {
  96. rt_kprintf("cv::imread failed\n");
  97. }
  98. std::vector<Object> objects;
  99. detect_yolov3(m, objects);
  100. draw_objects(m, objects);
  101. }
  102. MSH_CMD_EXPORT(mnet_yolov3_test, test yolo);
  103. int main(int argc, char** argv)
  104. {
  105. rt_kprintf("Hello RT-Thread NCNN\n");
  106. return 0;
  107. }