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])

Related Posts