Practical.kr

실용적인 소프트웨어를 만듭니다.

SAM과 YOLOv8을 활용한 객체 세그멘테이션

지난번 글에서 Supervision과 ByteTrack을 이용해서 도로에서 차량을 찾고(Yolo) 트래킹(ByteTracking) 하는 예제를 만들어 봤는데, 이글은 그 연장선에서 목표를 찾아서 찾아서 박스만 두르는 것이 아니고 정확한 형상을 추출(세그멘테이션)하는 기능을 구현하기 위하여 Segment Anything Model (SAM)과 YOLOv8을 결합하여 사람과 공(sports ball) 객체를 정확하게 탐지하고 세그멘테이션하는 시스템을 구현해봤다. UI는 Streamlit을 활용했다.

셈플 이미지는 어제 축구 경기의 한장면. 5-0 이라니… ㅠ,ㅠ


기술 스택

  • SAM (Segment Anything Model): Meta의 최신 세그멘테이션 모델
  • YOLOv8: Ultralytics의 객체 탐지 모델
  • Streamlit: 웹 애플리케이션 프레임워크
  • OpenCV: 이미지 처리
  • PyTorch: 딥러닝 프레임워크

시스템 아키텍처

1. YOLOv8 객체 탐지

# 선수(사람) (클래스 ID: 0) 또는 공 (클래스 ID: 32 - sports ball)
if (cls_id == 0 or cls_id == 32) and conf > 0.5:
    object_type = "person" if cls_id == 0 else "ball"

YOLOv8을 사용하여 COCO 데이터셋 기반으로 사람과 공을 탐지해서 YOLO를 이용해 바운딩 박스를 만든다.

3. SAM 세그멘테이션

predictor = SamPredictor(sam)
predictor.set_image(image_array)

input_box = np.array([x1, y1, x2, y2])
masks, scores, logits = predictor.predict(
    point_coords=None,
    point_labels=None,
    box=input_box[None, :],
    multimask_output=False,
)

그리고 세그멘테이션을 실행한다

핵심 구현 세부사항

1. 이중 모델 파이프라인

1단계: YOLO 탐지

  • 빠른 객체 탐지를 통한 ROI(Region of Interest) 식별
  • 바운딩 박스 좌표 추출
  • 객체 타입 분류 (사람/공)

2단계: SAM 세그멘테이션

  • YOLO의 바운딩 박스를 SAM의 입력으로 활용
  • 정밀한 픽셀 단위 마스크 생성
  • 객체 경계의 정확한 분할

2. 데이터 구조 설계

all_masks.append({
    'segmentation': best_mask,
    'area': np.sum(best_mask),
    'object_id': i+1,
    'object_type': detection['type'],
    'bbox': [x1, y1, x2, y2]
})

각 마스크 객체는 세그멘테이션 마스크, 면적, 객체 ID, 타입, 바운딩 박스 정보를 포함.
이 단계는 중간 단계로 다음은 동영상에서 각각의 프레임별도 위의 작업을 진행하여 동영상으로 다시 합치면 박스 오브젝트가 아닌 세그멘테이션 단위의 오브젝트 트래킹이 가능하다.

Tags: