Caffe2 - C++ API
A deep learning, cross platform ML framework
box_with_nms_limit_op.h
1 // Copyright 2004-present Facebook. All Rights Reserved.
2 
3 #ifndef BOX_WITH_NMS_AND_LIMIT_OP_H_
4 #define BOX_WITH_NMS_AND_LIMIT_OP_H_
5 
6 #include "caffe2/core/context.h"
7 #include "caffe2/core/operator.h"
8 
9 namespace caffe2 {
10 
11 // C++ implementation of function insert_box_results_with_nms_and_limit()
12 template <class Context>
13 class BoxWithNMSLimitOp final : public Operator<Context> {
14  public:
15  USE_OPERATOR_CONTEXT_FUNCTIONS;
16  BoxWithNMSLimitOp(const OperatorDef& operator_def, Workspace* ws)
17  : Operator<Context>(operator_def, ws),
18  score_thres_(
19  OperatorBase::GetSingleArgument<float>("score_thresh", 0.05)),
20  nms_thres_(OperatorBase::GetSingleArgument<float>("nms", 0.3)),
21  detections_per_im_(
22  OperatorBase::GetSingleArgument<int>("detections_per_im", 100)),
23  soft_nms_enabled_(
24  OperatorBase::GetSingleArgument<bool>("soft_nms_enabled", false)),
25  soft_nms_method_str_(OperatorBase::GetSingleArgument<std::string>(
26  "soft_nms_method",
27  "linear")),
28  soft_nms_sigma_(
29  OperatorBase::GetSingleArgument<float>("soft_nms_sigma", 0.5)),
30  soft_nms_min_score_thres_(OperatorBase::GetSingleArgument<float>(
31  "soft_nms_min_score_thres",
32  0.001)) {
33  CAFFE_ENFORCE(
34  soft_nms_method_str_ == "linear" || soft_nms_method_str_ == "gaussian",
35  "Unexpected soft_nms_method");
36  soft_nms_method_ = (soft_nms_method_str_ == "linear") ? 1 : 2;
37  }
38 
39  ~BoxWithNMSLimitOp() {}
40 
41  bool RunOnDevice() override;
42 
43  protected:
44  // TEST.SCORE_THRESH
45  float score_thres_ = 0.05;
46  // TEST.NMS
47  float nms_thres_ = 0.3;
48  // TEST.DETECTIONS_PER_IM
49  int detections_per_im_ = 100;
50  // TEST.SOFT_NMS.ENABLED
51  bool soft_nms_enabled_ = false;
52  // TEST.SOFT_NMS.METHOD
53  std::string soft_nms_method_str_ = "linear";
54  unsigned int soft_nms_method_ = 1; // linear
55  // TEST.SOFT_NMS.SIGMA
56  float soft_nms_sigma_ = 0.5;
57  // Lower-bound on updated scores to discard boxes
58  float soft_nms_min_score_thres_ = 0.001;
59 };
60 
61 } // namespace caffe2
62 #endif // BOX_WITH_NMS_AND_LIMIT_OP_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.