import numpy as np
import supervision as sv

from ultralytics import YOLO


# load model
model = YOLO('yolov8x-seg.pt')

# load first video frame
frames_generator = sv.get_video_frames_generator(source_path='walking.mp4')
frames_iterator = iter(frames_generator)
frame = next(frames_iterator)

# run inference
result = model.track(frame, imgsz=1920)[0]
detections = sv.Detections.from_ultralytics(result)

# map category id to category name
detections['names'] = [
    model.model.names[class_id]
    for class_id
    in detections.class_id
]

# split detections based on category
person_detections = detections[detections.class_id == 0]
nonperson_detections = detections[detections.class_id != 0]

# merge detections
merged_detections = sv.Detections.merge([person_detections, nonperson_detections])

print(merged_detections['names'])