44 lines
1014 B
C++
44 lines
1014 B
C++
#include "yolo.hpp"
|
|
|
|
#include <yaml-cpp/yaml.h>
|
|
|
|
#include "yolos/yolo11.hpp"
|
|
#include "yolos/yolov5.hpp"
|
|
#include "yolos/yolov8.hpp"
|
|
|
|
namespace auto_aim
|
|
{
|
|
YOLO::YOLO(const std::string & config_path, bool debug)
|
|
{
|
|
auto yaml = YAML::LoadFile(config_path);
|
|
auto yolo_name = yaml["yolo_name"].as<std::string>();
|
|
|
|
if (yolo_name == "yolov8") {
|
|
yolo_ = std::make_unique<YOLOV8>(config_path, debug);
|
|
}
|
|
|
|
else if (yolo_name == "yolo11") {
|
|
yolo_ = std::make_unique<YOLO11>(config_path, debug);
|
|
}
|
|
|
|
else if (yolo_name == "yolov5") {
|
|
yolo_ = std::make_unique<YOLOV5>(config_path, debug);
|
|
}
|
|
|
|
else {
|
|
throw std::runtime_error("Unknown yolo name: " + yolo_name + "!");
|
|
}
|
|
}
|
|
|
|
std::list<Armor> YOLO::detect(const cv::Mat & img, int frame_count)
|
|
{
|
|
return yolo_->detect(img, frame_count);
|
|
}
|
|
|
|
std::list<Armor> YOLO::postprocess(
|
|
double scale, cv::Mat & output, const cv::Mat & bgr_img, int frame_count)
|
|
{
|
|
return yolo_->postprocess(scale, output, bgr_img, frame_count);
|
|
}
|
|
|
|
} // namespace auto_aim
|