
--------------------------------------------------------------------------------
-- YOLO onnx reasoning module
--------------------------------------------------------------------------------

---@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 abscissa
---@field y number vertical coordinate

---@class yolo_box*
---@field x number abscissa
---@field y number vertical coordinate
---@field width number Width
---@field height number height

---@class yolo_keypoint*
---@field x number abscissa
---@field y number vertical coordinate
---@field score? number Confidence
---@field visible? boolean visible or not

---@class yolo_detection*
---@field id? integer Category ID
---@field class_id? integer Category ID
---@field label? string Category name
---@field class_name? string Category name
---@field score number Confidence
---@field confidence? number Confidence
---@field x number abscissa
---@field y number vertical coordinate
---@field width number Width
---@field height number height
---@field box? yolo_box* detection frame
---@field polygon? yolo_point*[] OBB polygon
---@field mask? ImageObject* segmentation mask
---@field keypoints? yolo_keypoint*[] Posture key points
---@field track_id? integer Tracking ID

---@class yolo_classification*
---@field id? integer Category ID
---@field class_id? integer Category ID
---@field label? string Category name
---@field class_name? string Category name
---@field score number Confidence

---@class yolo_detect_options*
---@field conf? number confidence threshold
---@field confidence? number confidence threshold
---@field iou? number NMS IoU threshold
---@field class_names? string[] Category name
---@field classes? integer[] Only keep specified categories
---@field topk? integer Classification return quantity
---@field resize_mode? 'letterbox'|'stretch' Zoom method
---@field letterbox_mode? 'center'|'top_left' letterbox Alignment
---@field pad_color? integer fill color
---@field task? yolo_task Task type

---@class yolo_open_options: yolo_detect_options*
---@field model_path? string model path
---@field compiled_model_path? string CoreML Compile model path
---@field model_dir? string Model catalog
---@field class_names? string[] Category name
---@field input_width? integer Enter width
---@field input_height? integer Enter height
---@field providers? string[] ONNX Runtime provider list
---@field fallback_to_cpu? boolean provider Fallback CPU on failure
---@field threads? integer Number of threads
---@field coreml_compute_units? 'all'|'cpu_only'|'cpu_and_gpu'|'cpu_and_neural_engine' CoreML computing unit
---@field coreml_enable_on_subgraph? boolean Whether to enable CoreML in subgraphs
---@field coreml_require_static_input_shapes? boolean Whether to require static input shape
---@field coreml_create_mlprogram? boolean Whether to create MLProgram

---@class yolo_model_info*
---@field model_path? string model path
---@field used_coreml? boolean Whether to use CoreML EP
---@field task? yolo_task Task type
---@field input_name? string Enter name
---@field input_shape? integer[] Enter shape
---@field input_width? integer Enter width
---@field input_height? integer Enter height
---@field input_names? string[] Enter name
---@field output_names? string[] output name
---@field providers? string[] provider list
---@field outputs? { name: string, shape: integer[] }[] Output information

---@class yolo_tracker_options*
---@field iou_threshold? number Associated IoU thresholds
---@field max_age? integer Maximum number of reserved frames
---@field min_hits? integer Minimum number of hits

---@class yolo_tracker*
---@field closed? boolean Is it closed
---@field config? yolo_tracker_options* Trace configuration
---@field tracks? onnxruntime_track[] tracking status
---@field next_id? integer Next tracking ID
local yolo_tracker = {}

---@param detections yolo_detection*[] Test results
---@param timestamp? integer Timestamp
---@return yolo_detection*[] detections Detection results with track_id
function yolo_tracker:update(detections, timestamp) end

---@return boolean ok Is it successful?
function yolo_tracker:reset() end

---@return onnxruntime_track[] tracks tracking status
function yolo_tracker:state() end

---@return boolean ok Is it successful?
function yolo_tracker:close() end

---@class yolo_onnx_instance*
---@field closed boolean Is it closed
---@field session onnxruntime_session* ONNX Runtime session
---@field decoder onnxruntime_decoder* ONNX Runtime decoder
---@field decoder_schema table decode schema
---@field output_names string[] output name
---@field model_info yolo_model_info* Model information
---@field input_width integer Enter width
---@field input_height integer Enter height
---@field task yolo_task Task type
---@field profile? table profile Run configuration
---@field confidence number Default confidence threshold
---@field iou number Default NMS IoU threshold
---@field max_det? integer Maximum number of tests
---@field tracker? yolo_tracker*|onnxruntime_tracker* tracker
---@field class_names string[] Category name
local yolo_instance = {}

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

--- Open the YOLO model
---@param opts? yolo_open_options* Options
---@return yolo_onnx_instance* model Model instance
function M.new(opts) end

--- Open the YOLO model
---@param opts? yolo_open_options* Options
---@return yolo_onnx_instance* model Model instance
function M.create(opts) end

--- Open the YOLO model
---@param opts? yolo_open_options* Options
---@return yolo_onnx_instance* model Model instance
function M.open(opts) end

--- Return to supported task list
---@return yolo_task[] tasks task list
function M.supported_tasks() end

--- Create tracker
---@param opts? yolo_tracker_options* Options
---@return yolo_tracker* tracker tracker
function M.tracker(opts) end

--- Run the model and return the backend raw output or decoded results by task
---@param input yolo_image_input Enter image path, image data, or image object
---@param arg1? any Parameter 1
---@param arg2? any Parameter 2
---@param arg3? any Parameter 3
---@return any result Inference results
function yolo_instance:run(input, arg1, arg2, arg3) end

--- Run target detection
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections Test results
function yolo_instance:detect(input, conf, iou, opts) end

--- Run target detection
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections Test results
function yolo_instance:predict(input, conf, iou, opts) end

--- Run classification
---@param input yolo_image_input Enter image path, image data, or image object
---@param opts? yolo_detect_options* Options
---@return yolo_classification*[] classes Classification results
function yolo_instance:classify(input, opts) end

--- Run instance split
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections Segmentation results
function yolo_instance:segment(input, conf, iou, opts) end

--- Running attitude estimation
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections Attitude results
function yolo_instance:pose(input, conf, iou, opts) end

--- Run target tracking
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections Track results
function yolo_instance:track(input, conf, iou, opts) end

--- Run rotating box detection
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections OBB Test results
function yolo_instance:detect_obb(input, conf, iou, opts) end

--- Run rotating box detection
---@param input yolo_image_input Enter image path, image data, or image object
---@param conf? number confidence threshold
---@param iou? number NMS IoU threshold
---@param opts? yolo_detect_options* Options
---@return yolo_detection*[] detections OBB Test results
function yolo_instance:obb(input, conf, iou, opts) end

--- Get model information
---@return yolo_model_info* info Model information
function yolo_instance:get_model_info() end

--- Close model instance
function yolo_instance:close() end

return M
