usls/models/yolo/
config.rs

1use crate::{
2    models::YOLOPredsFormat, Config, ResizeMode, Scale, Task, NAMES_COCO_80,
3    NAMES_COCO_KEYPOINTS_17, NAMES_DOTA_V1_15, NAMES_IMAGENET_1K, NAMES_YOLO_DOCLAYOUT_10,
4};
5
6impl Config {
7    /// Creates a base YOLO configuration with common settings.
8    ///
9    /// Sets up default input dimensions (640x640) and image processing parameters.
10    pub fn yolo() -> Self {
11        Self::default()
12            .with_name("yolo")
13            .with_model_ixx(0, 0, 1.into())
14            .with_model_ixx(0, 1, 3.into())
15            .with_model_ixx(0, 2, 640.into())
16            .with_model_ixx(0, 3, 640.into())
17            .with_resize_mode(ResizeMode::FitAdaptive)
18            .with_resize_filter("CatmullRom")
19    }
20
21    /// Creates a configuration for YOLO image classification.
22    ///
23    /// Configures the model for ImageNet classification with:
24    /// - 224x224 input size
25    /// - Exact resize mode with bilinear interpolation
26    /// - ImageNet 1000 class names
27    pub fn yolo_classify() -> Self {
28        Self::yolo()
29            .with_task(Task::ImageClassification)
30            .with_model_ixx(0, 2, 224.into())
31            .with_model_ixx(0, 3, 224.into())
32            .with_resize_mode(ResizeMode::FitExact)
33            .with_resize_filter("Bilinear")
34            .with_class_names(&NAMES_IMAGENET_1K)
35    }
36
37    /// Creates a configuration for YOLO object detection.
38    ///
39    /// Configures the model for COCO dataset object detection with 80 classes.
40    pub fn yolo_detect() -> Self {
41        Self::yolo()
42            .with_task(Task::ObjectDetection)
43            .with_class_names(&NAMES_COCO_80)
44    }
45
46    /// Creates a configuration for YOLO pose estimation.
47    ///
48    /// Configures the model for human keypoint detection with 17 COCO keypoints.
49    pub fn yolo_pose() -> Self {
50        Self::yolo()
51            .with_task(Task::KeypointsDetection)
52            .with_keypoint_names(&NAMES_COCO_KEYPOINTS_17)
53    }
54
55    /// Creates a configuration for YOLO instance segmentation.
56    ///
57    /// Configures the model for COCO dataset instance segmentation with 80 classes.
58    pub fn yolo_segment() -> Self {
59        Self::yolo()
60            .with_task(Task::InstanceSegmentation)
61            .with_class_names(&NAMES_COCO_80)
62    }
63
64    /// Creates a configuration for YOLO oriented object detection.
65    ///
66    /// Configures the model for detecting rotated objects with:
67    /// - 1024x1024 input size
68    /// - DOTA v1 dataset classes
69    pub fn yolo_obb() -> Self {
70        Self::yolo()
71            .with_model_ixx(0, 2, 1024.into())
72            .with_model_ixx(0, 3, 1024.into())
73            .with_task(Task::OrientedObjectDetection)
74            .with_class_names(&NAMES_DOTA_V1_15)
75    }
76
77    /// Creates a configuration for document layout analysis using YOLOv10.
78    ///
79    /// Configures the model for detecting document structure elements with:
80    /// - Variable input size up to 1024x1024
81    /// - 10 document layout classes
82    pub fn doclayout_yolo_docstructbench() -> Self {
83        Self::yolo_detect()
84            .with_version(10.into())
85            .with_model_ixx(0, 2, (640, 1024, 1024).into())
86            .with_model_ixx(0, 3, (640, 1024, 1024).into())
87            .with_class_confs(&[0.4])
88            .with_class_names(&NAMES_YOLO_DOCLAYOUT_10)
89            .with_model_file("doclayout-docstructbench.onnx") // TODO: batch_size > 1
90    }
91
92    pub fn fastsam_s() -> Self {
93        Self::yolo_segment()
94            .with_class_names(&["object"])
95            .with_scale(Scale::S)
96            .with_version(8.into())
97            .with_model_file("FastSAM-s.onnx")
98    }
99
100    pub fn fastsam_x() -> Self {
101        Self::yolo_segment()
102            .with_class_names(&["object"])
103            .with_scale(Scale::X)
104            .with_version(8.into())
105            .with_model_file("FastSAM-x.onnx")
106    }
107
108    pub fn ultralytics_rtdetr_l() -> Self {
109        Self::yolo_detect()
110            .with_yolo_preds_format(YOLOPredsFormat::n_a_cxcywh_clss_n())
111            .with_scale(Scale::L)
112            .with_model_file("rtdetr-l.onnx")
113    }
114
115    pub fn ultralytics_rtdetr_x() -> Self {
116        Self::yolo_detect()
117            .with_yolo_preds_format(YOLOPredsFormat::n_a_cxcywh_clss_n())
118            .with_scale(Scale::X)
119            .with_model_file("rtdetr-x.onnx")
120    }
121}
usls/models/yolo/config.rs

usls/models/yolo/
config.rs
