ultralytics_inference/

task.rs

// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

//! Task definitions for YOLO models.
//!
//! This module defines the different tasks that YOLO models can perform,
//! along with their associated capabilities and string representations.

use std::fmt;
use std::str::FromStr;

/// YOLO model task types.
///
/// Each task type corresponds to a different computer vision problem
/// that YOLO models can solve. The task type determines the expected
/// model outputs and post-processing steps.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum Task {
    /// Object detection.
    /// Predicts bounding boxes and class labels for objects in an image.
    #[default]
    Detect,
    /// Instance segmentation.
    /// Predicts bounding boxes, class labels, and pixel-level masks for objects.
    Segment,
    /// Pose estimation.
    /// Predicts bounding boxes and skeletal keypoints for objects (e.g., humans).
    Pose,
    /// Image classification.
    /// Predicts class probabilities for the entire image (no localization).
    Classify,
    /// Oriented bounding box detection (OBB).
    /// Predicts rotated bounding boxes for objects, useful for aerial imagery etc.
    Obb,
    /// Semantic segmentation.
    /// Assigns a class label to every pixel in the image.
    Semantic,
}

impl Task {
    /// Get the string representation used in ONNX model metadata
    /// (e.g. `"detect"`, `"segment"`).
    #[must_use]
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::Detect => "detect",
            Self::Segment => "segment",
            Self::Pose => "pose",
            Self::Classify => "classify",
            Self::Obb => "obb",
            Self::Semantic => "semantic",
        }
    }

    /// ONNX filename suffix for this task, used to construct `{family}n{suffix}.onnx`
    /// (e.g. `yolo26n-seg.onnx`, `yolo11n-pose.onnx`, `yolov8n.onnx`).
    ///
    /// ```
    /// use ultralytics_inference::Task;
    /// assert_eq!(Task::Detect.model_suffix(), "");
    /// assert_eq!(Task::Segment.model_suffix(), "-seg");
    /// ```
    #[must_use]
    pub const fn model_suffix(&self) -> &'static str {
        match self {
            Self::Detect => "",
            Self::Segment => "-seg",
            Self::Pose => "-pose",
            Self::Classify => "-cls",
            Self::Obb => "-obb",
            Self::Semantic => "-sem",
        }
    }

    /// Default nano `YOLO26` model filename for this task.
    ///
    /// Used by the CLI to auto-pick a model when `--model` is omitted but `--task` is set.
    /// `YOLO26`, `YOLO11`, and `YOLOv8` variants are all auto-downloadable.
    ///
    /// ```
    /// use ultralytics_inference::Task;
    /// assert_eq!(Task::Detect.default_model(), "yolo26n.onnx");
    /// assert_eq!(Task::Segment.default_model(), "yolo26n-seg.onnx");
    /// ```
    #[must_use]
    pub fn default_model(&self) -> String {
        format!("yolo26n{}.onnx", self.model_suffix())
    }

    /// Returns `true` when the task outputs bounding boxes: Detect, Segment, Pose, and Obb.
    #[must_use]
    pub const fn has_boxes(&self) -> bool {
        matches!(self, Self::Detect | Self::Segment | Self::Pose | Self::Obb)
    }

    /// Returns `true` only for the Segment task, which outputs per-instance segmentation masks.
    #[must_use]
    pub const fn has_masks(&self) -> bool {
        matches!(self, Self::Segment)
    }

    /// Returns `true` only for the Pose task, which outputs skeletal keypoints.
    #[must_use]
    pub const fn has_keypoints(&self) -> bool {
        matches!(self, Self::Pose)
    }

    /// Returns `true` only for the Classify task, which outputs global class probabilities.
    #[must_use]
    pub const fn has_probs(&self) -> bool {
        matches!(self, Self::Classify)
    }

    /// Returns `true` only for the Obb task, which outputs oriented (rotated) bounding boxes.
    #[must_use]
    pub const fn has_obb(&self) -> bool {
        matches!(self, Self::Obb)
    }

    /// Returns `true` only for the `Semantic` task, which outputs a per-pixel class label map.
    #[must_use]
    pub const fn has_semantic_mask(&self) -> bool {
        matches!(self, Self::Semantic)
    }
}

impl fmt::Display for Task {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl FromStr for Task {
    type Err = TaskParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "detect" | "detection" => Ok(Self::Detect),
            "segment" | "segmentation" => Ok(Self::Segment),
            "pose" | "keypoint" | "keypoints" => Ok(Self::Pose),
            "classify" | "classification" | "cls" => Ok(Self::Classify),
            "obb" | "oriented" => Ok(Self::Obb),
            "semantic" | "semantic_segmentation" | "semseg" => Ok(Self::Semantic),
            _ => Err(TaskParseError(s.to_string())),
        }
    }
}

/// Error returned when parsing an invalid task string.
#[derive(Debug, Clone)]
pub struct TaskParseError(String);

impl fmt::Display for TaskParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "invalid task '{}', expected one of: detect, segment, pose, classify, obb, semantic",
            self.0
        )
    }
}

impl std::error::Error for TaskParseError {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_task_from_str() {
        assert_eq!("detect".parse::<Task>().unwrap(), Task::Detect);
        assert_eq!("segment".parse::<Task>().unwrap(), Task::Segment);
        assert_eq!("pose".parse::<Task>().unwrap(), Task::Pose);
        assert_eq!("classify".parse::<Task>().unwrap(), Task::Classify);
        assert_eq!("obb".parse::<Task>().unwrap(), Task::Obb);

        assert_eq!("semantic".parse::<Task>().unwrap(), Task::Semantic);

        // Alternative names
        assert_eq!("detection".parse::<Task>().unwrap(), Task::Detect);
        assert_eq!("segmentation".parse::<Task>().unwrap(), Task::Segment);
        assert_eq!("keypoints".parse::<Task>().unwrap(), Task::Pose);
        assert_eq!("cls".parse::<Task>().unwrap(), Task::Classify);
        assert_eq!(
            "semantic_segmentation".parse::<Task>().unwrap(),
            Task::Semantic
        );
    }

    #[test]
    fn test_task_display() {
        assert_eq!(Task::Detect.to_string(), "detect");
        assert_eq!(Task::Segment.to_string(), "segment");
        assert_eq!(Task::Semantic.to_string(), "semantic");
    }

    #[test]
    fn test_task_capabilities() {
        assert!(Task::Detect.has_boxes());
        assert!(!Task::Detect.has_masks());
        assert!(Task::Segment.has_masks());
        assert!(Task::Pose.has_keypoints());
        assert!(Task::Classify.has_probs());
        assert!(Task::Obb.has_obb());
        assert!(Task::Semantic.has_semantic_mask());
        assert!(!Task::Detect.has_semantic_mask());
    }

    #[test]
    fn test_task_suffix_and_default_model() {
        let cases = [
            (Task::Detect, "detect", "", "yolo26n.onnx"),
            (Task::Segment, "segment", "-seg", "yolo26n-seg.onnx"),
            (Task::Pose, "pose", "-pose", "yolo26n-pose.onnx"),
            (Task::Classify, "classify", "-cls", "yolo26n-cls.onnx"),
            (Task::Obb, "obb", "-obb", "yolo26n-obb.onnx"),
            (Task::Semantic, "semantic", "-sem", "yolo26n-sem.onnx"),
        ];
        for (task, name, suffix, model) in cases {
            assert_eq!(task.as_str(), name);
            assert_eq!(task.model_suffix(), suffix);
            assert_eq!(task.default_model(), model);
        }
    }

    #[test]
    fn test_task_from_str_aliases_and_errors() {
        assert_eq!("KEYPOINT".parse::<Task>().unwrap(), Task::Pose);
        assert_eq!("oriented".parse::<Task>().unwrap(), Task::Obb);
        assert_eq!("semseg".parse::<Task>().unwrap(), Task::Semantic);
        assert_eq!("Classification".parse::<Task>().unwrap(), Task::Classify);
        assert!("not_a_task".parse::<Task>().is_err());
    }
}