diff --git a/modules/yolo_config_parser.cpp b/modules/yolo_config_parser.cpp index 18b98c5..c9aaef4 100644 --- a/modules/yolo_config_parser.cpp +++ b/modules/yolo_config_parser.cpp @@ -48,7 +48,7 @@ DEFINE_string(config_file_path1, "not-specified", "[REQUIRED] Darknet cfg file") DEFINE_string(wts_file_path1, "not-specified", "[REQUIRED] Darknet weights file"); DEFINE_string(config_file_path2, "not-specified", "[REQUIRED] Darknet cfg file"); DEFINE_string(wts_file_path2, "not-specified", "[REQUIRED] Darknet weights file"); -DEFINE_string(labels_file_path, "configs/bdd100k.names", "[REQUIRED] Object class labels file"); +DEFINE_string(labels_file_path, "../configs/bdd100k.names", "[REQUIRED] Object class labels file"); DEFINE_string(precision, "kFLOAT", "[OPTIONAL] Inference precision. Choose from kFLOAT, kHALF and kINT8."); DEFINE_string(deviceType, "kGPU", @@ -94,6 +94,9 @@ DEFINE_bool(prof, false, DEFINE_string(dump, "not-specified", "[OPTIONAL] Path to dump predictions for mAP calculation"); +DEFINE_string(output, "not-specified", + "[OPTIONAL] Path to output predictions for pseudo-labeling"); + DEFINE_bool(mp, false, "[OPTIONAL] Flag to multi-precision"); DEFINE_int64(dla, -1, "[OPTIONAL] DLA"); @@ -269,3 +272,9 @@ get_cuda_flg(void) { return FLAGS_cuda; } + +std::string +get_output_path(void) +{ + return FLAGS_output; +} diff --git a/modules/yolo_config_parser.h b/modules/yolo_config_parser.h index a4b806c..e52bc2d 100644 --- a/modules/yolo_config_parser.h +++ b/modules/yolo_config_parser.h @@ -66,4 +66,6 @@ bool get_multi_precision_flg(void); bool get_cuda_flg(void); +std::string +get_output_path(void); #endif //_YOLO_CONFIG_PARSER_ diff --git a/samples/sample_detector.cpp b/samples/sample_detector.cpp index 5bc7d35..7bc1aaa 100644 --- a/samples/sample_detector.cpp +++ b/samples/sample_detector.cpp @@ -43,6 +43,34 @@ write_prediction(std::string dumpPath, std::string filename, std::vector names, std::vector objects, int width, int height) +{ + int pos = filename.find_last_of("."); + std::string body = filename.substr(0, pos); + std::string dstName = body + ".txt"; + std::ofstream writing_file; + fs::path p = outputPath; + fs::create_directory(p); + p.append(dstName); + writing_file.open(p.string(), std::ios::out); + for (const auto & object : objects) { + const auto left = object.rect.x; + const auto top = object.rect.y; + const auto right = std::clamp(left + object.rect.width, 0, width); + const auto bottom = std::clamp(top + object.rect.height, 0, height); + std::string writing_text = format("%d %f %f %f %f", + object.id, + (left+right)/2/(double)width, + (top+bottom)/2/(double)height, + object.rect.width/(double)width, + object.rect.height/(double)height + ); + writing_file << writing_text << std::endl; + } + writing_file.close(); +} + int main(int argc, char** argv) { @@ -58,6 +86,7 @@ int main(int argc, char** argv) std::string save_path = getSaveDetectionsPath(); bool dont_show = get_dont_show_flg(); const std::string dumpPath = get_dump_path(); + const std::string outputPath = get_output_path(); const bool cuda = get_cuda_flg(); Config config; config.net_type = YOLOV4; @@ -97,14 +126,22 @@ int main(int argc, char** argv) batch_img.push_back(src); } detector->detect(batch_img, batch_res, cuda); - detector->segment(batch_img); + if (!dont_show) { + detector->segment(batch_img); + } if (dumpPath != "not-specified") { fs::path p (file.path()); std::string filename = p.filename().string(); std::vector names = get_names(); write_prediction(dumpPath, filename, names, batch_res[0], src.cols, src.rows); - } + } + if (outputPath != "not-specified") { + fs::path p (file.path()); + std::string filename = p.filename().string(); + std::vector names = get_names(); + write_outputs(outputPath, filename, names, batch_res[0], src.cols, src.rows); + } //disp if (dont_show == true) { continue; @@ -139,7 +176,9 @@ int main(int argc, char** argv) std::vector batch_img; batch_img.push_back(frame); detector->detect(batch_img, batch_res, cuda); - detector->segment(batch_img); + if (!dont_show) { + detector->segment(batch_img); + } if (dont_show == true) { continue;