이미지 세그멘테이션
이미지 세그멘테이션 (Image Segmentation)
이미지 세그멘테이션 모델은 이미지에서 관심 영역에 해당하는 부분들을 분리해요. 이 모델들은 각 픽셀에 라벨을 할당하는 방식으로 동작하죠. 세그멘테이션에는 의미론적 세그멘테이션(semantic), 인스턴스 세그멘테이션(instance), 파놉틱 세그멘테이션(panoptic) 같은 여러 종류가 있어요.
이 가이드에서 할 일은 이렇습니다:
시작하기 전에 필요한 라이브러리가 모두 설치돼 있는지 확인하세요:
# uncomment to install the necessary libraries
!pip install -q datasets transformers evaluate accelerate trackio
모델을 커뮤니티와 공유하고 업로드할 수 있도록 Hugging Face 계정에 로그인하길 권장해요. 요청이 오면 토큰을 입력해 로그인합니다:
>>> from huggingface_hub import notebook_login
>>> notebook_login()
세그멘테이션의 종류
의미론적 세그멘테이션은 이미지의 모든 단일 픽셀에 라벨이나 클래스를 할당해요. 의미론적 세그멘테이션 모델의 출력을 살펴볼게요. 이미지에서 만나는 객체의 모든 인스턴스에 같은 클래스를 할당하는데, 예를 들어 모든 고양이를 "cat-1", "cat-2"처럼 나누지 않고 전부 "cat"으로 라벨링해요. transformers의 이미지 세그멘테이션 파이프라인을 쓰면 의미론적 세그멘테이션 모델을 빠르게 추론할 수 있어요. 예시 이미지를 살펴볼게요.
from transformers import pipeline
from PIL import Image
import requests
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/segmentation_input.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image
nvidia/segformer-b1-finetuned-cityscapes-1024-1024을 사용할 거예요.
semantic_segmentation = pipeline("image-segmentation", "nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
results = semantic_segmentation(image)
results
세그멘테이션 파이프라인의 출력에는 예측된 각 클래스에 대한 마스크가 포함돼 있어요.
[{'score': None,
'label': 'road',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'sidewalk',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'building',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'wall',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'pole',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'traffic sign',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'vegetation',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'terrain',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'sky',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': None,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>}]
car 클래스의 마스크를 살펴보면, 모든 자동차가 같은 마스크로 분류된 걸 볼 수 있어요.
results[-1]["mask"]
인스턴스 세그멘테이션은 모든 픽셀을 분류하는 게 목표가 아니라, 주어진 이미지에서 객체의 모든 인스턴스에 대한 마스크를 예측하는 데 목표를 두어요. 객체 탐지와 아주 비슷하게 동작하는데, 각 인스턴스에 바운딩 박스가 있는 대신 세그멘테이션 마스크가 있는 거죠. 여기서는 facebook/mask2former-swin-large-cityscapes-instance를 사용할 거예요.
instance_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-instance")
results = instance_segmentation(image)
results
아래에서 보듯이 여러 대의 자동차가 분류됐고, 자동차와 사람 인스턴스에 속한 픽셀 외에는 분류가 없어요.
[{'score': 0.999944,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999945,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999652,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.903529,
'label': 'person',
'mask': <PIL.Image.Image image mode=L size=612x415>}]
아래에서 자동차 마스크 중 하나를 확인해 봅니다.
results[2]["mask"]
파놉틱 세그멘테이션은 의미론적 세그멘테이션과 인스턴스 세그멘테이션을 결합해요. 모든 픽셀을 클래스와 그 클래스의 인스턴스로 분류하고, 클래스의 각 인스턴스마다 여러 마스크가 있어요. 여기서는 facebook/mask2former-swin-large-cityscapes-panoptic을 사용할 수 있습니다.
panoptic_segmentation = pipeline("image-segmentation", "facebook/mask2former-swin-large-cityscapes-panoptic")
results = panoptic_segmentation(image)
results
아래에서 보듯이 클래스가 더 많아요. 모든 픽셀이 클래스 중 하나로 분류되는 걸 나중에 그림으로 확인해 볼게요.
[{'score': 0.999981,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999958,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.99997,
'label': 'vegetation',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999575,
'label': 'pole',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999958,
'label': 'building',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999634,
'label': 'road',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.996092,
'label': 'sidewalk',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.999221,
'label': 'car',
'mask': <PIL.Image.Image image mode=L size=612x415>},
{'score': 0.99987,
'label': 'sky',
'mask': <PIL.Image.Image image mode=L size=612x415>}]
모든 세그멘테이션 종류를 나란히 비교해 볼게요.
모든 세그멘테이션 종류를 봤으니, 의미론적 세그멘테이션을 위한 모델 파인튜닝을 깊이 파고들어 볼게요.
의미론적 세그멘테이션의 흔한 실무 응용으로는 자율주행차가 보행자와 중요한 교통 정보를 식별하도록 훈련하는 것, 의료 영상에서 세포와 이상을 식별하는 것, 위성 영상에서 환경 변화를 모니터링하는 것 등이 있어요.
세그멘테이션 모델 파인튜닝
이제 이렇게 해 볼 거예요:
- SceneParse150 데이터셋으로 SegFormer를 파인튜닝합니다.
- 파인튜닝한 모델로 추론(inference)을 수행합니다.
이 작업과 호환되는 모든 아키텍처와 체크포인트를 보려면 task-page를 확인해 보세요.
SceneParse150 데이터셋 불러오기
🤗 Datasets 라이브러리에서 SceneParse150 데이터셋의 더 작은 하위 집합부터 불러와 볼게요. 전체 데이터셋으로 훈련 시간을 더 쓰기 전에, 작은 집합으로 실험해서 모든 게 잘 동작하는지 확인할 기회를 갖는 거예요.
>>> from datasets import load_dataset
>>> ds = load_dataset("merve/scene_parse_150", split="train[:50]")
데이터셋의 train 스플릿을 [~datasets.Dataset.train_test_split] 메서드로 훈련·테스트 세트로 나눕니다:
>>> ds = ds.train_test_split(test_size=0.2)
>>> train_ds = ds["train"]
>>> test_ds = ds["test"]
그다음 예시를 하나 살펴볼게요:
>>> train_ds[0]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=512x683 at 0x7F9B0C201F90>,
'annotation': <PIL.PngImagePlugin.PngImageFile image mode=L size=512x683 at 0x7F9B0C201DD0>,
'scene_category': 368}
# view the image
>>> train_ds[0]["image"]
image: 장면의 PIL 이미지.annotation: 세그멘테이션 맵의 PIL 이미지로, 모델의 타깃이기도 해요.scene_category: "kitchen"이나 "office"처럼 이미지 장면을 설명하는 카테고리 id. 이 가이드에서는image와annotation만 필요하며, 둘 다 PIL 이미지입니다.
나중에 모델을 설정할 때 쓸 수 있도록 라벨 id를 라벨 클래스로 매핑하는 딕셔너리도 만들고 싶을 거예요. Hub에서 매핑을 다운로드해 id2label과 label2id 딕셔너리를 만듭니다:
>>> import json
>>> from pathlib import Path
>>> from huggingface_hub import hf_hub_download
>>> repo_id = "huggingface/label-files"
>>> filename = "ade20k-id2label.json"
>>> id2label = json.loads(Path(hf_hub_download(repo_id, filename, repo_type="dataset")).read_text())
>>> id2label = {int(k): v for k, v in id2label.items()}
>>> label2id = {v: k for k, v in id2label.items()}
>>> num_labels = len(id2label)
커스텀 데이터셋
노트북 대신 run_semantic_segmentation.py 스크립트로 훈련하고 싶다면, 나만의 데이터셋을 만들거나 사용할 수도 있어요. 스크립트가 요구하는 것은:
-
"image"와 "label"이라는 두 개의 [
~datasets.Image] 컬럼을 가진 [~datasets.DatasetDict]from datasets import Dataset, DatasetDict, Image image_paths_train = ["path/to/image_1.jpg/jpg", "path/to/image_2.jpg/jpg", ..., "path/to/image_n.jpg/jpg"] label_paths_train = ["path/to/annotation_1.png", "path/to/annotation_2.png", ..., "path/to/annotation_n.png"] image_paths_validation = [...] label_paths_validation = [...] def create_dataset(image_paths, label_paths): dataset = Dataset.from_dict({"image": sorted(image_paths), "label": sorted(label_paths)}) dataset = dataset.cast_column("image", Image()) dataset = dataset.cast_column("label", Image()) return dataset # step 1: create Dataset objects train_dataset = create_dataset(image_paths_train, label_paths_train) validation_dataset = create_dataset(image_paths_validation, label_paths_validation) # step 2: create DatasetDict dataset = DatasetDict({ "train": train_dataset, "validation": validation_dataset, } ) # step 3: push to Hub (assumes you have ran the hf auth login command in a terminal/notebook) dataset.push_to_hub("your-name/dataset-repo") # optionally, you can push to a private repo on the Hub # dataset.push_to_hub("name of repo on the hub", private=True) -
클래스 정수를 클래스 이름으로 매핑하는 id2label 딕셔너리
import json # simple example id2label = {0: 'cat', 1: 'dog'} with open('id2label.json', 'w') as fp: json.dump(id2label, fp)
예시로, 위 단계들로 만든 example dataset을 확인해 보세요.
전처리
다음 단계는 이미지와 주석을 모델용으로 준비할 SegFormer 이미지 프로세서를 로드하는 거예요. 이 데이터셋 같은 일부 데이터셋은 영(0) 인덱스를 배경 클래스로 사용해요. 하지만 배경 클래스는 실제로 150개 클래스에 포함되지 않아서, 모든 라벨에서 1을 빼도록 do_reduce_labels=True를 설정해야 해요. 영 인덱스는 255로 교체돼 SegFormer의 손실 함수에서 무시됩니다:
>>> from transformers import AutoImageProcessor
>>> checkpoint = "nvidia/mit-b0"
>>> image_processor = AutoImageProcessor.from_pretrained(checkpoint, do_reduce_labels=True)
이미지 데이터셋에 몇 가지 데이터 증강을 적용해 과적합에 더 강건한 모델을 만드는 건 흔한 일이에요. 이 가이드에서는 torchvision의 ColorJitter 함수로 이미지의 색상 속성을 무작위로 바꿀 건데, 원하는 이미지 라이브러리를 써도 됩니다.
>>> from torchvision.transforms import ColorJitter
>>> jitter = ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)
이제 이미지와 주석을 모델용으로 준비하는 전처리 함수 두 개를 만듭니다. 이 함수들은 이미지를 pixel_values로, 주석을 labels로 변환해요. 훈련 세트에서는 이미지를 이미지 프로세서에 넘기기 전에 jitter를 적용하고, 테스트 세트에서는 이미지 프로세서가 images를 크롭·정규화하고 labels만 크롭해요. 테스트 중에는 데이터 증강을 적용하지 않기 때문이죠.
>>> def train_transforms(example_batch):
... images = [jitter(x) for x in example_batch["image"]]
... labels = [x for x in example_batch["annotation"]]
... inputs = image_processor(images, labels)
... return inputs
>>> def val_transforms(example_batch):
... images = [x for x in example_batch["image"]]
... labels = [x for x in example_batch["annotation"]]
... inputs = image_processor(images, labels)
... return inputs
전체 데이터셋에 jitter를 적용하려면 🤗 Datasets의 [~datasets.Dataset.set_transform] 함수를 쓰세요. 변환이 그때그때 적용되어 더 빠르고 디스크 공간도 덜 차지해요:
>>> train_ds.set_transform(train_transforms)
>>> test_ds.set_transform(val_transforms)
평가
훈련 중에 지표를 포함하면 모델 성능을 평가하는 데 자주 도움이 돼요. 🤗 Evaluate 라이브러리로 평가 방법을 빠르게 불러올 수 있습니다. 이 작업에서는 mean Intersection over Union (IoU) 지표를 로드합니다 (지표를 로드하고 계산하는 방법을 더 배우려면 🤗 Evaluate quick tour를 확인해 보세요):
>>> import evaluate
>>> metric = evaluate.load("mean_iou")
그다음 지표를 [~evaluate.EvaluationModule.compute]하는 함수를 만듭니다. 예측을 먼저 로짓으로 변환하고, 라벨 크기에 맞게 리셰이프한 뒤에 [~evaluate.EvaluationModule.compute]를 호출해야 해요:
>>> import numpy as np
>>> import torch
>>> from torch import nn
>>> def compute_metrics(eval_pred):
... with torch.no_grad():
... logits, labels = eval_pred
... logits_tensor = torch.from_numpy(logits)
... logits_tensor = nn.functional.interpolate(
... logits_tensor,
... size=labels.shape[-2:],
... mode="bilinear",
... align_corners=False,
... ).argmax(dim=1)
... pred_labels = logits_tensor.detach().cpu().numpy()
... metrics = metric.compute(
... predictions=pred_labels,
... references=labels,
... num_labels=num_labels,
... ignore_index=255,
... reduce_labels=False,
... )
... for key, value in metrics.items():
... if isinstance(value, np.ndarray):
... metrics[key] = value.tolist()
... return metrics
compute_metrics 함수가 준비됐으니, 훈련을 설정할 때 다시 쓰게 될 거예요.
훈련
[Trainer]로 모델을 파인튜닝하는 방법이 익숙하지 않다면 여기의 기본 튜토리얼을 먼저 확인해 보세요!
모델 훈련을 시작할 준비가 됐어요! [AutoModelForSemanticSegmentation]으로 SegFormer를 로드하고, 모델에 라벨 id와 라벨 클래스 사이의 매핑을 넘겨 주세요:
>>> from transformers import AutoModelForSemanticSegmentation, TrainingArguments, Trainer
>>> model = AutoModelForSemanticSegmentation.from_pretrained(checkpoint, id2label=id2label, label2id=label2id)
이 시점에 남은 단계는 세 개뿐이에요:
- [
TrainingArguments]에 훈련 하이퍼파라미터를 정의합니다. 사용하지 않는 컬럼을 제거하면image컬럼이 사라질 수 있으니 주의해야 해요.image컬럼이 없으면pixel_values를 만들 수 없거든요. 이 동작을 막으려면remove_unused_columns=False로 설정하세요! 유일한 필수 파라미터는 모델을 저장할 위치를 정하는output_dir이에요.push_to_hub=True로 설정하면 이 모델을 Hub에 push합니다 (모델을 업로드하려면 Hugging Face에 로그인해야 해요). [Trainer]는 각 에폭이 끝날 때 IoU 지표를 평가하고 훈련 체크포인트를 저장해요. - 훈련 인자를 모델, 데이터셋, 토크나이저, 데이터 콜레이터,
compute_metrics함수와 함께 [Trainer]로 넘깁니다. - [
~Trainer.train]을 호출해 모델을 파인튜닝합니다.
>>> training_args = TrainingArguments(
... output_dir="segformer-b0-scene-parse-150",
... learning_rate=6e-5,
... num_train_epochs=50,
... per_device_train_batch_size=2,
... per_device_eval_batch_size=2,
... save_total_limit=3,
... eval_strategy="steps",
... save_strategy="steps",
... save_steps=20,
... eval_steps=20,
... logging_steps=1,
... report_to="trackio",
... run_name="scene-parse-150",
... eval_accumulation_steps=5,
... remove_unused_columns=False,
... push_to_hub=True,
... )
>>> trainer = Trainer(
... model=model,
... args=training_args,
... train_dataset=train_ds,
... eval_dataset=test_ds,
... compute_metrics=compute_metrics,
... )
>>> trainer.train()
훈련이 끝나면 [~transformers.Trainer.push_to_hub] 메서드로 모델을 Hub에 공유해 누구나 쓸 수 있게 하세요:
>>> trainer.push_to_hub()
추론
좋아요, 이제 모델을 파인튜닝했으니 추론에 쓸 수 있어요!
데이터셋을 다시 로드하고 추론용 이미지를 불러옵니다.
>>> from datasets import load_dataset
>>> ds = load_dataset("scene_parse_150", split="train[:50]")
>>> ds = ds.train_test_split(test_size=0.2)
>>> test_ds = ds["test"]
>>> image = ds["test"][0]["image"]
>>> image
이제 파이프라인 없이 추론하는 방법을 볼게요. 이미지 프로세서로 이미지를 처리하고 pixel_values를 GPU에 올립니다:
>>> from accelerate import Accelerator
>>> device = Accelerator().device
>>> encoding = image_processor(image, return_tensors="pt")
>>> pixel_values = encoding.pixel_values.to(device)
입력을 모델에 넘기고 logits을 얻습니다:
>>> outputs = model(pixel_values=pixel_values)
>>> logits = outputs.logits.cpu()
다음으로 로짓을 원본 이미지 크기로 다시 스케일합니다:
>>> upsampled_logits = nn.functional.interpolate(
... logits,
... size=image.size[::-1],
... mode="bilinear",
... align_corners=False,
... )
>>> pred_seg = upsampled_logits.argmax(dim=1)[0]
결과를 시각화하려면, 각 클래스를 RGB 값으로 매핑하는 dataset color palette를 ade_palette()로 로드합니다.
def ade_palette():
return np.asarray([
[0, 0, 0],
[120, 120, 120],
[180, 120, 120],
[6, 230, 230],
[80, 50, 50],
[4, 200, 3],
[120, 120, 80],
[140, 140, 140],
[204, 5, 255],
[230, 230, 230],
[4, 250, 7],
[224, 5, 255],
[235, 255, 7],
[150, 5, 61],
[120, 120, 70],
[8, 255, 51],
[255, 6, 82],
[143, 255, 140],
[204, 255, 4],
[255, 51, 7],
[204, 70, 3],
[0, 102, 200],
[61, 230, 250],
[255, 6, 51],
[11, 102, 255],
[255, 7, 71],
[255, 9, 224],
[9, 7, 230],
[220, 220, 220],
[255, 9, 92],
[112, 9, 255],
[8, 255, 214],
[7, 255, 224],
[255, 184, 6],
[10, 255, 71],
[255, 41, 10],
[7, 255, 255],
[224, 255, 8],
[102, 8, 255],
[255, 61, 6],
[255, 194, 7],
[255, 122, 8],
[0, 255, 20],
[255, 8, 41],
[255, 5, 153],
[6, 51, 255],
[235, 12, 255],
[160, 150, 20],
[0, 163, 255],
[140, 140, 140],
[250, 10, 15],
[20, 255, 0],
[31, 255, 0],
[255, 31, 0],
[255, 224, 0],
[153, 255, 0],
[0, 0, 255],
[255, 71, 0],
[0, 235, 255],
[0, 173, 255],
[31, 0, 255],
[11, 200, 200],
[255, 82, 0],
[0, 255, 245],
[0, 61, 255],
[0, 255, 112],
[0, 255, 133],
[255, 0, 0],
[255, 163, 0],
[255, 102, 0],
[194, 255, 0],
[0, 143, 255],
[51, 255, 0],
[0, 82, 255],
[0, 255, 41],
[0, 255, 173],
[10, 0, 255],
[173, 255, 0],
[0, 255, 153],
[255, 92, 0],
[255, 0, 255],
[255, 0, 245],
[255, 0, 102],
[255, 173, 0],
[255, 0, 20],
[255, 184, 184],
[0, 31, 255],
[0, 255, 61],
[0, 71, 255],
[255, 0, 204],
[0, 255, 194],
[0, 255, 82],
[0, 10, 255],
[0, 112, 255],
[51, 0, 255],
[0, 194, 255],
[0, 122, 255],
[0, 255, 163],
[255, 153, 0],
[0, 255, 10],
[255, 112, 0],
[143, 255, 0],
[82, 0, 255],
[163, 255, 0],
[255, 235, 0],
[8, 184, 170],
[133, 0, 255],
[0, 255, 92],
[184, 0, 255],
[255, 0, 31],
[0, 184, 255],
[0, 214, 255],
[255, 0, 112],
[92, 255, 0],
[0, 224, 255],
[112, 224, 255],
[70, 184, 160],
[163, 0, 255],
[153, 0, 255],
[71, 255, 0],
[255, 0, 163],
[255, 204, 0],
[255, 0, 143],
[0, 255, 235],
[133, 255, 0],
[255, 0, 235],
[245, 0, 255],
[255, 0, 122],
[255, 245, 0],
[10, 190, 212],
[214, 255, 0],
[0, 204, 255],
[20, 0, 255],
[255, 255, 0],
[0, 153, 255],
[0, 41, 255],
[0, 255, 204],
[41, 0, 255],
[41, 255, 0],
[173, 0, 255],
[0, 245, 255],
[71, 0, 255],
[122, 0, 255],
[0, 255, 184],
[0, 92, 255],
[184, 255, 0],
[0, 133, 255],
[255, 214, 0],
[25, 194, 194],
[102, 255, 0],
[92, 0, 255],
])
그다음 이미지와 예측된 세그멘테이션 맵을 결합해 그릴 수 있어요:
>>> import matplotlib.pyplot as plt
>>> import numpy as np
>>> color_seg = np.zeros((pred_seg.shape[0], pred_seg.shape[1], 3), dtype=np.uint8)
>>> palette = np.array(ade_palette())
>>> for label, color in enumerate(palette):
... color_seg[pred_seg == label, :] = color
>>> color_seg = color_seg[..., ::-1] # convert to BGR
>>> img = np.array(image) * 0.5 + color_seg * 0.5 # plot the image with the segmentation map
>>> img = img.astype(np.uint8)
>>> plt.figure(figsize=(15, 10))
>>> plt.imshow(img)
>>> plt.show()
더 알아보기 (Learn more)
- task-page: 이미지 세그멘테이션과 호환되는 모든 아키텍처·체크포인트 목록.
- run_semantic_segmentation.py: 세그멘테이션 훈련용 스크립트.
- 🤗 Evaluate quick tour: 지표 로드·계산 방법.