[YOLO] YOLOv11 Object Segmentation

1. Install

1.1 Conda env create

conda create -n yolov11 python=3.8 -y

• conda 환경을 생성해 준다
• python version은 개인의 CUDA version을 참고하여 맞춰주면 된다

 

1.2 Pytorch install

https://pytorch.org/get-started/previous-versions/

Previous PyTorch Versions

• 위 Pytorch 페이지에서 자신의 CUDA version에 맞는 pip install을 설치
• (기타 opencv-python, numpy 등 필요 패키지 설치)

 

 

2. 필요 패키지 import

import torch
import torch.optim
import model
import numpy as np
import cv2
from ultralytics import YOLO

 

 

3. YOLO model 불러오기

yolo = YOLO('yolo11x-seg.pt')
yolo.eval()

• segmentation을 위해 'yolo11x-seg.pt'를 불러옴 (model path를 입력하거나 이름으로 다운로드할 수 있다)

 

 

 

4. Predict (Segmentation)

image_path = './test_image.jpg'
image = cv2.imread(image_path)
image_mask = np.zeros(image.shape, np.uint8)

results = model(image)[0]

• opencv를 통해 image를 불러옴
• np.zeros 함수를 통해 image와 같은 크기의 모든 값이 0인, mask 이미지 미리 생성
• model(image)로 predict(segmentation) 실행

 

 

5. 결과 후처리

for mask in results.masks:
    xy = mask.xy[0].astype(np.int32)
    xy = xy.reshape(-1, 1, 2)
    
    color = list(np.random.random(size=3) * 256)
    cv2.drawContours(image_mask, [xy], -1, color, cv2.FILLED)
    
image_result = cv2.addWeighted(image, 0.6, image_mask, 0.5, 0)

cv2.imshow('masks', image_mask)
cv2.imshow('result', image_result)
cv2.waitKey(0)

• mask.xy : 텐서로 표현된 픽셀 좌표의 키포인트 목록
• xy.reshape를 통해 cv2에 맞는 형식으로 변경
• cv2.drawContours 함수를 통해 image_mask에 mask를 그려줌
• cv2.addWeighted 함수를 통해 원본 이미지에 image_mask를 가중합

결과 이미지

 

 

전체 코드

from ultralytics import YOLO
import numpy as np
import cv2


model = YOLO('yolo11x-seg.pt')

image_path = './test_image.jpg'
image = cv2.imread(image_path)
image_mask = np.zeros(image.shape, np.uint8)

results = model(image)[0]

for mask in results.masks:
    xy = mask.xy[0].astype(np.int32)
    xy = xy.reshape(-1, 1, 2)
    
    color = list(np.random.random(size=3) * 256)
    cv2.drawContours(image_mask, [xy], -1, color, cv2.FILLED)
    
image_result = cv2.addWeighted(image, 0.6, image_mask, 0.5, 0)

cv2.imshow('masks', image_mask)
cv2.imshow('result', image_result)
cv2.waitKey(0)