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