// Copyright chenyy 2026. Licensed under the MIT License.
// TensorRT inference engine implementation

#include "armor_yolo_detect/yolo_tensorrt.hpp"
#include "armor_yolo_detect/trt_logger.hpp"

#include <cuda_runtime.h>
#include <NvInferPlugin.h>
#include <NvInferVersion.h>
#include <NvOnnxParser.h>

#include <algorithm>
#include <array>
#include <cstring>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <numeric>
#include <vector>

#include <opencv2/imgproc.hpp>

namespace {
inline uint16_t floatToHalf(float f) {
  uint32_t x;
  std::memcpy(&x, &f, sizeof(float));
  uint32_t sign = (x >> 16) & 0x8000;
  uint32_t exp = (x & 0x7f800000) >> 23;
  uint32_t mant = x & 0x007fffff;
  if (exp == 0) return sign;
  if (exp == 255) return sign | 0x7c00 | (mant ? 1 : 0);
  int new_exp = static_cast<int>(exp) - 127 + 15;
  if (new_exp <= 0) return sign;
  if (new_exp >= 31) return sign | 0x7c00;
  return sign | (static_cast<uint32_t>(new_exp) << 10) | (mant >> 13);
}
} // namespace

namespace armor_yolo_detect {

YoloTensorRT::~YoloTensorRT() {
  if (cuda_stream_ != nullptr) {
    cudaStreamDestroy(cuda_stream_);
    cuda_stream_ = nullptr;
  }

  if (device_buffer_[0]) {
    cudaFree(device_buffer_[0]);
    device_buffer_[0] = nullptr;
  }
  if (device_buffer_[1]) {
    cudaFree(device_buffer_[1]);
    device_buffer_[1] = nullptr;
  }

  if (context_) {
#if NV_TENSORRT_MAJOR >= 10
    delete context_;
#else
    context_->destroy();
#endif
    context_ = nullptr;
  }
  if (engine_) {
#if NV_TENSORRT_MAJOR >= 10
    delete engine_;
#else
    engine_->destroy();
#endif
    engine_ = nullptr;
  }
  if (runtime_) {
#if NV_TENSORRT_MAJOR >= 10
    delete runtime_;
#else
    runtime_->destroy();
#endif
    runtime_ = nullptr;
  }
}

bool YoloTensorRT::initFromEngine(const std::string& engine_path) {
  std::ifstream engine_file(engine_path, std::ios::binary);
  if (!engine_file.good()) {
    std::cerr << "Error: Failed to open engine file: " << engine_path << std::endl;
    return false;
  }

  engine_file.seekg(0, std::ifstream::end);
  const size_t engine_size = static_cast<size_t>(engine_file.tellg());
  engine_file.seekg(0, std::ifstream::beg);

  std::vector<char> engine_data(engine_size);
  engine_file.read(engine_data.data(), static_cast<std::streamsize>(engine_size));
  engine_file.close();

  if (cudaSetDevice(0) != cudaSuccess) {
    std::cerr << "Error: cudaSetDevice(0) failed" << std::endl;
    return false;
  }

  if (cuda_stream_ == nullptr && cudaStreamCreate(&cuda_stream_) != cudaSuccess) {
    std::cerr << "Error: cudaStreamCreate failed" << std::endl;
    return false;
  }

  runtime_ = nvinfer1::createInferRuntime(gLogger);
  if (!runtime_) {
    std::cerr << "Error: Failed to create TensorRT runtime" << std::endl;
    return false;
  }

  engine_ = runtime_->deserializeCudaEngine(engine_data.data(), engine_size);
  if (!engine_) {
    std::cerr << "Error: Failed to deserialize TensorRT engine" << std::endl;
    return false;
  }

  context_ = engine_->createExecutionContext();
  if (!context_) {
    std::cerr << "Error: Failed to create TensorRT execution context" << std::endl;
    return false;
  }

#if NV_TENSORRT_MAJOR >= 10
  io_tensor_count_ = static_cast<int>(engine_->getNbIOTensors());
  for (int i = 0; i < io_tensor_count_; ++i) {
    const char* tensor_name = engine_->getIOTensorName(i);
    const auto mode = engine_->getTensorIOMode(tensor_name);
    if (mode == nvinfer1::TensorIOMode::kINPUT) {
      input_index_ = i;
      input_tensor_name_ = tensor_name;
      if (engine_->getTensorDataType(tensor_name) == nvinfer1::DataType::kHALF) {
        is_input_half_ = true;
      }
    } else if (mode == nvinfer1::TensorIOMode::kOUTPUT) {
      output_index_ = i;
      output_tensor_name_ = tensor_name;
    }
  }
#else
  input_index_ = engine_->getBindingIndex("images");
  output_index_ = 1 - input_index_;
  io_tensor_count_ = 2;
  if (engine_->getBindingDataType(input_index_) == nvinfer1::DataType::kHALF) {
    is_input_half_ = true;
  }
#endif

#if NV_TENSORRT_MAJOR >= 10
  if (input_tensor_name_.empty() || output_tensor_name_.empty()) {
    std::cerr << "Error: Failed to resolve TensorRT I/O tensor names" << std::endl;
    return false;
  }
  nvinfer1::Dims output_dims = engine_->getTensorShape(output_tensor_name_.c_str());
#else
  nvinfer1::Dims output_dims = engine_->getBindingDimensions(output_index_);
#endif

  if (output_dims.nbDims >= 2) {
    output_rows_ = output_dims.d[output_dims.nbDims - 2];
    output_cols_ = output_dims.d[output_dims.nbDims - 1];
  }
  if (output_rows_ <= 0 || output_cols_ <= 0) {
    output_rows_ = OUTPUT_ROWS;
    output_cols_ = OUTPUT_COLS;
  }

  const size_t input_size =
      static_cast<size_t>(INPUT_WIDTH) * static_cast<size_t>(INPUT_HEIGHT) *
      static_cast<size_t>(INPUT_C) * sizeof(float);
  const size_t output_size =
      static_cast<size_t>(output_rows_) * static_cast<size_t>(output_cols_) * sizeof(float);

  if (cudaMalloc(&device_buffer_[input_index_], input_size) != cudaSuccess) {
    std::cerr << "Error: cudaMalloc failed for input buffer" << std::endl;
    return false;
  }
  if (cudaMalloc(&device_buffer_[output_index_], output_size) != cudaSuccess) {
    std::cerr << "Error: cudaMalloc failed for output buffer" << std::endl;
    return false;
  }

  input_buffer_.resize(static_cast<size_t>(INPUT_WIDTH) * static_cast<size_t>(INPUT_HEIGHT) *
                       static_cast<size_t>(INPUT_C));
  output_buffer_.resize(static_cast<size_t>(output_rows_) * static_cast<size_t>(output_cols_));

  return true;
}

bool YoloTensorRT::buildFromOnnx(
    const std::string& onnx_path,
    const std::string& engine_path,
    int max_batch_size) {
  (void)max_batch_size;

  initLibNvInferPlugins(&gLogger, "");

  auto builder = nvinfer1::createInferBuilder(gLogger);
  if (!builder) {
    std::cerr << "Error: Failed to create TensorRT builder" << std::endl;
    return false;
  }

#if NV_TENSORRT_MAJOR >= 10
  auto network = builder->createNetworkV2(0U);
#else
  const auto explicit_batch =
      1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
  auto network = builder->createNetworkV2(explicit_batch);
#endif
  if (!network) {
    std::cerr << "Error: Failed to create TensorRT network" << std::endl;
    return false;
  }

  auto parser = nvonnxparser::createParser(*network, gLogger);
  if (!parser) {
    std::cerr << "Error: Failed to create ONNX parser" << std::endl;
    return false;
  }

  if (!parser->parseFromFile(onnx_path.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kWARNING))) {
    std::cerr << "Error: Failed to parse ONNX model" << std::endl;
    return false;
  }

  auto config = builder->createBuilderConfig();
#if NV_TENSORRT_MAJOR >= 10
  config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1ULL << 30);
#else
  builder->setMaxBatchSize(max_batch_size);
  builder->setMaxWorkspaceSize(1ULL << 30);
#endif

  std::cout << "Building TensorRT engine from ONNX..." << std::endl;

#if NV_TENSORRT_MAJOR >= 10
  auto serialized_engine = builder->buildSerializedNetwork(*network, *config);
  if (!serialized_engine) {
    std::cerr << "Error: Failed to build serialized TensorRT engine" << std::endl;
    return false;
  }

  runtime_ = nvinfer1::createInferRuntime(gLogger);
  if (!runtime_) {
    std::cerr << "Error: Failed to create TensorRT runtime" << std::endl;
    return false;
  }

  engine_ = runtime_->deserializeCudaEngine(serialized_engine->data(), serialized_engine->size());
#else
  engine_ = builder->buildEngineWithConfig(*network, *config);
  auto serialized_engine = engine_ ? engine_->serialize() : nullptr;
#endif

  if (!engine_) {
    std::cerr << "Error: Failed to build TensorRT engine" << std::endl;
    return false;
  }

  std::ofstream engine_file(engine_path, std::ios::binary);
  engine_file.write(static_cast<const char*>(serialized_engine->data()),
                    static_cast<std::streamsize>(serialized_engine->size()));
  engine_file.close();

  context_ = engine_->createExecutionContext();
  if (!context_) {
    std::cerr << "Error: Failed to create TensorRT execution context" << std::endl;
    return false;
  }

#if NV_TENSORRT_MAJOR >= 10
  io_tensor_count_ = static_cast<int>(engine_->getNbIOTensors());
  for (int i = 0; i < io_tensor_count_; ++i) {
    const char* tensor_name = engine_->getIOTensorName(i);
    const auto mode = engine_->getTensorIOMode(tensor_name);
    if (mode == nvinfer1::TensorIOMode::kINPUT) {
      input_index_ = i;
      input_tensor_name_ = tensor_name;
      if (engine_->getTensorDataType(tensor_name) == nvinfer1::DataType::kHALF) {
        is_input_half_ = true;
      }
    } else if (mode == nvinfer1::TensorIOMode::kOUTPUT) {
      output_index_ = i;
      output_tensor_name_ = tensor_name;
    }
  }
  nvinfer1::Dims output_dims = engine_->getTensorShape(output_tensor_name_.c_str());
#else
  input_index_ = engine_->getBindingIndex("images");
  output_index_ = 1 - input_index_;
  io_tensor_count_ = 2;
  if (engine_->getBindingDataType(input_index_) == nvinfer1::DataType::kHALF) {
    is_input_half_ = true;
  }
  nvinfer1::Dims output_dims = engine_->getBindingDimensions(output_index_);
#endif

  if (output_dims.nbDims >= 2) {
    output_rows_ = output_dims.d[output_dims.nbDims - 2];
    output_cols_ = output_dims.d[output_dims.nbDims - 1];
  }
  if (output_rows_ <= 0 || output_cols_ <= 0) {
    output_rows_ = OUTPUT_ROWS;
    output_cols_ = OUTPUT_COLS;
  }

  const size_t input_size =
      static_cast<size_t>(INPUT_WIDTH) * static_cast<size_t>(INPUT_HEIGHT) *
      static_cast<size_t>(INPUT_C) * sizeof(float);
  const size_t output_size =
      static_cast<size_t>(output_rows_) * static_cast<size_t>(output_cols_) * sizeof(float);

  if (cuda_stream_ == nullptr && cudaStreamCreate(&cuda_stream_) != cudaSuccess) {
    std::cerr << "Error: cudaStreamCreate failed" << std::endl;
    return false;
  }

  if (cudaMalloc(&device_buffer_[input_index_], input_size) != cudaSuccess) {
    std::cerr << "Error: cudaMalloc failed for input buffer" << std::endl;
    return false;
  }
  if (cudaMalloc(&device_buffer_[output_index_], output_size) != cudaSuccess) {
    std::cerr << "Error: cudaMalloc failed for output buffer" << std::endl;
    return false;
  }

  input_buffer_.resize(static_cast<size_t>(INPUT_WIDTH) * static_cast<size_t>(INPUT_HEIGHT) *
                       static_cast<size_t>(INPUT_C));
  output_buffer_.resize(static_cast<size_t>(output_rows_) * static_cast<size_t>(output_cols_));

#if NV_TENSORRT_MAJOR >= 10
  delete config;
  delete parser;
  delete network;
  delete builder;
#else
  config->destroy();
  parser->destroy();
  network->destroy();
  builder->destroy();
#endif

  return true;
}

cv::Mat YoloTensorRT::preprocess(const cv::Mat& img) {
  cv::Mat resized;
  cv::resize(img, resized, cv::Size(INPUT_WIDTH, INPUT_HEIGHT));

  cv::Mat rgb;
  if (resized.channels() == 3) {
    // DO NOT convert to RGB. The model natively expects BGR!
    // cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB);
    rgb = resized;
  } else {
    cv::cvtColor(resized, rgb, cv::COLOR_GRAY2BGR);
  }

  cv::Mat normalized;
  rgb.convertTo(normalized, CV_32FC3, 1.0 / 255.0);
  return normalized;
}

bool YoloTensorRT::inferInternal(const float* input, float* output, int batch_size) {
  if (cuda_stream_ == nullptr) {
    std::cerr << "Error: CUDA stream is not initialized" << std::endl;
    return false;
  }

  const size_t num_input_elements =
      static_cast<size_t>(INPUT_WIDTH) * static_cast<size_t>(INPUT_HEIGHT) *
      static_cast<size_t>(INPUT_C) * static_cast<size_t>(batch_size);

  const void* copy_src = input;
  size_t input_bytes = num_input_elements * sizeof(float);

  if (is_input_half_) {
    if (half_input_buffer_.size() < num_input_elements) {
      half_input_buffer_.resize(num_input_elements);
    }
    for (size_t i = 0; i < num_input_elements; ++i) {
      half_input_buffer_[i] = floatToHalf(input[i]);
    }
    copy_src = half_input_buffer_.data();
    input_bytes = num_input_elements * sizeof(uint16_t);
  }

  if (cudaMemcpyAsync(
          device_buffer_[input_index_], copy_src, input_bytes, cudaMemcpyHostToDevice, cuda_stream_) !=
      cudaSuccess) {
    std::cerr << "Error: cudaMemcpyAsync HostToDevice failed" << std::endl;
    return false;
  }

  size_t output_bytes =
      static_cast<size_t>(output_rows_) * static_cast<size_t>(output_cols_) * sizeof(float) *
      static_cast<size_t>(batch_size);

#if NV_TENSORRT_MAJOR >= 10
  nvinfer1::Dims input_dims = context_->getTensorShape(input_tensor_name_.c_str());
  bool has_dynamic_dim = false;
  for (int i = 0; i < input_dims.nbDims; ++i) {
    if (input_dims.d[i] == -1) {
      has_dynamic_dim = true;
      break;
    }
  }

  if (has_dynamic_dim) {
    nvinfer1::Dims4 fixed_input_dims(batch_size, INPUT_C, INPUT_HEIGHT, INPUT_WIDTH);
    if (!context_->setInputShape(input_tensor_name_.c_str(), fixed_input_dims)) {
      std::cerr << "Error: setInputShape failed for input tensor: " << input_tensor_name_ << std::endl;
      return false;
    }
  }

  nvinfer1::Dims output_dims = context_->getTensorShape(output_tensor_name_.c_str());
  if (output_dims.nbDims >= 2) {
    const int resolved_rows = output_dims.d[output_dims.nbDims - 2];
    const int resolved_cols = output_dims.d[output_dims.nbDims - 1];
    if (resolved_rows > 0 && resolved_cols > 0) {
      output_rows_ = resolved_rows;
      output_cols_ = resolved_cols;
    }
  }

  const size_t required_output_elems =
      static_cast<size_t>(output_rows_) * static_cast<size_t>(output_cols_);
  if (required_output_elems > output_buffer_.size()) {
    output_buffer_.resize(required_output_elems);
  }
  output_bytes = required_output_elems * sizeof(float) * static_cast<size_t>(batch_size);

  if (!context_->setTensorAddress(input_tensor_name_.c_str(), device_buffer_[input_index_])) {
    std::cerr << "Error: setTensorAddress failed for input tensor" << std::endl;
    return false;
  }
  if (!context_->setTensorAddress(output_tensor_name_.c_str(), device_buffer_[output_index_])) {
    std::cerr << "Error: setTensorAddress failed for output tensor" << std::endl;
    return false;
  }
  if (!context_->enqueueV3(cuda_stream_)) {
    std::cerr << "Error: TensorRT enqueueV3 failed" << std::endl;
    return false;
  }
#else
  if (!context_->execute(batch_size, device_buffer_)) {
    std::cerr << "Error: TensorRT execute failed" << std::endl;
    return false;
  }
#endif

  if (cudaMemcpyAsync(
          output, device_buffer_[output_index_], output_bytes, cudaMemcpyDeviceToHost, cuda_stream_) !=
      cudaSuccess) {
    std::cerr << "Error: cudaMemcpyAsync DeviceToHost failed" << std::endl;
    return false;
  }

  if (cudaStreamSynchronize(cuda_stream_) != cudaSuccess) {
    std::cerr << "Error: cudaStreamSynchronize failed" << std::endl;
    return false;
  }

  return true;
}

std::vector<YoloObject> YoloTensorRT::postprocess(const float* output, int detect_color, bool dump_this_frame) {
  std::vector<YoloObject> objects;

  int dropped_by_conf = 0;
  int dropped_by_color = 0;

  struct DumpCandidate {
    int row = -1;
    float conf = 0.0f;
    float color_score = 0.0f;
    int color_id = -1;
    int num_id = -1;
    float score = 0.0f;
  };

  std::vector<DumpCandidate> dump_candidates;
  if (dump_this_frame) {
    dump_candidates.reserve(static_cast<size_t>(dump_topk_));
  }

  for (int i = 0; i < output_rows_; i++) {
    const float* row = output + i * output_cols_;

    const float conf = sigmoid(row[8]);
    if (conf < conf_threshold_) {
      dropped_by_conf++;
      continue;
    }

    int color_id = 0;
    float max_color_score = sigmoid(row[9]);
    for (int c = 1; c < 4; c++) {
      const float color_score = sigmoid(row[9 + c]);
      if (color_score > max_color_score) {
        max_color_score = color_score;
        color_id = c;
      }
    }

    if (color_id == 2 || color_id == 3 ||
        (detect_color >= 0 && detect_color == 0 && color_id == 1) ||
        (detect_color >= 0 && detect_color == 1 && color_id == 0)) {
      dropped_by_color++;
      continue;
    }

    const int num_class_count = std::max(0, output_cols_ - 13);
    if (num_class_count <= 0) {
      continue;
    }

    int num_id = 0;
    float max_num_score = sigmoid(row[13]);
    for (int c = 1; c < num_class_count; c++) {
      const float num_score = sigmoid(row[13 + c]);
      if (num_score > max_num_score) {
        max_num_score = num_score;
        num_id = c;
      }
    }

    YoloObject obj;
    obj.prob = conf; // Only use detection config, not color conf
    obj.label = num_id;
    obj.color = static_cast<ModelColor>(color_id);

    float min_x = std::numeric_limits<float>::max();
    float max_x = std::numeric_limits<float>::lowest();
    float min_y = std::numeric_limits<float>::max();
    float max_y = std::numeric_limits<float>::lowest();

    for (int k = 0; k < 4; ++k) {
      const float x = row[k * 2];
      const float y = row[k * 2 + 1];
      obj.landmarks[k * 2] = x;
      obj.landmarks[k * 2 + 1] = y;
      min_x = std::min(min_x, x);
      max_x = std::max(max_x, x);
      min_y = std::min(min_y, y);
      max_y = std::max(max_y, y);
    }

    obj.rect = cv::Rect_<float>(min_x, min_y, max_x - min_x, max_y - min_y);

    objects.push_back(obj);

    if (dump_this_frame) {
      dump_candidates.push_back(DumpCandidate{i, conf, max_color_score, color_id, num_id, obj.prob});
      std::sort(dump_candidates.begin(), dump_candidates.end(),
                [](const DumpCandidate& a, const DumpCandidate& b) {
                  return a.score > b.score;
                });
      if (static_cast<int>(dump_candidates.size()) > dump_topk_) {
        dump_candidates.resize(static_cast<size_t>(dump_topk_));
      }
    }
  }

  std::vector<YoloObject> result = objects;

  if (dump_this_frame) {
    std::cerr << std::fixed << std::setprecision(4)
              << "[YoloTensorRT][dump] frame=" << infer_counter_
              << " output_shape=[" << output_rows_ << "," << output_cols_ << "]"
              << " topk=" << dump_candidates.size()
              << " dropped_conf=" << dropped_by_conf
              << " dropped_color=" << dropped_by_color
              << " final_after_postprocess=" << result.size()
              << std::endl;
  }

  return result;
}

std::vector<int> YoloTensorRT::applyNMS(
    const std::vector<cv::Rect>& boxes,
    const std::vector<float>& scores,
    float nms_threshold) {
  std::vector<int> indices;
  if (boxes.empty()) {
    return indices;
  }

  std::vector<int> order(boxes.size());
  std::iota(order.begin(), order.end(), 0);
  std::sort(order.begin(), order.end(),
            [&scores](int i, int j) { return scores[i] > scores[j]; });

  std::vector<bool> suppressed(boxes.size(), false);

  for (size_t i = 0; i < order.size(); i++) {
    const int idx = order[i];
    if (suppressed[idx]) {
      continue;
    }

    indices.push_back(idx);

    for (size_t j = i + 1; j < order.size(); j++) {
      const int idx_j = order[j];
      if (suppressed[idx_j]) {
        continue;
      }

      const cv::Rect inter = boxes[idx] & boxes[idx_j];
      const float inter_area = static_cast<float>(inter.area());
      const float union_area = static_cast<float>(boxes[idx].area() + boxes[idx_j].area()) - inter_area;
      const float iou = (union_area > 0.0f) ? inter_area / union_area : 0.0f;

      if (iou > nms_threshold) {
        suppressed[idx_j] = true;
      }
    }
  }

  return indices;
}

std::vector<YoloObject> YoloTensorRT::infer(const cv::Mat& img, int detect_color) {
  if (img.empty()) {
    return {};
  }

  const cv::Mat preprocessed = preprocess(img);

  const int h = preprocessed.rows;
  const int w = preprocessed.cols;
  const int c = preprocessed.channels();

  const size_t channel_size = static_cast<size_t>(h) * static_cast<size_t>(w);
  const size_t required_input_elems = channel_size * static_cast<size_t>(c);
  if (input_buffer_.size() < required_input_elems) {
    input_buffer_.resize(required_input_elems);
  }

  std::array<cv::Mat, INPUT_C> chw_channels = {
      cv::Mat(h, w, CV_32F, input_buffer_.data() + channel_size * 0),
      cv::Mat(h, w, CV_32F, input_buffer_.data() + channel_size * 1),
      cv::Mat(h, w, CV_32F, input_buffer_.data() + channel_size * 2)};
  cv::split(preprocessed, chw_channels);

  if (!inferInternal(input_buffer_.data(), output_buffer_.data(), 1)) {
    return {};
  }

  infer_counter_++;
  const bool dump_this_frame =
      dump_model_output_ &&
      (dump_interval_ <= 1 || (infer_counter_ % static_cast<uint64_t>(dump_interval_) == 0));

  auto objects = postprocess(output_buffer_.data(), detect_color, dump_this_frame);

  const float scale_x = static_cast<float>(img.cols) / static_cast<float>(INPUT_WIDTH);
  const float scale_y = static_cast<float>(img.rows) / static_cast<float>(INPUT_HEIGHT);
  const float max_x = static_cast<float>(std::max(0, img.cols - 1));
  const float max_y = static_cast<float>(std::max(0, img.rows - 1));

  for (auto& object : objects) {
    for (int k = 0; k < 4; ++k) {
      const int x_idx = k * 2;
      const int y_idx = x_idx + 1;
      object.landmarks[x_idx] = std::clamp(object.landmarks[x_idx] * scale_x, 0.0f, max_x);
      object.landmarks[y_idx] = std::clamp(object.landmarks[y_idx] * scale_y, 0.0f, max_y);
    }

    const float x1 = std::min(std::min(object.landmarks[0], object.landmarks[2]),
                              std::min(object.landmarks[4], object.landmarks[6]));
    const float y1 = std::min(std::min(object.landmarks[1], object.landmarks[3]),
                              std::min(object.landmarks[5], object.landmarks[7]));
    const float x2 = std::max(std::max(object.landmarks[0], object.landmarks[2]),
                              std::max(object.landmarks[4], object.landmarks[6]));
    const float y2 = std::max(std::max(object.landmarks[1], object.landmarks[3]),
                              std::max(object.landmarks[5], object.landmarks[7]));

    object.rect.x = x1;
    object.rect.y = y1;
    object.rect.width = std::max(0.0f, x2 - x1);
    object.rect.height = std::max(0.0f, y2 - y1);
  }

  std::vector<cv::Rect> boxes;
  std::vector<float> scores;
  boxes.reserve(objects.size());
  scores.reserve(objects.size());
  for (const auto& obj : objects) {
    boxes.push_back(cv::Rect(static_cast<int>(obj.rect.x), static_cast<int>(obj.rect.y),
                             static_cast<int>(obj.rect.width), static_cast<int>(obj.rect.height)));
    scores.push_back(obj.prob);
  }

  std::vector<int> keep_indices = applyNMS(boxes, scores, nms_threshold_);

  std::vector<YoloObject> final_objects;
  final_objects.reserve(keep_indices.size());
  for (int idx : keep_indices) {
    final_objects.push_back(objects[idx]);
  }

  return final_objects;
}

}  // namespace armor_yolo_detect
