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    /// Returns `true` only for the `Depth` task, which outputs a per-pixel depth map.
131    #[must_use]
132    pub const fn has_depth(&self) -> bool {
133        matches!(self, Self::Depth)
134    }
135}
136
137impl fmt::Display for Task {
138    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139        f.write_str(self.as_str())
140    }
141}
142
143impl FromStr for Task {
144    type Err = TaskParseError;
145
146    fn from_str(s: &str) -> Result<Self, Self::Err> {
147        match s.to_lowercase().as_str() {
148            "detect" | "detection" => Ok(Self::Detect),
149            "segment" | "segmentation" => Ok(Self::Segment),
150            "pose" | "keypoint" | "keypoints" => Ok(Self::Pose),
151            "classify" | "classification" | "cls" => Ok(Self::Classify),
152            "obb" | "oriented" => Ok(Self::Obb),
153            "semantic" | "semantic_segmentation" | "semseg" => Ok(Self::Semantic),
154            "depth" | "depth_estimation" => Ok(Self::Depth),
155            _ => Err(TaskParseError(s.to_string())),
156        }
157    }
158}
159
160/// Error returned when parsing an invalid task string.
161#[derive(Debug, Clone)]
162pub struct TaskParseError(String);
163
164impl fmt::Display for TaskParseError {
165    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        write!(
167            f,
168            "invalid task '{}', expected one of: detect, segment, pose, classify, obb, semantic, depth",
169            self.0
170        )
171    }
172}
173
174impl std::error::Error for TaskParseError {}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn test_task_from_str() {
182        assert_eq!("detect".parse::<Task>().unwrap(), Task::Detect);
183        assert_eq!("segment".parse::<Task>().unwrap(), Task::Segment);
184        assert_eq!("pose".parse::<Task>().unwrap(), Task::Pose);
185        assert_eq!("classify".parse::<Task>().unwrap(), Task::Classify);
186        assert_eq!("obb".parse::<Task>().unwrap(), Task::Obb);
187
188        assert_eq!("semantic".parse::<Task>().unwrap(), Task::Semantic);
189        assert_eq!("depth".parse::<Task>().unwrap(), Task::Depth);
190
191        // Alternative names
192        assert_eq!("detection".parse::<Task>().unwrap(), Task::Detect);
193        assert_eq!("segmentation".parse::<Task>().unwrap(), Task::Segment);
194        assert_eq!("keypoints".parse::<Task>().unwrap(), Task::Pose);
195        assert_eq!("cls".parse::<Task>().unwrap(), Task::Classify);
196        assert_eq!(
197            "semantic_segmentation".parse::<Task>().unwrap(),
198            Task::Semantic
199        );
200        assert_eq!("depth_estimation".parse::<Task>().unwrap(), Task::Depth);
201    }
202
203    #[test]
204    fn test_task_display() {
205        assert_eq!(Task::Detect.to_string(), "detect");
206        assert_eq!(Task::Segment.to_string(), "segment");
207        assert_eq!(Task::Semantic.to_string(), "semantic");
208    }
209
210    #[test]
211    fn test_task_capabilities() {
212        assert!(Task::Detect.has_boxes());
213        assert!(!Task::Detect.has_masks());
214        assert!(Task::Segment.has_masks());
215        assert!(Task::Pose.has_keypoints());
216        assert!(Task::Classify.has_probs());
217        assert!(Task::Obb.has_obb());
218        assert!(Task::Semantic.has_semantic_mask());
219        assert!(!Task::Detect.has_semantic_mask());
220        assert!(Task::Depth.has_depth());
221        assert!(!Task::Semantic.has_depth());
222        assert!(!Task::Depth.has_semantic_mask());
223    }
224
225    #[test]
226    fn test_task_suffix_and_default_model() {
227        let cases = [
228            (Task::Detect, "detect", "", "yolo26n.onnx"),
229            (Task::Segment, "segment", "-seg", "yolo26n-seg.onnx"),
230            (Task::Pose, "pose", "-pose", "yolo26n-pose.onnx"),
231            (Task::Classify, "classify", "-cls", "yolo26n-cls.onnx"),
232            (Task::Obb, "obb", "-obb", "yolo26n-obb.onnx"),
233            (Task::Semantic, "semantic", "-sem", "yolo26n-sem.onnx"),
234            (Task::Depth, "depth", "-depth", "yolo26n-depth.onnx"),
235        ];
236        for (task, name, suffix, model) in cases {
237            assert_eq!(task.as_str(), name);
238            assert_eq!(task.model_suffix(), suffix);
239            assert_eq!(task.default_model(), model);
240        }
241    }
242
243    #[test]
244    fn test_task_from_str_aliases_and_errors() {
245        assert_eq!("KEYPOINT".parse::<Task>().unwrap(), Task::Pose);
246        assert_eq!("oriented".parse::<Task>().unwrap(), Task::Obb);
247        assert_eq!("semseg".parse::<Task>().unwrap(), Task::Semantic);
248        assert_eq!("Classification".parse::<Task>().unwrap(), Task::Classify);
249        assert!("not_a_task".parse::<Task>().is_err());
250    }
251}
ultralytics_inference/task.rs

ultralytics_inference/
task.rs
