I want to filter detections so that only the classes "No Helmet", "Person", "Rider" are detected.

[REZIZ-TER]

Search before asking

• I have searched the Supervision issues and found no similar feature requests.

Question

My model has 4 classes: "No Helmet", "Person", "Rider", "Wear a helmet"
I want to filter detections so that only the classes "No Helmet", "Person", "Rider" are detected. How do I edit or add code?

I set selected_classes = [0, 1, 2] to be "No Helmet", "Person", "Rider"
by following the example of Supervision:

import numpy as np
import supervision as s
selected_classes = [0, 2, 3]
detections = sv.Detections(...)
detections = detections[np.isin(detections.class_id, selected_classes)]

and found error

Traceback (most recent call last):
   File "d:\Private\Y3Project\python_project\PyCode\preath.py", line 187, in <module>
     main()
   File "d:\Private\Y3Project\python_project\PyCode\preath.py", line 143, in main
     mask = zone.trigger(detections=detections)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "D:\anaconda3\envs\myenv\Lib\site-packages\supervision\detection\tools\polygon_zone.py", line 82, in trigger
     clipped_detections = replace(detections, xyxy=clipped_xyxy)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "D:\anaconda3\envs\myenv\Lib\dataclasses.py", line 1501, in replace
     return obj.__class__(**changes)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "<string>", line 9, in __init__
   File "D:\anaconda3\envs\myenv\Lib\site-packages\supervision\detection\core.py", line 120, in __post_init__
     validate_detections_fields(
   File "D:\anaconda3\envs\myenv\Lib\site-packages\supervision\validators\__init__.py", line 126, in validate_detections_fields
     validate_tracker_id(tracker_id, n)
   File "D:\anaconda3\envs\myenv\Lib\site-packages\supervision\validators\__init__.py", line 77, in validate_tracker_id
     raise ValueError(
ValueError: tracker_id must be a 1D np.ndarray with shape (2,), but got shape (3,)

My code below.

import cv2
import argparse
from ultralytics import YOLO
import supervision as sv
import numpy as np
import time
from pymongo import MongoClient
from datetime import datetime
from pytz import timezone
import base64
import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')
torch.cuda.set_device(0)

mongo_uri = "myMongo"
database_name = "databasename"
collection_name = "collectionname"
# สร้าง MongoClient
client = MongoClient(mongo_uri)
database = client[database_name]
collection = database[collection_name]

cooldown_period = 10  # 1 minute

ZONE_POLYGON = np.array([
    [0.1, 0.1],
    [0.9, 0.1],
    [0.9, 0.9],
    [0.1, 0.9],
    [0.1, 0.1]
])

model_path = "D:\\Private\\Y3Project\\python_project\\runs\\detect\\train\\weights\\best.pt"
#rstp_url = 'rtsp://user:user@192.xxx.xx.x:554/'
myvdo = "D:\Private\mydrive\myvdo\ss100met.mp4"


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="YOLOv8 live")
    parser.add_argument(
        "--webcam-resolution",
        default=[1280, 720],
        nargs=2,
        type=int
    )
    args = parser.parse_args()
    return args


def save_image_to_mongodb(image_binary, count_no_helmet, count_rider):
    # Convert time to Thailand timezone
    thai_timezone = timezone('Asia/Bangkok')
    upload_time = datetime.now(thai_timezone).strftime('%Y-%m-%d')

    # Save image and time in MongoDB
    image_data = {
        "image": base64.b64encode(image_binary).decode('utf-8'),
        "upload_time": upload_time,
        "count_no_helmet": count_no_helmet,
        "count_rider": count_rider
    }

    result = collection.insert_one(image_data)
    print(f"Image uploaded successfully. Object ID: {result.inserted_id}")


def main():
    last_save_time = time.time()
    args = parse_arguments()
    frame_width, frame_height = args.webcam_resolution

    model = YOLO(model_path)
    cap = cv2.VideoCapture(myvdo)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, frame_width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, frame_height)

    box_annotator = sv.RoundBoxAnnotator(
        thickness=2
    )

    label_annotator = sv.LabelAnnotator(
        text_position=sv.Position.TOP_CENTER,
        text_thickness=2,
        text_scale=1
    )

    zone_polygon = (ZONE_POLYGON * np.array(args.webcam_resolution)).astype(int)
    zone = sv.PolygonZone(polygon=zone_polygon, frame_resolution_wh=tuple(args.webcam_resolution))
    zone_annotator = sv.PolygonZoneAnnotator(
        zone=zone,
        color=sv.Color.RED,
        thickness=2,
        text_thickness=4,
        text_scale=2
    )

    # Initialize FPS display text
    fps_text = "FPS: calculating..."

    fps_start_time = time.time()
    fps_counter = 0
    
    while True:
        ret, frame = cap.read()
        fps_counter += 1
        bg_color = (0, 0, 0)  # Black background
        text_color = (255, 255, 255)  # White text

        for result in model.track(source=frame, stream=True, persist=True):
            frame = result.orig_img
            selected_classes = [0, 1, 2]
            detections = sv.Detections.from_ultralytics(result)
            detections = detections[np.isin(detections.class_id,selected_classes)]

            if result.boxes.id is not None:
                detections.tracker_id = result.boxes.id.cpu().numpy().astype(int)

            labels = []
            if detections.tracker_id is not None:
                labels = [
                    f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}"
                    for class_id, confidence, tracker_id
                    in zip(detections.class_id, detections.confidence, detections.tracker_id)

                ]

            if labels:
                frame = box_annotator.annotate(
                    scene=frame.copy(),
                    detections=detections
                )
                frame = label_annotator.annotate(
                    scene=frame.copy(),
                    detections=detections,
                    labels=labels
                )
            else:
                print("No Labels")

            mask = zone.trigger(detections=detections)
            frame = zone_annotator.annotate(scene=frame)
            
            # ตรวจสอบว่า detections มีขนาดไม่เท่ากับ 0
            if len(detections) > 0:
                count_no_helmet = np.count_nonzero((detections.class_id == 0) & (detections.confidence > 0.5) & mask) #No Helmet
                count_rider = np.count_nonzero((detections.class_id == 2) & (detections.confidence > 0.5) & mask) #Rider

                # Print ค่าจำนวนออกมา
                print(f"จำนวนคนไม่สวมหมวก: {count_no_helmet}")
                print(f"จำนวนคนผู้ขับขี่: {count_rider}")

                if count_no_helmet >= 1 and count_rider >= 1:
                    current_time = time.time()
                    if current_time - last_save_time >= cooldown_period:
                        image_binary = cv2.imencode('.jpg', frame)[1].tobytes()
                        save_image_to_mongodb(image_binary, count_no_helmet, count_rider)
                        print("Save Images Successfully")
                        last_save_time = current_time
            else:
                count_no_helmet = 0
                count_rider = 0
                
        # Update FPS text
        fps_counter += 1
        if time.time() - fps_start_time >= 1:
            fps = fps_counter / (time.time() - fps_start_time)
            fps_counter = 0
            fps_start_time = time.time()
            fps_text = f"FPS: {fps:.2f}"     
            
        cv2.rectangle(frame, (10, 10), (200, 50), bg_color, -1)
        cv2.putText(frame, fps_text, (15, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
        cv2.imshow("yolov8", frame)
        if (cv2.waitKey(30) == 27):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

Additional

No response

[SkalskiP]

Hi @REZIZ-TER 👋🏻, let me convert this issue into discussion and place it in Q&A section.

[SkalskiP]

@REZIZ-TER the problem is not filtering but the order of operations. In your code, you filter detections and then reassign the tracker_id value. This leads to a situation where you have more tracker_ids than boxes because you already dropped some of the boxes.

Replace this:

selected_classes = [0, 1, 2]
detections = sv.Detections.from_ultralytics(result)
detections = detections[np.isin(detections.class_id,selected_classes)]

if result.boxes.id is not None:
    detections.tracker_id = result.boxes.id.cpu().numpy().astype(int)

With this:

selected_classes = [0, 1, 2]
detections = sv.Detections.from_ultralytics(result)

if result.boxes.id is not None:
    detections.tracker_id = result.boxes.id.cpu().numpy().astype(int)

detections = detections[np.isin(detections.class_id,selected_classes)]