[Python] How to create a binary mask from a yolo8 segmentation result

I want to segment an image using yolo8 and then create a mask for all objects in the image with specific class.

I have developed this code:

img=cv2.imread('images/bus.jpg')
model = YOLO('yolov8m-seg.pt')
results = model.predict(source=img.copy(), save=False, save_txt=False)
class_ids = np.array(results[0].boxes.cls.cpu(), dtype="int")
for i in range(len(class_ids)):
if class_ids==0:
empty_image = np.zeros((height, width,3), dtype=np.uint8)
res_plotted = results[0].plot(boxes=0, img=empty_image)



In the above code, res_plotted is the mask for one object, in RGB. I want to add all of these images to each other and create a mask for all objects with class 0 (it is a pedestrian in this example)

My questions:

1. How can I complete this code?
2. Is there any better way to achieve this without having a loop?
3. Is there any utility function in the yolo8 library to do this?

Continue reading...