Segmentation

Ryuzy 2025. 3. 10. 20:15

728x90

1. Segmentation

Segmentation은 컴퓨터 비전에서 이미지나 영상을 픽셀 단위로 분할하여 각 영역이 무엇을 나타내는지 구분하는 기술입니다. 이는 크게 Semantic Segmentation과 Instance Segmentation으로 나뉘는데, Semantic Segmentation은 같은 종류의 객체를 동일한 클래스로 분류하는 반면, Instance Segmentation은 같은 클래스 내에서도 개별 객체를 구분합니다. 이를 통해 의료 영상 분석, 자율주행, 위성 이미지 처리 등 다양한 분야에서 정밀한 객체 인식을 수행할 수 있습니다. Segmentation 모델로는 U-Net, DeepLab, Mask R-CNN 등이 널리 사용됩니다.

2. 스타벅스 데이터셋

아래 이미지와 어노테이션이 있는 압축파일을 다운받고 예제를 진행합니다.

starbucks.zip

0.31MB

!pip install -q ultralytics opencv-python

import ultralytics
ultralytics.checks()

import os
import random
import shutil
from tqdm import tqdm
import cv2
import glob
import json
import matplotlib.pyplot as plt
import numpy as np
from pycocotools import mask as maskUtils
from torchvision import transforms

%pwd

%cd /content/drive/MyDrive/본인의 파일 경로

data_root = '/content/drive/MyDrive/본인의 파일 경로/resource/starbucks'

data_list = glob.glob(f'{data_root}/data/*.jpg') + glob.glob(f'{data_root}/data/*.jpeg')
data_list

# COCO JSON 파일 로드
def load_coco_annotations(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data

load_coco_annotations(f'{data_root}/data/instances_default.json')

# 이미지 파일 불러오기
def load_image(image_path):
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f'이미지를 찾을 수 없음: {image_path}')
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

load_image(data_list[2])

# 폴리곤 및 RLE 주석을 이미지 위에 그리기
# RLE(Run-Length Encoding): 이미지 압축 및 데이터 압축을 위한 기법 중 하나.
# 특히 COCO Datasets의 객체 분할시 정보를 저장할 때 많이 사용됨
def draw_annotations(image, annotations, image_id):
    for ann in annotations:
        if ann['image_id'] == image_id and 'segmentation' in ann:
            segmentation = ann['segmentation']

            # RLE 형식(딕셔너리 타입)
            if isinstance(segmentation, dict) and 'counts' in segmentation:
                try:
                    if isinstance(segmentation['counts'], list):
                        rle = maskUtils.frPyObjects(segmentation, segmentation['size'][0], segmentation['size'][1])
                        mask = maskUtils.decode(rle)
                    else:
                        mask = maskUtils.decode(rle)

                    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
                    cv2.drawContours(image, contours, -1, (255, 0, 0), 2)
                except Exception as e:
                    print(f'에러가 발생: {e}')

            # 폴리곤 형식(리스트 타입)
            elif isinstance(segmentation, list):
                for seg in segmentation:
                    if isinstance(seg, list) and len(seg) >= 6:
                        points = np.array(seg).reshape(-1, 2).astype(np.int32)
                        cv2.polylines(image, [points], isClosed=True, color=(0, 255, 0), thickness=2)

    return image

# 메인 실행 함수
def visualize_coco(json_path, image_folder, image_id):
    coco_data = load_coco_annotations(json_path)
    image_info = next((img for img in coco_data['images'] if img['id'] == image_id), None)
    if not image_info:
        raise ValueError(f'Image ID {image_id} not found in COCO JSON file')

    image_path = os.path.join(image_folder, image_info['file_name'])
    image = load_image(image_path)
    annotated_image = draw_annotations(image, coco_data['annotations'], image_id)

    plt.figure(figsize=(8, 6))
    plt.imshow(annotated_image)
    plt.axis('off')
    plt.title(f'Image ID: {image_id} with Annotations')
    plt.show()

visualize_coco(f'{data_root}/data/instances_default.json', f'{data_root}/data', 20)

visualize_coco(f'{data_root}/data/instances_default.json', f'{data_root}/data', 21)

# COCO JSON 파일 경로와 YOLO 저장할 폴더
coco_json_path = f'{data_root}/data/instances_default.json'
yolo_output_folder = f'{data_root}/data'
image_folder = f'{data_root}/data/'

# YOLO 라벨 폴더 생성(없으면 생성)
os.makedirs(yolo_output_folder, exist_ok=True)

# COCO JSON 로드
with open(coco_json_path, 'r') as f:
    coco_data = json.load(f)

# 이미지 정보 불러오기
image_dict = {img['id']: img for img in coco_data['images']}

# COCO -> YOLO 변환
for ann in coco_data['annotations']:
    image_id = ann['image_id']
    category_id = ann['category_id'] - 1 # YOLO는 class_id가 0부터 시작
    segmentation = ann['segmentation']

    image_info = image_dict.get(image_id)
    if not image_info:
        continue

    img_width, img_height = image_info['width'], image_info['height']
    yolo_label_path = os.path.join(yolo_output_folder, f"{os.path.splitext(image_info['file_name'])[0]}.txt")

    yolo_lines = []

    # RLE 형식
    if isinstance(segmentation, dict) and 'counts' in segmentation:
        rle = maskUtils.frPyObjects(segmentation, img_height, img_width)
        mask = maskUtils.decode(rle)

        contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        for contour in contours:
            if len(contour) >= 3:
                normalized_points = [(x / img_width, y / img_height) for x, y in contour.reshape(-1, 2)]
                yolo_line = f'{category_id} ' + ' '.join([f'{x:.6f} {y:.6f}' for x, y in normalized_points])
                yolo_lines.append(yolo_line)

    elif isinstance(segmentation, list):
        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_line = f'{category_id} ' + ' '.join([f'{x:.6f} {y:6f}' for x, y in normalized_points])
                yolo_lines.append(yolo_line)

    if yolo_lines:
        with open(yolo_label_path, 'w') as f:
            f.write('\n'.join(yolo_lines))

print(f'COCO -> YOLO 변환 완료!')

random.seed(2025)
data_list = glob.glob(f"{data_root}/data/*.jpg") + glob.glob(f"{data_root}/data/*.jpeg")

random.shuffle(data_list)
test_ratio = 0.2
num_data = len(data_list)

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

len(test_list), len(valid_list), len(train_list)

file_root = f'{data_root}/data'

train_root = f'{data_root}/train'
valid_root = f'{data_root}/valid'
test_root = f'{data_root}/test'

for folder in [train_root, valid_root, test_root]:
    if not os.path.exists(folder):
        os.makedirs(folder)

# 파일 복사 함수(이미지 + 동일한 이름의 txt 파일)
def copy_files(file_list, dest_folder):
    for file_path in file_list:
        file_name = os.path.basename(file_path) # 파일 이름 추출 apple.jpg
        file_base, ext = os.path.splitext(file_name) # apple jpg

        # 원본(이미지) 파일 복사
        dest_path = os.path.join(dest_folder, file_name)
        shutil.copy(file_path, dest_path)

        # 텍스트 파일도 복사
        txt_file_path = os.path.join(os.path.dirname(file_path), f'{file_base}.txt')
        if os.path.exists(txt_file_path):
            dest_txt_path = os.path.join(dest_folder, f'{file_base}.txt')
            shutil.copy(txt_file_path, dest_txt_path)

copy_files(train_list, train_root)
copy_files(valid_list, valid_root)
copy_files(test_list, test_root)

train_root = f'/content/drive/MyDrive/본인의 파일 경로/resource/starbucks/train'
valid_root = f'/content/drive/MyDrive/본인의 파일 경로/resource/starbucks/valid'
test_root = f'/content/drive/MyDrive/본인의 파일 경로/resource/starbucks/test'

import yaml

data = dict()
data['train'] = train_root
data['val'] = valid_root
data['test'] = test_root
data['nc'] = 2
data['names'] = ['logo', 'text']

with open(f'./resource/starbucks/starbucks.yaml', 'w') as f:
    yaml.dump(data, f)

%pwd

%cd /content/drive/MyDrive/KDT 시즌4/13. 컴퓨터 비전/resource/starbucks

from ultralytics import YOLO

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

results = model.train(
    data='./starbucks.yaml',
    epochs=10,
    batch=8,
    imgsz=224,
    device=0,
    workers=2,
    amp=False,
    name='starbucks_s'
)

model = YOLO('runs/segment/starbucks_s/weights/best.pt')

results = model.predict(
    source='test',
    imgsz=224,
    conf=0.25,
    device=0,
    save=True,
    save_conf=True
)

728x90