yolo.hpp 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. #ifndef __YOLO_HPP__
  2. #define __YOLO_HPP__
  3. #include <future>
  4. #include <memory>
  5. #include <string>
  6. #include <vector>
  7. namespace yolo {
  8. enum class Type : int {
  9. V5 = 0,
  10. X = 1,
  11. V3 = 2,
  12. V7 = 3,
  13. V8 = 5,
  14. V8Seg = 6 // yolov8 instance segmentation
  15. };
  16. struct InstanceSegmentMap {
  17. int width = 0, height = 0; // width % 8 == 0
  18. unsigned char *data = nullptr; // is width * height memory
  19. InstanceSegmentMap(int width, int height);
  20. virtual ~InstanceSegmentMap();
  21. };
  22. struct Box {
  23. float left, top, right, bottom, confidence;
  24. int class_label;
  25. std::shared_ptr<InstanceSegmentMap> seg; // valid only in segment task
  26. Box() = default;
  27. Box(float left, float top, float right, float bottom, float confidence, int class_label)
  28. : left(left),
  29. top(top),
  30. right(right),
  31. bottom(bottom),
  32. confidence(confidence),
  33. class_label(class_label) {}
  34. };
  35. struct Image {
  36. const void *bgrptr = nullptr;
  37. int width = 0, height = 0;
  38. Image() = default;
  39. Image(const void *bgrptr, int width, int height) : bgrptr(bgrptr), width(width), height(height) {}
  40. };
  41. typedef std::vector<Box> BoxArray;
  42. // [Preprocess]: 0.50736 ms
  43. // [Forward]: 3.96410 ms
  44. // [BoxDecode]: 0.12016 ms
  45. // [SegmentDecode]: 0.15610 ms
  46. class Infer {
  47. public:
  48. virtual BoxArray forward(const Image &image, void *stream = nullptr) = 0;
  49. virtual std::vector<BoxArray> forwards(const std::vector<Image> &images,
  50. void *stream = nullptr) = 0;
  51. };
  52. std::shared_ptr<Infer> load(const std::string &engine_file, Type type,
  53. float confidence_threshold = 0.25f, float nms_threshold = 0.5f);
  54. Infer *loadraw(const std::string &engine_file, Type type,
  55. float confidence_threshold = 0.25f, float nms_threshold = 0.5f);
  56. const char *type_name(Type type);
  57. std::tuple<uint8_t, uint8_t, uint8_t> hsv2bgr(float h, float s, float v);
  58. std::tuple<uint8_t, uint8_t, uint8_t> random_color(int id);
  59. }; // namespace yolo
  60. #endif // __YOLO_HPP__