Instant Segmentation by YOLOv8

If we want to use YOLOv8 for object detection, here is a good example.

What if I want to use YOLOv8 to segment a picture, crop out the object, and paste (only paste the object, not the pixels near it) it to a new picture? I wrote an example:

import cv2
import numpy as np

from ultralytics import YOLO

IMG_SIZE = 2048

def main(path):
    filename, file_extension = os.path.splitext(path)

    # Load segment model of yolov8
    model = YOLO("yolov8x-seg.pt")
    img = cv2.imread(path)

    results = model(img, imgsz=(img.shape[1], img.shape[0]))

    count = 0
    for res in results:
        for mask in res.masks.xy:
            polygan = mask.reshape((-1, 1, 2)).astype(np.int32)
            x, y, w, h = cv2.boundingRect(polygan)

            # Create mask with all value of 255
            binary_mask = np.ones(img.shape, dtype=np.uint8) * 255
            # Fill the polygan (the object we want) with zero
            cv2.fillPoly(binary_mask, [polygan], (0, 0, 0))
            # Add zero polygan with origin image could keep object, and push background to 255
            out_img = cv2.add(img, binary_mask)[y:y+h, x:x+w]

            cv2.imwrite(f"{filename}_{count}{file_extension}", out_img)
            count += 1
    print(f"Total: {count}")

if __name__ == "__main__":
    main(sys.argv[1])

import cv2

import numpy as np

​

from ultralytics import YOLO

​

IMG_SIZE = 2048

​

def main(path):

    filename, file_extension = os.path.splitext(path)

​

    # Load segment model of yolov8

    model = YOLO("yolov8x-seg.pt")

    img = cv2.imread(path)

​

    results = model(img, imgsz=(img.shape[1], img.shape[0]))

​

    count = 0

    for res in results:

        for mask in res.masks.xy:

            polygan = mask.reshape((-1, 1, 2)).astype(np.int32)

            x, y, w, h = cv2.boundingRect(polygan)

​

            # Create mask with all value of 255

            binary_mask = np.ones(img.shape, dtype=np.uint8) * 255

            # Fill the polygan (the object we want) with zero

            cv2.fillPoly(binary_mask, [polygan], (0, 0, 0))

            # Add zero polygan with origin image could keep object, and push background to 255

            out_img = cv2.add(img, binary_mask)[y:y+h, x:x+w]

​

            cv2.imwrite(f"{filename}_{count}{file_extension}", out_img)

            count += 1

    print(f"Total: {count}")

​

if __name__ == "__main__":

    main(sys.argv[1])