-
Notifications
You must be signed in to change notification settings - Fork 122
/
engine.h
49 lines (40 loc) · 1.57 KB
/
engine.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
#ifndef EDGETPU_CPP_DETECTION_ENGINE_H_
#define EDGETPU_CPP_DETECTION_ENGINE_H_
#include <vector>
#include "src/cpp/basic/basic_engine.h"
#include "src/cpp/bbox_utils.h"
namespace coral {
class DetectionEngine : public BasicEngine {
public:
// Loads detection model. Now we only support SSD model with postprocessing
// operator.
// - 'model_path' : the file path of the model.
explicit DetectionEngine(const std::string& model_path)
: BasicEngine(model_path) {
Validate();
}
// Loads detection model and specifies EdgeTpu to use.
// - 'model_path' : the file path of the model.
// - 'device_path' : the device path of EdgeTpu.
explicit DetectionEngine(const std::string& model_path,
const std::string& device_path)
: BasicEngine(model_path, device_path) {
Validate();
}
// Detects objects with input tensor.
// - 'input' : vector of uint8, input to the model.
// - 'threshold' : float, minimum confidence threshold for returned
// predictions. For example, use 0.5 to receive only predictions
// with a confidence equal-to or higher-than 0.5.
// - 'top_k': int, the maximum number of predictions to return.
//
// The function will return a vector of predictions which is sorted by
// <score, label_id> in descending order.
std::vector<DetectionCandidate> DetectWithInputTensor(
const std::vector<uint8_t>& input, float threshold = 0.0, int top_k = 3);
private:
// Checks the format of the model.
void Validate();
};
} // namespace coral
#endif // EDGETPU_CPP_DETECTION_ENGINE_H_