jamjamjon/usls

YOLO - Restrict detection classes

pplanel opened this issue · 3 comments

Hi, thanks for this library, great work!
As the title says, how could I restrict the detections classes?
And the options/parameters prefixed with i1-275 are related to what?

Thanks!

Hi,

First, the options struct does not provide an option for restricting detection classes. Therefore, to filter out certain classes, you would need to extract the desired classes from the detection results. Here’s an example code snippet for reference:

let ys = model.forward(&xs, args.profile)?; 

for y in ys.iter() {
    if let Some(bboxes) = y.bboxes() {
        for bbox in bboxes {
            let class_id = bbox.id();
            todo!()
        }
    }
}

I will add the feature you mentioned in the next version, and the possible API could be Options::new().exclude_classes() and retain_classes().

Regarding the other question: Are you referring to Options::new().with_ixx(0, 0, (1, args.batch_size as _, 4).into())? This is used to set the input values for the ONNX model under dynamic dimensions. Here, i stands for input, the first x represents the x-th input, and the second x indicates the x-th axis of that input.

You can now use Options::new().exclude_classes() to filter out some classes, and use Options::new().retain_classes() to specify detecting only the provided classes in v0.0.18.

Refer to the YOLO demo

usls/examples/yolo/main.rs

Lines 163 to 164 in 1d59638

.exclude_classes(&[0])
// .retain_classes(&[0, 5])

Thanks!