main-mobilenet.cpp 4.0 KB

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