augmentedstartups / AS-One

Easy & Modular Computer Vision Detectors, Trackers & SAM - Run YOLOv9,v8,v7,v6,v5,R,X in under 10 lines of code.
https://www.augmentedstartups.com
GNU General Public License v3.0
603 stars 102 forks source link

How can we detect and track only a certain class (for e.g. person)? #4

Closed learninguser closed 1 year ago

learninguser commented 1 year ago

Hi, Can you please let me know how can we detect and track only humans in a video sequence?

MNMaqsood commented 1 year ago

Currently you have to filter the ids based on class_id. Let us also add this in next update

MNMaqsood commented 1 year ago

This has been done in new github updates.

To filter classes in detector pass filter_classes list containing class names.

detector = Detector(asone.YOLOV7_E6_ONNX, use_cuda=True).get_detector() # Set use_cuda to False for cpu

filter_classes = ['person'] # Set to None to detect all classes

dets, img_info = detector.detect(img, , filter_classes=filter_classes)

To filter classes in detector pass filter_classes list containing class names.

filter_classes = ['person'] # set to None to track all classes

# Get tracking function
track_fn = dt_obj.track_video('data/sample_videos/test.mp4', output_dir='data/results', save_result=True, display=True, filter_classes=filter_classes)

Check README.md and main.py