feat: ✨ Add class filtering support to inference and CLI arguments (#72)

Signed-off-by: Onuralp SEZER <onuralp@ultralytics.com>

Author: onuralpszr

README.md

@@ -401,6 +401,7 @@ ONNX Runtime threading is set to auto (`num_threads: 0`) which lets ORT choose o
 - [x] FP16 half-precision inference
 - [x] Batch inference support
 - [x] Rectangular inference support and optimization
+- [x] Class filtering support
 
 ### In Progress
 

src/cli/args.rs

@@ -22,6 +22,7 @@ use clap::{Args, Parser, Subcommand};
     --show                 Display results in a window
     --device <DEVICE>      Device (cpu, cuda:0, mps, coreml, directml:0, openvino, tensorrt:0, xnnpack)
     --verbose              Show verbose output [default: true]
+    --classes <CLASSES>    Filter by class IDs (e.g., "0", "0,1,2", "[0, 1]")
 
 Examples:
     ultralytics-inference predict
@@ -31,7 +32,7 @@ Examples:
     ultralytics-inference predict --source 0 --conf 0.5 --show
     ultralytics-inference predict --source assets/ --save --half
     ultralytics-inference predict --source image.jpg --device cuda:0
-    ultralytics-inference predict --source image.jpg --device mps"#)]
+    ultralytics-inference predict --source image.jpg --classes 0"#)]
 pub struct Cli {
     #[command(subcommand)]
     /// Subcommand to execute.
@@ -104,6 +105,43 @@ pub struct PredictArgs {
     /// Show verbose output
     #[arg(long, default_value_t = true, action = clap::ArgAction::Set)]
     pub verbose: bool,
+
+    /// Filter by class IDs (e.g. 0 or "0,1,2" or "[0, 1, 2]")
+    ///
+    /// Supported formats:
+    /// - Single integer: --classes 0
+    /// - Comma-separated list: --classes "0,1,2"
+    /// - List syntax: --classes "[0, 1, 2]"
+    ///
+    /// Note: When passing a list directly without quotes, avoid spaces to prevent
+    /// shell argument parsing issues (e.g. use --classes 0,1 not --classes 0, 1).
+    #[arg(long, allow_hyphen_values = true)]
+    pub classes: Option<String>,
+}
+
+/// Parse class IDs from various formats: "1,2,3", "[1,2,3]", "(1,2,3)"
+///
+/// # Errors
+///
+/// Returns a `String` error message if any segment of the string cannot be parsed as a `usize`.
+pub fn parse_classes(s: &str) -> Result<Vec<usize>, String> {
+    // Remove brackets, parentheses, and quotes
+    let cleaned = s
+        .trim()
+        .trim_matches(|c| c == '[' || c == ']' || c == '(' || c == ')' || c == '"' || c == '\'');
+
+    if cleaned.is_empty() {
+        return Ok(Vec::new());
+    }
+
+    cleaned
+        .split(',')
+        .map(|part| {
+            part.trim()
+                .parse::<usize>()
+                .map_err(|e| format!("Invalid class ID '{}': {}", part.trim(), e))
+        })
+        .collect()
 }
 
 #[cfg(test)]

src/cli/predict.rs

@@ -84,6 +84,21 @@ pub fn run_prediction(args: &PredictArgs) {
         config = config.with_device(d.clone());
     }
 
+    // Apply class filter if specified
+    if let Some(classes_str) = &args.classes {
+        match crate::cli::args::parse_classes(classes_str) {
+            Ok(classes) => {
+                if !classes.is_empty() {
+                    config = config.with_classes(classes);
+                }
+            }
+            Err(e) => {
+                error!("Error parsing classes: {e}");
+                process::exit(1);
+            }
+        }
+    }
+
     let mut model = match YOLOModel::load_with_config(model_path, config) {
         Ok(m) => m,
         Err(e) => {
@@ -423,7 +438,7 @@ fn format_class_counts(
 /// Format detection summary like "4 persons, 1 bus".
 #[allow(clippy::option_if_let_else)]
 fn format_detection_summary(result: &Results) -> String {
-    if let Some(ref boxes) = result.boxes {
+    let summary = if let Some(ref boxes) = result.boxes {
         format_class_counts(&boxes.cls(), boxes.len(), &result.names)
     } else if let Some(ref obb) = result.obb {
         format_class_counts(&obb.cls(), obb.len(), &result.names)
@@ -439,6 +454,12 @@ fn format_detection_summary(result: &Results) -> String {
         parts.join(", ")
     } else {
         String::new()
+    };
+
+    if summary.is_empty() {
+        "(no detections)".to_string()
+    } else {
+        summary
     }
 }
 
@@ -527,7 +548,7 @@ mod tests {
         result.boxes = Some(boxes);
 
         let summary = format_detection_summary(&result);
-        assert!(summary.is_empty());
+        assert_eq!(summary, "(no detections)");
     }
 
     /// Test `format_detection_summary` with OBB detections.
@@ -573,7 +594,7 @@ mod tests {
         result.obb = Some(obb);
 
         let summary = format_detection_summary(&result);
-        assert!(summary.is_empty());
+        assert_eq!(summary, "(no detections)");
     }
 
     /// Test `format_detection_summary` with classification probs.
@@ -616,7 +637,7 @@ mod tests {
         );
 
         let summary = format_detection_summary(&result);
-        assert!(summary.is_empty());
+        assert_eq!(summary, "(no detections)");
     }
 
     /// Test `format_detection_summary` with unknown class (uses "object" fallback).

src/inference.rs

@@ -69,6 +69,9 @@ pub struct InferenceConfig {
     /// Whether to use minimal padding (rectangular inference).
     /// Defaults to `true` to match Ultralytics Python.
     pub rect: bool,
+    /// Class IDs to filter predictions. If `None`, all classes are returned.
+    /// Useful for focusing on specific objects in multi-class detection tasks.
+    pub classes: Option<Vec<usize>>,
 }
 
 impl Default for InferenceConfig {
@@ -85,6 +88,7 @@ impl Default for InferenceConfig {
             save: Self::DEFAULT_SAVE,
             save_frames: Self::DEFAULT_SAVE_FRAMES,
             rect: Self::DEFAULT_RECT,
+            classes: None,
         }
     }
 }
@@ -302,6 +306,47 @@ impl InferenceConfig {
         self.rect = rect;
         self
     }
+
+    /// Set the class IDs to filter predictions.
+    ///
+    /// Only detections belonging to the specified classes will be returned.
+    ///
+    /// # Arguments
+    ///
+    /// * `classes` - A vector of class IDs to keep.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use ultralytics_inference::InferenceConfig;
+    ///
+    /// // Only detect persons (class 0) and cars (class 2)
+    /// let config = InferenceConfig::new()
+    ///     .with_classes(vec![0, 2]);
+    /// ```
+    ///
+    /// # Returns
+    ///
+    /// * The modified `InferenceConfig`.
+    #[must_use]
+    pub fn with_classes(mut self, classes: Vec<usize>) -> Self {
+        self.classes = Some(classes);
+        self
+    }
+    /// Check if a class should be included in the results.
+    ///
+    /// # Arguments
+    ///
+    /// * `class_id` - The class index to check.
+    ///
+    /// # Returns
+    ///
+    /// * `true` if the class should be kept.
+    /// * `false` if the class should be filtered out.
+    #[must_use]
+    pub fn keep_class(&self, class_id: usize) -> bool {
+        self.classes.as_ref().is_none_or(|c| c.contains(&class_id))
+    }
 }
 
 #[cfg(test)]
@@ -331,4 +376,22 @@ mod tests {
         assert_eq!(config.imgsz, Some((640, 640)));
         assert_eq!(config.num_threads, 8);
     }
+
+    #[test]
+    fn test_keep_class() {
+        let config = InferenceConfig::default();
+        // Default: no filtering -> keep all
+        assert!(config.keep_class(0));
+        assert!(config.keep_class(100));
+
+        let config_filtered = InferenceConfig::new().with_classes(vec![1, 3]);
+        // Class 1 is in list -> keep
+        assert!(config_filtered.keep_class(1));
+        // Class 3 is in list -> keep
+        assert!(config_filtered.keep_class(3));
+        // Class 0 is NOT in list -> filter out (keep = false)
+        assert!(!config_filtered.keep_class(0));
+        // Class 2 is NOT in list -> filter out (keep = false)
+        assert!(!config_filtered.keep_class(2));
+    }
 }

src/postprocessing.rs

@@ -290,6 +290,13 @@ fn extract_detect_boxes(
 
         for (idx, &score) in max_scores.iter().enumerate() {
             if score > conf_thresh {
+                let best_class = max_classes[idx];
+
+                // Filter by class if specified
+                if !config.keep_class(best_class) {
+                    continue;
+                }
+
                 let cx = unsafe { *output.get_unchecked(idx) };
                 let cy = unsafe { *output.get_unchecked(num_predictions + idx) };
                 let w = unsafe { *output.get_unchecked(2 * num_predictions + idx) };
@@ -303,7 +310,7 @@ fn extract_detect_boxes(
                 candidates.push(Candidate {
                     bbox: [x1, y1, x2, y2],
                     score,
-                    class: max_classes[idx],
+                    class: best_class,
                 });
             }
         }
@@ -343,6 +350,11 @@ fn extract_detect_boxes(
             }
 
             if found {
+                // Filter by class if specified
+                if !config.keep_class(best_class) {
+                    continue;
+                }
+
                 let cx = unsafe { *output.get_unchecked(base) };
                 let cy = unsafe { *output.get_unchecked(base + 1) };
                 let w = unsafe { *output.get_unchecked(base + 2) };
@@ -604,6 +616,11 @@ fn postprocess_segment(
         let scaled = scale_coords(&[x1, y1, x2, y2], preprocess.scale, preprocess.padding);
         let clipped = clip_coords(&scaled, preprocess.orig_shape);
 
+        // Filter by class if specified
+        if !config.keep_class(best_class) {
+            continue;
+        }
+
         candidates.push((
             [clipped[0], clipped[1], clipped[2], clipped[3]],
             best_score,
@@ -941,6 +958,11 @@ fn postprocess_pose(
             keypoints.push([scaled_x, scaled_y, kpt_conf]);
         }
 
+        // Filter by class if specified
+        if !config.keep_class(best_class) {
+            continue;
+        }
+
         candidates.push((
             [clipped[0], clipped[1], clipped[2], clipped[3]],
             best_score,
@@ -1178,6 +1200,11 @@ fn postprocess_obb(
         #[allow(clippy::cast_precision_loss)]
         let clipped_cy = scaled_cy.max(0.0).min(oh as f32);
 
+        // Filter by class if specified
+        if !config.keep_class(best_class) {
+            continue;
+        }
+
         candidates.push((
             [clipped_cx, clipped_cy, scaled_w, scaled_h, angle],
             best_score,

tests/integration_test.rs

@@ -29,6 +29,7 @@ fn test_run_prediction_e2e() {
         device: None,
         verbose: true,
         rect: false,
+        classes: None,
     };
 
     // This should run successfully (download model/images and predict)