ultralytics_inference/
task.rs

1// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
3//! Task definitions for YOLO models.
4//!
5//! This module defines the different tasks that YOLO models can perform,
6//! along with their associated capabilities and string representations.
7
8use std::fmt;
9use std::str::FromStr;
10
11/// YOLO model task types.
12///
13/// Each task type corresponds to a different computer vision problem
14/// that YOLO models can solve. The task type determines the expected
15/// model outputs and post-processing steps.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
17pub enum Task {
18    /// Object detection.
19    /// Predicts bounding boxes and class labels for objects in an image.
20    #[default]
21    Detect,
22    /// Instance segmentation.
23    /// Predicts bounding boxes, class labels, and pixel-level masks for objects.
24    Segment,
25    /// Pose estimation.
26    /// Predicts bounding boxes and skeletal keypoints for objects (e.g., humans).
27    Pose,
28    /// Image classification.
29    /// Predicts class probabilities for the entire image (no localization).
30    Classify,
31    /// Oriented bounding box detection (OBB).
32    /// Predicts rotated bounding boxes for objects, useful for aerial imagery etc.
33    Obb,
34    /// Semantic segmentation.
35    /// Assigns a class label to every pixel in the image.
36    Semantic,
37    /// Monocular depth estimation.
38    /// Predicts a per-pixel depth map (in meters) for the whole image.
39    Depth,
40}
41
42impl Task {
43    /// Get the string representation used in ONNX model metadata
44    /// (e.g. `"detect"`, `"segment"`).
45    #[must_use]
46    pub const fn as_str(&self) -> &'static str {
47        match self {
48            Self::Detect => "detect",
49            Self::Segment => "segment",
50            Self::Pose => "pose",
51            Self::Classify => "classify",
52            Self::Obb => "obb",
53            Self::Semantic => "semantic",
54            Self::Depth => "depth",
55        }
56    }
57
58    /// ONNX filename suffix for this task, used to construct `{family}n{suffix}.onnx`
59    /// (e.g. `yolo26n-seg.onnx`, `yolo11n-pose.onnx`, `yolov8n.onnx`).
60    ///
61    /// ```
62    /// use ultralytics_inference::Task;
63    /// assert_eq!(Task::Detect.model_suffix(), "");
64    /// assert_eq!(Task::Segment.model_suffix(), "-seg");
65    /// ```
66    #[must_use]
67    pub const fn model_suffix(&self) -> &'static str {
68        match self {
69            Self::Detect => "",
70            Self::Segment => "-seg",
71            Self::Pose => "-pose",
72            Self::Classify => "-cls",
73            Self::Obb => "-obb",
74            Self::Semantic => "-sem",
75            Self::Depth => "-depth",
76        }
77    }
78
79    /// Default nano `YOLO26` model filename for this task.
80    ///
81    /// Used by the CLI to auto-pick a model when `--model` is omitted but `--task` is set.
82    /// `YOLO26`, `YOLO11`, and `YOLOv8` variants are all auto-downloadable.
83    ///
84    /// ```
85    /// use ultralytics_inference::Task;
86    /// assert_eq!(Task::Detect.default_model(), "yolo26n.onnx");
87    /// assert_eq!(Task::Segment.default_model(), "yolo26n-seg.onnx");
88    /// ```
89    #[must_use]
90    pub fn default_model(&self) -> String {
91        format!("yolo26n{}.onnx", self.model_suffix())
92    }
93
94    /// Returns `true` when the task outputs bounding boxes: Detect, Segment, Pose, and Obb.
95    #[must_use]
96    pub const fn has_boxes(&self) -> bool {
97        matches!(self, Self::Detect | Self::Segment | Self::Pose | Self::Obb)
98    }
99
100    /// Returns `true` only for the Segment task, which outputs per-instance segmentation masks.
101    #[must_use]
102    pub const fn has_masks(&self) -> bool {
103        matches!(self, Self::Segment)
104    }
105
106    /// Returns `true` only for the Pose task, which outputs skeletal keypoints.
107    #[must_use]
108    pub const fn has_keypoints(&self) -> bool {
109        matches!(self, Self::Pose)
110    }
111
112    /// Returns `true` only for the Classify task, which outputs global class probabilities.
113    #[must_use]
114    pub const fn has_probs(&self) -> bool {
115        matches!(self, Self::Classify)
116    }
117
118    /// Returns `true` only for the Obb task, which outputs oriented (rotated) bounding boxes.
119    #[must_use]
120    pub const fn has_obb(&self) -> bool {
121        matches!(self, Self::Obb)
122    }
123
124    /// Returns `true` only for the `Semantic` task, which outputs a per-pixel class label map.
125    #[must_use]
126    pub const fn has_semantic_mask(&self) -> bool {
127        matches!(self, Self::Semantic)
128    }
129}
130
131impl fmt::Display for Task {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        f.write_str(self.as_str())
134    }
135}
136
137impl FromStr for Task {
138    type Err = TaskParseError;
139
140    fn from_str(s: &str) -> Result<Self, Self::Err> {
141        match s.to_lowercase().as_str() {
142            "detect" | "detection" => Ok(Self::Detect),
143            "segment" | "segmentation" => Ok(Self::Segment),
144            "pose" | "keypoint" | "keypoints" => Ok(Self::Pose),
145            "classify" | "classification" | "cls" => Ok(Self::Classify),
146            "obb" | "oriented" => Ok(Self::Obb),
147            "semantic" | "semantic_segmentation" | "semseg" => Ok(Self::Semantic),
148            "depth" | "depth_estimation" => Ok(Self::Depth),
149            _ => Err(TaskParseError(s.to_string())),
150        }
151    }
152}
153
154/// Error returned when parsing an invalid task string.
155#[derive(Debug, Clone)]
156pub struct TaskParseError(String);
157
158impl fmt::Display for TaskParseError {
159    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
160        write!(
161            f,
162            "invalid task '{}', expected one of: detect, segment, pose, classify, obb, semantic, depth",
163            self.0
164        )
165    }
166}
167
168impl std::error::Error for TaskParseError {}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn test_task_from_str() {
176        assert_eq!("detect".parse::<Task>().unwrap(), Task::Detect);
177        assert_eq!("segment".parse::<Task>().unwrap(), Task::Segment);
178        assert_eq!("pose".parse::<Task>().unwrap(), Task::Pose);
179        assert_eq!("classify".parse::<Task>().unwrap(), Task::Classify);
180        assert_eq!("obb".parse::<Task>().unwrap(), Task::Obb);
181
182        assert_eq!("semantic".parse::<Task>().unwrap(), Task::Semantic);
183        assert_eq!("depth".parse::<Task>().unwrap(), Task::Depth);
184
185        // Alternative names
186        assert_eq!("detection".parse::<Task>().unwrap(), Task::Detect);
187        assert_eq!("segmentation".parse::<Task>().unwrap(), Task::Segment);
188        assert_eq!("keypoints".parse::<Task>().unwrap(), Task::Pose);
189        assert_eq!("cls".parse::<Task>().unwrap(), Task::Classify);
190        assert_eq!(
191            "semantic_segmentation".parse::<Task>().unwrap(),
192            Task::Semantic
193        );
194        assert_eq!("depth_estimation".parse::<Task>().unwrap(), Task::Depth);
195    }
196
197    #[test]
198    fn test_task_display() {
199        assert_eq!(Task::Detect.to_string(), "detect");
200        assert_eq!(Task::Segment.to_string(), "segment");
201        assert_eq!(Task::Semantic.to_string(), "semantic");
202    }
203
204    #[test]
205    fn test_task_capabilities() {
206        assert!(Task::Detect.has_boxes());
207        assert!(!Task::Detect.has_masks());
208        assert!(Task::Segment.has_masks());
209        assert!(Task::Pose.has_keypoints());
210        assert!(Task::Classify.has_probs());
211        assert!(Task::Obb.has_obb());
212        assert!(Task::Semantic.has_semantic_mask());
213        assert!(!Task::Detect.has_semantic_mask());
214    }
215
216    #[test]
217    fn test_task_suffix_and_default_model() {
218        let cases = [
219            (Task::Detect, "detect", "", "yolo26n.onnx"),
220            (Task::Segment, "segment", "-seg", "yolo26n-seg.onnx"),
221            (Task::Pose, "pose", "-pose", "yolo26n-pose.onnx"),
222            (Task::Classify, "classify", "-cls", "yolo26n-cls.onnx"),
223            (Task::Obb, "obb", "-obb", "yolo26n-obb.onnx"),
224            (Task::Semantic, "semantic", "-sem", "yolo26n-sem.onnx"),
225            (Task::Depth, "depth", "-depth", "yolo26n-depth.onnx"),
226        ];
227        for (task, name, suffix, model) in cases {
228            assert_eq!(task.as_str(), name);
229            assert_eq!(task.model_suffix(), suffix);
230            assert_eq!(task.default_model(), model);
231        }
232    }
233
234    #[test]
235    fn test_task_from_str_aliases_and_errors() {
236        assert_eq!("KEYPOINT".parse::<Task>().unwrap(), Task::Pose);
237        assert_eq!("oriented".parse::<Task>().unwrap(), Task::Obb);
238        assert_eq!("semseg".parse::<Task>().unwrap(), Task::Semantic);
239        assert_eq!("Classification".parse::<Task>().unwrap(), Task::Classify);
240        assert!("not_a_task".parse::<Task>().is_err());
241    }
242}
ultralytics_inference/task.rs

ultralytics_inference/
task.rs
