
--------------------------------------------------------------------------------
-- YOLO onnx 推理模块
--------------------------------------------------------------------------------

---@type onnxruntime*
local _onnxruntime = require("onnxruntime")

---@type image_cv*
local _image_cv = require("image.cv")

---@type yolo_tracker_module*
local _yolo_tracker = require("yolo_tracker")

---@alias yolo_task 'detect'|'segment'|'pose'|'classify'|'obb'|'track'
---@alias yolo_image_input string|ImageObject*

---@class yolo_point*
---@field x number 横坐标
---@field y number 纵坐标

---@class yolo_box*
---@field x number 横坐标
---@field y number 纵坐标
---@field width number 宽度
---@field height number 高度

---@class yolo_keypoint*
---@field x number 横坐标
---@field y number 纵坐标
---@field score? number 置信度
---@field visible? boolean 是否可见

---@class yolo_detection*
---@field id? integer 类别 ID
---@field class_id? integer 类别 ID
---@field label? string 类别名称
---@field class_name? string 类别名称
---@field score number 置信度
---@field confidence? number 置信度
---@field x number 横坐标
---@field y number 纵坐标
---@field width number 宽度
---@field height number 高度
---@field box? yolo_box* 检测框
---@field polygon? yolo_point*[] OBB 多边形
---@field mask? ImageObject* 分割掩膜
---@field keypoints? yolo_keypoint*[] 姿态关键点
---@field track_id? integer 跟踪 ID

---@class yolo_classification*
---@field id? integer 类别 ID
---@field class_id? integer 类别 ID
---@field label? string 类别名称
---@field class_name? string 类别名称
---@field score number 置信度

---@class yolo_detect_options*
---@field conf? number 置信度阈值
---@field confidence? number 置信度阈值
---@field iou? number NMS IoU 阈值
---@field class_names? string[] 类别名称
---@field classes? integer[] 只保留指定类别
---@field topk? integer 分类返回数量
---@field resize_mode? 'letterbox'|'stretch' 缩放方式
---@field letterbox_mode? 'center'|'top_left' letterbox 对齐方式
---@field pad_color? integer 填充颜色
---@field task? yolo_task 任务类型

---@class yolo_open_options: yolo_detect_options*
---@field model_path? string 模型路径
---@field compiled_model_path? string CoreML 编译模型路径
---@field model_dir? string 模型目录
---@field class_names? string[] 类别名称
---@field input_width? integer 输入宽度
---@field input_height? integer 输入高度
---@field providers? string[] ONNX Runtime provider 列表
---@field fallback_to_cpu? boolean provider 失败时回退 CPU
---@field threads? integer 线程数
---@field coreml_compute_units? 'all'|'cpu_only'|'cpu_and_gpu'|'cpu_and_neural_engine' CoreML 计算单元
---@field coreml_enable_on_subgraph? boolean 是否在子图启用 CoreML
---@field coreml_require_static_input_shapes? boolean 是否要求静态输入形状
---@field coreml_create_mlprogram? boolean 是否创建 MLProgram

---@class yolo_model_info*
---@field model_path? string 模型路径
---@field used_coreml? boolean 是否使用 CoreML EP
---@field task? yolo_task 任务类型
---@field input_name? string 输入名
---@field input_shape? integer[] 输入形状
---@field input_width? integer 输入宽度
---@field input_height? integer 输入高度
---@field input_names? string[] 输入名
---@field output_names? string[] 输出名
---@field providers? string[] provider 列表
---@field outputs? { name: string, shape: integer[] }[] 输出信息

---@class yolo_tracker_options*
---@field iou_threshold? number 关联 IoU 阈值
---@field max_age? integer 最大保留帧数
---@field min_hits? integer 最小命中次数

---@class yolo_tracker*
---@field closed? boolean 是否已关闭
---@field config? yolo_tracker_options* 跟踪配置
---@field tracks? onnxruntime_track[] 跟踪状态
---@field next_id? integer 下一个跟踪 ID
local yolo_tracker = {}

---@param detections yolo_detection*[] 检测结果
---@param timestamp? integer 时间戳
---@return yolo_detection*[] detections 带 track_id 的检测结果
function yolo_tracker:update(detections, timestamp) end

---@return boolean ok 是否成功
function yolo_tracker:reset() end

---@return onnxruntime_track[] tracks 跟踪状态
function yolo_tracker:state() end

---@return boolean ok 是否成功
function yolo_tracker:close() end

---@class yolo_onnx_instance*
---@field closed boolean 是否已关闭
---@field session onnxruntime_session* ONNX Runtime 会话
---@field decoder onnxruntime_decoder* ONNX Runtime 解码器
---@field decoder_schema table 解码 schema
---@field output_names string[] 输出名
---@field model_info yolo_model_info* 模型信息
---@field input_width integer 输入宽度
---@field input_height integer 输入高度
---@field task yolo_task 任务类型
---@field profile? table profile 运行配置
---@field confidence number 默认置信度阈值
---@field iou number 默认 NMS IoU 阈值
---@field max_det? integer 最大检测数
---@field tracker? yolo_tracker*|onnxruntime_tracker* 跟踪器
---@field class_names string[] 类别名称
local yolo_instance = {}

---@class yolo_onnx*
---@type yolo_onnx*
local M = {}

--- 打开 YOLO 模型
---@param opts? yolo_open_options* 选项
---@return yolo_onnx_instance* model 模型实例
function M.new(opts) end

--- 打开 YOLO 模型
---@param opts? yolo_open_options* 选项
---@return yolo_onnx_instance* model 模型实例
function M.create(opts) end

--- 打开 YOLO 模型
---@param opts? yolo_open_options* 选项
---@return yolo_onnx_instance* model 模型实例
function M.open(opts) end

--- 返回支持的任务列表
---@return yolo_task[] tasks 任务列表
function M.supported_tasks() end

--- 创建跟踪器
---@param opts? yolo_tracker_options* 选项
---@return yolo_tracker* tracker 跟踪器
function M.tracker(opts) end

--- 运行模型，返回后端原始输出或按任务解码后的结果
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param arg1? any 参数 1
---@param arg2? any 参数 2
---@param arg3? any 参数 3
---@return any result 推理结果
function yolo_instance:run(input, arg1, arg2, arg3) end

--- 运行目标检测
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections 检测结果
function yolo_instance:detect(input, conf, iou, opts) end

--- 运行目标检测
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections 检测结果
function yolo_instance:predict(input, conf, iou, opts) end

--- 运行分类
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param opts? yolo_detect_options* 选项
---@return yolo_classification*[] classes 分类结果
function yolo_instance:classify(input, opts) end

--- 运行实例分割
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections 分割结果
function yolo_instance:segment(input, conf, iou, opts) end

--- 运行姿态估计
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections 姿态结果
function yolo_instance:pose(input, conf, iou, opts) end

--- 运行目标跟踪
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections 跟踪结果
function yolo_instance:track(input, conf, iou, opts) end

--- 运行旋转框检测
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections OBB 检测结果
function yolo_instance:detect_obb(input, conf, iou, opts) end

--- 运行旋转框检测
---@param input yolo_image_input 输入图片路径、图片数据或图片对象
---@param conf? number 置信度阈值
---@param iou? number NMS IoU 阈值
---@param opts? yolo_detect_options* 选项
---@return yolo_detection*[] detections OBB 检测结果
function yolo_instance:obb(input, conf, iou, opts) end

--- 获取模型信息
---@return yolo_model_info* info 模型信息
function yolo_instance:get_model_info() end

--- 关闭模型实例
function yolo_instance:close() end

return M
