Merged from Alexey/darknet

uf-mil · Sep 19, 2021 · c5bcec0 · c5bcec0
1 parent 5d213fb
commit c5bcec0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/include/darknet.h b/include/darknet.h
@@ -770,7 +770,7 @@ typedef struct network {
     float label_smooth_eps;
     int resize_step;
     int attention;
-    int adversarial;    
+    int adversarial;
     float adversarial_lr;
     float max_chart_loss;
     int letter_box;
@@ -821,7 +821,7 @@ typedef struct network {
     size_t *max_input16_size;
     size_t *max_output16_size;
     int wait_stream;
-    
+
     void *cuda_graph;
     void *cuda_graph_exec;
     int use_cuda_graph;
@@ -899,6 +899,7 @@ typedef struct ious {
 typedef struct detection{
     box bbox;
     int classes;
+    int best_class_idx;
     float *prob;
     float *mask;
     float objectness;
@@ -908,7 +909,7 @@ typedef struct detection{
     float *embeddings;  // embeddings for tracking
     int embedding_size;
     float sim;
-    int track_id;    
+    int track_id;
 } detection;
 
 // network.c -batch inference
@@ -1017,7 +1018,6 @@ typedef struct box_label {
 // parser.c
 LIB_API network *load_network(char *cfg, char *weights, int clear);
 LIB_API network *load_network_custom(char *cfg, char *weights, int clear, int batch);
-LIB_API network *load_network(char *cfg, char *weights, int clear);
 LIB_API void free_network(network net);
 LIB_API void free_network_ptr(network* net);
 
@@ -1047,7 +1047,7 @@ LIB_API void reset_rnn(network *net);
 LIB_API float *network_predict_image(network *net, image im);
 LIB_API float *network_predict_image_letterbox(network *net, image im);
 LIB_API float validate_detector_map(char *datacfg, char *cfgfile, char *weightfile, float thresh_calc_avg_iou, const float iou_thresh, const int map_points, int letter_box, network *existing_net);
-LIB_API void train_detector(char *datacfg, char *cfgfile, char *weightfile, int *gpus, int ngpus, int clear, int dont_show, int calc_map, int mjpeg_port, int show_imgs, int benchmark_layers, char* chart_path);
+LIB_API void train_detector(char *datacfg, char *cfgfile, char *weightfile, int *gpus, int ngpus, int clear, int dont_show, int calc_map, float thresh, float iou_thresh, int mjpeg_port, int show_imgs, int benchmark_layers, char* chart_path);
 LIB_API void test_detector(char *datacfg, char *cfgfile, char *weightfile, char *filename, float thresh,
     float hier_thresh, int dont_show, int ext_output, int save_labels, char *outfile, int letter_box, int benchmark_layers);
 LIB_API int network_width(network *net);

diff --git a/include/yolo_v2_class.hpp b/include/yolo_v2_class.hpp
@@ -56,7 +56,7 @@ struct bbox_t_container {
 #include <opencv2/imgproc/imgproc_c.h>   // C
 #endif
 
-extern "C" LIB_API int init(const char *configurationFilename, const char *weightsFilename, int gpu);
+extern "C" LIB_API int init(const char *configurationFilename, const char *weightsFilename, int gpu, int batch_size);
 extern "C" LIB_API int detect_image(const char *filename, bbox_t_container &container);
 extern "C" LIB_API int detect_mat(const uint8_t* data, const size_t data_length, bbox_t_container &container);
 extern "C" LIB_API int dispose();
@@ -76,11 +76,12 @@ class Detector {
     float nms = .4;
     bool wait_stream;
 
-    LIB_API Detector(std::string cfg_filename, std::string weight_filename, int gpu_id = 0);
+    LIB_API Detector(std::string cfg_filename, std::string weight_filename, int gpu_id = 0, int batch_size = 1);
     LIB_API ~Detector();
 
     LIB_API std::vector<bbox_t> detect(std::string image_filename, float thresh = 0.2, bool use_mean = false);
     LIB_API std::vector<bbox_t> detect(image_t img, float thresh = 0.2, bool use_mean = false);
+    LIB_API std::vector<std::vector<bbox_t>> detectBatch(image_t img, int batch_size, int width, int height, float thresh, bool make_nms = true);
     static LIB_API image_t load_image(std::string image_filename);
     static LIB_API void free_image(image_t m);
     LIB_API int get_net_width() const;