[CV Archive] YOLOv8 Segmentation: 스타벅스 로고 추출하기

안녕하세요!

오늘은 스타벅스 로고와 텍스트를 탐지하는 실습 코드를 바탕으로, Segmentation의 개념부터 YOLOv8을 활용한 실제 학습 과정까지 완벽하게 정리해 보겠습니다.

---

1. Segmentation 개념

Segmentation(분할)은 이미지를 픽셀 단위로 분석하여 각 픽셀이 어떤 객체에 속하는지 분류하는 기술입니다. 단순히 사각형 박스를 치는 Object Detection보다 훨씬 정밀한 작업입니다.

주요 분류

• Semantic Segmentation: 같은 클래스라면 개별 객체를 구분하지 않고 하나의 덩어리로 인식합니다. (예: 모든 '사람'은 같은 색상)
• Instance Segmentation: 같은 클래스라도 개별 객체를 각각 따로 구분합니다. (예: 사람 A, 사람 B를 서로 다른 색상으로 구분)
• 실습 모델: 오늘 우리가 사용하는 YOLOv8-seg는 대표적인 Instance Segmentation 모델입니다.

---

1-1. 환경 설정 및 데이터 로드

실습을 위해 필요한 라이브러리를 설치하고, COCO 형식의 데이터셋을 불러옵니다.

# 필요 라이브러리 설치
!pip install -q ultralytics opencv-python

import os
import random
import shutil
import json
import glob
import cv2
import matplotlib.pyplot as plt
import numpy as np
from ultralytics import YOLO

# 데이터 경로 설정 (본인의 환경에 맞게 수정)
data_root = '/content/starbucks' 

# COCO 어노테이션 로드 함수
def load_coco_annotations(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data

---

1-2. 데이터 시각화 (Annotation 확인)

데이터가 제대로 라벨링 되었는지 확인하는 것은 학습 전 필수 단계입니다. cv2.polylines를 활용해 다각형(Polygon)을 그려봅니다.

def draw_annotations(image, annotations, image_id):
    for ann in annotations:
        if ann['image_id'] == image_id and 'segmentation' in ann:
            segmentation = ann['segmentation']
            for seg in segmentation:
                if isinstance(seg, list) and len(seg) >= 6:
                    # [x1, y1, x2, y2...] 형태를 [[x1, y1], [x2, y2]...] 형태로 변환
                    points = np.array(seg).reshape(-1, 2).astype(np.int32)
                    cv2.polylines(image, [points], isClosed=True, color=(0, 255, 0), thickness=2)
    return image

# 시각화 실행 예시
# visualize_coco(json_path, image_folder, image_id=7)

---

1-3. COCO 형식을 YOLO 형식으로 변환

YOLO 모델은 0~1 사이의 정규화된(Normalized) 좌표값이 포함된 .txt 파일을 사용합니다.

변환 공식

$$x_{normalized} = \frac{x_{pixel}}{Width}, \quad y_{normalized} = \frac{y_{pixel}}{Height}$$

# COCO JSON -> YOLO TXT 변환 핵심 로직
for ann in coco_data['annotations']:
    image_id = ann['image_id']
    category_id = ann['category_id'] - 1 # 0번부터 시작하도록 조정
    segmentation = ann['segmentation']
    
    image_info = image_dict.get(image_id)
    img_width, img_height = image_info['width'], image_info['height']
    
    yolo_lines = []
    for seg in segmentation:
        if isinstance(seg, list) and len(seg) >= 6:
            # 좌표 정규화
            normalized_points = [
                (seg[i] / img_width, seg[i+1] / img_height) for i in range(0, len(seg), 2)
            ]
            # YOLO 포맷: class_id x1 y1 x2 y2 ...
            yolo_line = f'{category_id} ' + ' '.join([f'{x:.6f} {y:.6f}' for x, y in normalized_points])
            yolo_lines.append(yolo_line)

---

1-4. 데이터셋 분할 (Train / Valid / Test)

모델의 일반화 성능을 위해 데이터를 섞고 6:2:2 등의 비율로 나눕니다.

# random.seed를 고정하면 매번 동일한 분할 결과를 얻을 수 있습니다.
random.seed(2026)
data_list = glob.glob(f'{data_root}/*.jpg')
random.shuffle(data_list)

# 비율에 따른 분할 예시
test_ratio = 0.2
test_list = data_list[:int(num_data*test_ratio)]
valid_list = data_list[int(num_data*test_ratio):int(num_data*test_ratio)*2]
train_list = data_list[int(num_data*test_ratio)*2:]

---

1-5. YAML 설정 및 모델 학습

YOLOv8은 .yaml 파일을 통해 데이터 경로를 파악합니다.

starbucks.yaml 예시

train: /path/to/train
val: /path/to/valid
test: /path/to/test
nc: 2
names: ['logo', 'text']

학습 코드

# 사전 학습된 가중치 로드 (Segmentation 전용)
model = YOLO('yolov8s-seg.pt')

# 모델 학습
results = model.train(
    data=f'{data_root}/starbucks.yaml',
    epochs=10,
    batch=4,
    imgsz=224,
    device=0, # GPU 사용
    name='starbucks_s'
)

---

1-6. 추론 및 결과 확인 (Prediction)

학습이 완료된 베스트 가중치를 사용하여 새로운 이미지에서 로고와 텍스트를 분리해 냅니다.

# 최적의 모델 로드
model = YOLO('/content/runs/segment/starbucks_s/weights/best.pt')

# 예측 수행
results = model.predict(
    source=f'{data_root}/test',
    imgsz=224,
    conf=0.3,
    save=True
)